diff --git a/ENV_VARIABLES.md b/ENV_VARIABLES.md index 094576278..e65080b09 100644 --- a/ENV_VARIABLES.md +++ b/ENV_VARIABLES.md @@ -926,12 +926,37 @@ conflicts closed. How often the post-processor dispatcher polls for completed runs needing post-processing. This is a **backfill/catch-up** mechanism — the primary dispatch path is event-driven (coder workers enqueue directly on run completion). The 30s default keeps idle RU consumption low while still catching missed events or version-upgrade backfills within a reasonable window. Reduce temporarily for large backfills. +### SCHEDULER_HANDLER_CACHE_TTL_MS +**Default:** `300000` +**Type:** integer (milliseconds) + +TTL for the DAG dispatcher's in-memory cache of handler topology (the `services` collection, `type: "post-process-handler"`). `loadHandlers()` runs on the notify/poll hot path (2–3 reads per notify plus every poll cycle) but topology only changes on deploys, so caching it eliminates nearly all of those reads (meaningful on CosmosDB where each query costs RUs). Cache coherence comes primarily from `HandlerDispatcher.registerHandler` invalidating the cache the moment topology changes — the scheduler is a singleton (`replicas: 1`, `Recreate`) and the sole writer of the collection — so this TTL is only a coarse backstop for out-of-band writes or a future scale-out to more than one replica. Kept well above `SCHEDULER_PP_POLL_INTERVAL_MS` so the poll loop doesn't force a fresh read every cycle. + ### QUEUE_NAME_POST_PROCESSOR **Default:** `post-processor-queue` **Type:** string Azure Storage Queue name used by both the scheduler (to enqueue post-processing work) and the post-processor worker (to dequeue). Must match between the two services. +### SCHEDULER_URL +**Default:** (not set) +**Type:** URL string + +Optional base URL for the scheduler's HTTP notification endpoints. When set, coder workers call `POST /notify/run-terminal` after a run reaches a terminal state, and post-processing handlers call `POST /notify/handler-complete` after a handler finishes. If unset, workers fall back to the legacy direct post-processor queue enqueue path and the scheduler's polling loop catches up. + +- **Docker Compose:** `http://scheduler:8080` +- **Kubernetes:** `http://scheduler-service.scoped.svc.cluster.local:8080` +- **Local dev:** Leave unset to rely on legacy queue fallback unless you are also running the scheduler locally + +### API_URL +**Default:** `http://api:80` +**Type:** URL string + +Base URL the **scheduler** uses to reach the API's `POST /api/v1/reports/trigger` endpoint. When a run's post-processing handler DAG drains (every registered handler reaches a terminal state), the scheduler calls this endpoint so the API can evaluate each report template's trigger and create the templated report job(s). Reports are deliberately not DAG handlers — see [docs/architecture/post-processing.md](docs/architecture/post-processing.md#report-triggering-on-dag-drain). + +- **Docker Compose:** `http://api:80` +- **Kubernetes:** `http://api.scoped.svc.cluster.local:80` + ### SCOPE_REAPER_ENABLED **Default:** `false` **Type:** boolean (`true` to enable) diff --git a/apps/scheduler/src/handler-dispatcher.test.ts b/apps/scheduler/src/handler-dispatcher.test.ts new file mode 100644 index 000000000..32d954243 --- /dev/null +++ b/apps/scheduler/src/handler-dispatcher.test.ts @@ -0,0 +1,472 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +import { describe, it, expect, vi, beforeEach, afterEach } from "vitest"; +import type { HandlerServiceDocument } from "shared"; +import { HandlerDispatcher } from "./handler-dispatcher.js"; + +const HANDLERS: HandlerServiceDocument[] = [ + { + _id: "pp-atif", + type: "post-process-handler", + version: 1, + queue: "post-processor-queue", + selector: "atif", + autoBackfill: true, + dependsOn: [], + }, + { + _id: "pp-example", + type: "post-process-handler", + version: 1, + queue: "pp-example-queue", + selector: "example", + autoBackfill: false, + dependsOn: ["pp-atif"], + }, +]; + +function createMockCollection() { + return { + findOne: vi.fn(), + findOneAndUpdate: vi.fn(), + updateOne: vi.fn().mockResolvedValue({ matchedCount: 1 }), + find: vi.fn(), + } as any; +} + +function createMockDb(handlers: HandlerServiceDocument[]) { + return { + collection: vi.fn().mockReturnValue({ + find: vi.fn().mockReturnValue({ toArray: vi.fn().mockResolvedValue(handlers) }), + updateOne: vi.fn().mockResolvedValue(undefined), + }), + } as any; +} + +function makeDispatcher( + collection: any, + handlers: HandlerServiceDocument[] = HANDLERS, + apiUrl = "http://api.test", +) { + const db = createMockDb(handlers); + const createQueueClient = vi.fn(); + return new HandlerDispatcher(collection, db, createQueueClient, 60_000, 30, apiUrl); +} + +function makeDispatcherWithQueue( + collection: any, + handlers: HandlerServiceDocument[] = HANDLERS, + apiUrl = "http://api.test", +) { + const db = createMockDb(handlers); + const queueClient = { + createIfNotExists: vi.fn().mockResolvedValue(undefined), + sendMessage: vi.fn().mockResolvedValue(undefined), + }; + const createQueueClient = vi.fn().mockReturnValue(queueClient); + const dispatcher = new HandlerDispatcher(collection, db, createQueueClient, 60_000, 30, apiUrl); + return { dispatcher, queueClient, createQueueClient }; +} + +describe("HandlerDispatcher — drain detection", () => { + let collection: ReturnType; + + beforeEach(() => { + collection = createMockCollection(); + }); + + function isDrained( + handlers: HandlerServiceDocument[], + handlerStatus: Record | undefined, + ): boolean { + const d = makeDispatcher(collection, handlers) as any; + const graph = d.buildGraph(handlers); + return d.isDagDrained(handlers, graph, handlerStatus); + } + + it("is drained when all handlers are done", () => { + expect( + isDrained(HANDLERS, { "pp-atif": { status: "done" }, "pp-example": { status: "done" } }), + ).toBe(true); + }); + + it("is NOT drained when the leaf is still pending", () => { + expect(isDrained(HANDLERS, { "pp-atif": { status: "done" } })).toBe(false); + }); + + it("is NOT drained when a handler is processing", () => { + expect( + isDrained(HANDLERS, { + "pp-atif": { status: "done" }, + "pp-example": { status: "processing" }, + }), + ).toBe(false); + }); + + it("is drained when a failed dependency blocks its descendant", () => { + // pp-atif failed → pp-example can never run (blocked) → terminal. + expect(isDrained(HANDLERS, { "pp-atif": { status: "failed" } })).toBe(true); + }); + + it("is drained when the leaf itself failed", () => { + expect( + isDrained(HANDLERS, { + "pp-atif": { status: "done" }, + "pp-example": { status: "failed" }, + }), + ).toBe(true); + }); + + it("is vacuously drained when there are no handlers", () => { + expect(isDrained([], undefined)).toBe(true); + }); +}); + +describe("HandlerDispatcher — maybeTriggerReports", () => { + let collection: ReturnType; + const fetchMock = vi.fn(); + + beforeEach(() => { + collection = createMockCollection(); + fetchMock.mockReset(); + vi.stubGlobal("fetch", fetchMock); + }); + + afterEach(() => { + vi.unstubAllGlobals(); + }); + + it("POSTs to /api/v1/reports/trigger once the DAG drains", async () => { + collection.findOne.mockResolvedValue({ + _id: "req-1", + run: { + _id: "run-1", + status: "done", + handlerStatus: { "pp-atif": { status: "done" }, "pp-example": { status: "done" } }, + }, + }); + collection.findOneAndUpdate.mockResolvedValue({ _id: "req-1" }); // claim succeeds + fetchMock.mockResolvedValue({ ok: true, status: 201, text: async () => "" }); + + const d = makeDispatcher(collection) as any; + await d.maybeTriggerReports("req-1", "run-1"); + + expect(collection.findOneAndUpdate).toHaveBeenCalledWith( + expect.objectContaining({ + _id: "req-1", + "run._id": "run-1", + "run.status": "done", + "run.reportsTriggeredAt": { $exists: false }, + }), + expect.objectContaining({ $set: expect.objectContaining({ "run.reportsTriggeredAt": expect.any(Date) }) }), + expect.objectContaining({ returnDocument: "after" }), + ); + expect(fetchMock).toHaveBeenCalledTimes(1); + const [url, init] = fetchMock.mock.calls[0]; + expect(url).toBe("http://api.test/api/v1/reports/trigger"); + expect(init.method).toBe("POST"); + expect(JSON.parse(init.body)).toEqual({ requestId: "req-1" }); + }); + + it("does not trigger when the DAG has not drained", async () => { + collection.findOne.mockResolvedValue({ + _id: "req-1", + run: { _id: "run-1", status: "done", handlerStatus: { "pp-atif": { status: "done" } } }, + }); + + const d = makeDispatcher(collection) as any; + await d.maybeTriggerReports("req-1", "run-1"); + + expect(collection.findOneAndUpdate).not.toHaveBeenCalled(); + expect(fetchMock).not.toHaveBeenCalled(); + }); + + it("does not trigger when reports were already triggered", async () => { + collection.findOne.mockResolvedValue({ + _id: "req-1", + run: { _id: "run-1", status: "done", reportsTriggeredAt: new Date(), handlerStatus: {} }, + }); + + const d = makeDispatcher(collection) as any; + await d.maybeTriggerReports("req-1", "run-1"); + + expect(collection.findOneAndUpdate).not.toHaveBeenCalled(); + expect(fetchMock).not.toHaveBeenCalled(); + }); + + it("does not trigger when the current attempt differs from the notified runId", async () => { + // A stale notify for run-1 arrives, but the request has since been retried + // into run-2 — the guard must not fire for the wrong attempt. + collection.findOne.mockResolvedValue({ + _id: "req-1", + run: { + _id: "run-2", + status: "done", + handlerStatus: { "pp-atif": { status: "done" }, "pp-example": { status: "done" } }, + }, + }); + + const d = makeDispatcher(collection) as any; + await d.maybeTriggerReports("req-1", "run-1"); + + expect(collection.findOneAndUpdate).not.toHaveBeenCalled(); + expect(fetchMock).not.toHaveBeenCalled(); + }); + + it("does not POST twice when the guard claim is lost to a concurrent path", async () => { + collection.findOne.mockResolvedValue({ + _id: "req-1", + run: { + _id: "run-1", + status: "done", + handlerStatus: { "pp-atif": { status: "done" }, "pp-example": { status: "done" } }, + }, + }); + collection.findOneAndUpdate.mockResolvedValue(null); // another path claimed it first + + const d = makeDispatcher(collection) as any; + await d.maybeTriggerReports("req-1", "run-1"); + + expect(fetchMock).not.toHaveBeenCalled(); + }); + + it("rolls back the guard when the trigger POST fails", async () => { + collection.findOne.mockResolvedValue({ + _id: "req-1", + run: { + _id: "run-1", + status: "done", + handlerStatus: { "pp-atif": { status: "done" }, "pp-example": { status: "done" } }, + }, + }); + collection.findOneAndUpdate.mockResolvedValue({ _id: "req-1" }); + fetchMock.mockResolvedValue({ ok: false, status: 500, text: async () => "boom" }); + + const d = makeDispatcher(collection) as any; + d.reportRetryOptions = { maxRetries: 1, baseDelayMs: 1, isRetryable: () => true }; + await d.maybeTriggerReports("req-1", "run-1"); + + // Guard rolled back via $unset (scoped to the attempt) so the poll net retries. + expect(collection.updateOne).toHaveBeenCalledWith( + { _id: "req-1", "run._id": "run-1" }, + { $unset: { "run.reportsTriggeredAt": "" } }, + ); + }); + + it("ignores runs that are not done", async () => { + collection.findOne.mockResolvedValue({ _id: "req-1", run: { _id: "run-1", status: "processing" } }); + + const d = makeDispatcher(collection) as any; + await d.maybeTriggerReports("req-1", "run-1"); + + expect(collection.findOneAndUpdate).not.toHaveBeenCalled(); + expect(fetchMock).not.toHaveBeenCalled(); + }); +}); + +describe("HandlerDispatcher — onHandlerComplete", () => { + let collection: ReturnType; + + beforeEach(() => { + collection = createMockCollection(); + }); + + it("only accepts completion when current handler status is processing", async () => { + const d = makeDispatcher(collection) as any; + const maybeTriggerReportsSpy = vi + .spyOn(d, "maybeTriggerReports") + .mockResolvedValue(undefined); + + await d.onHandlerComplete("req-1", "run-1", "pp-atif", "failed"); + + expect(collection.updateOne).toHaveBeenCalledWith( + { _id: "req-1", "run.handlerStatus.pp-atif.status": "processing" }, + { + $set: { + "run.handlerStatus.pp-atif.status": "failed", + "run.handlerStatus.pp-atif.updatedAt": expect.any(Date), + }, + }, + ); + expect(maybeTriggerReportsSpy).toHaveBeenCalledWith("req-1", "run-1"); + }); + + it("ignores stale or duplicate completion notifications", async () => { + collection.updateOne.mockResolvedValue({ matchedCount: 0 }); + const d = makeDispatcher(collection) as any; + const maybeTriggerReportsSpy = vi + .spyOn(d, "maybeTriggerReports") + .mockResolvedValue(undefined); + const loadHandlersSpy = vi.spyOn(d, "loadHandlers"); + + await d.onHandlerComplete("req-1", "run-1", "pp-atif", "done"); + + expect(maybeTriggerReportsSpy).not.toHaveBeenCalled(); + expect(loadHandlersSpy).not.toHaveBeenCalled(); + }); +}); + +describe("HandlerDispatcher — handler topology caching", () => { + function makeWithSpies(ttlMs = 60_000, handlers: HandlerServiceDocument[] = HANDLERS) { + const collection = createMockCollection(); + const toArray = vi.fn().mockResolvedValue(handlers); + const find = vi.fn().mockReturnValue({ toArray }); + const servicesUpdateOne = vi.fn().mockResolvedValue(undefined); + const db = { + collection: vi.fn().mockReturnValue({ find, updateOne: servicesUpdateOne }), + } as any; + const dispatcher = new HandlerDispatcher( + collection, + db, + vi.fn(), + 60_000, + 30, + "http://api.test", + ttlMs, + ) as any; + return { dispatcher, find, toArray, servicesUpdateOne }; + } + + it("reuses cached topology within the TTL window", async () => { + const { dispatcher, toArray } = makeWithSpies(60_000); + + await dispatcher.loadHandlers(); + await dispatcher.loadHandlers(); + + expect(toArray).toHaveBeenCalledTimes(1); + }); + + it("re-reads topology once the TTL has elapsed", async () => { + const { dispatcher, toArray } = makeWithSpies(0); // TTL 0 → every read is stale + + await dispatcher.loadHandlers(); + await dispatcher.loadHandlers(); + + expect(toArray).toHaveBeenCalledTimes(2); + }); + + it("invalidates the cache when a handler is registered", async () => { + const { dispatcher, toArray } = makeWithSpies(60_000); + + await dispatcher.loadHandlers(); // read #1 — populates cache + await dispatcher.registerHandler({ + _id: "pp-new", + type: "post-process-handler", + version: 1, + queue: "pp-new-queue", + selector: "new", + autoBackfill: false, + dependsOn: [], + }); // read #2 (forceFresh validation) + invalidate + await dispatcher.loadHandlers(); // read #3 — cache was invalidated + + expect(toArray).toHaveBeenCalledTimes(3); + }); +}); + +describe("HandlerDispatcher — buildDispatchFilter", () => { + let collection: ReturnType; + + beforeEach(() => { + collection = createMockCollection(); + }); + + it("allows done runs to be considered for autoBackfill version bumps", () => { + const d = makeDispatcher(collection) as any; + const autoBackfillHandler = HANDLERS[0]; + + const filter = d.buildDispatchFilter(autoBackfillHandler, HANDLERS, d.buildGraph(HANDLERS)); + + expect(filter[`run.handlerStatus.${autoBackfillHandler._id}.status`]).toEqual({ + $nin: ["queued", "processing"], + }); + expect(filter.$or).toEqual([ + { [`run.handlerStatus.${autoBackfillHandler._id}.version`]: { $exists: false } }, + { [`run.handlerStatus.${autoBackfillHandler._id}.version`]: { $lt: autoBackfillHandler.version } }, + ]); + }); + + it("keeps done runs excluded for non-backfill handlers", () => { + const d = makeDispatcher(collection) as any; + const noBackfillHandler = HANDLERS[1]; + + const filter = d.buildDispatchFilter(noBackfillHandler, HANDLERS, d.buildGraph(HANDLERS)); + + expect(filter[`run.handlerStatus.${noBackfillHandler._id}.status`]).toEqual({ + $nin: ["queued", "processing", "done"], + }); + }); + + describe("HandlerDispatcher — dispatchIfEligible", () => { + let collection: ReturnType; + + beforeEach(() => { + collection = createMockCollection(); + }); + + it("allows autoBackfill re-dispatch for outdated done versions", async () => { + collection.findOne.mockResolvedValue({ + _id: "req-1", + run: { + _id: "run-1", + status: "done", + handlerStatus: { + "pp-atif": { status: "done", version: 0 }, + }, + }, + }); + collection.updateOne.mockResolvedValue({ matchedCount: 1 }); + + const { dispatcher, queueClient } = makeDispatcherWithQueue(collection); + await (dispatcher as any).dispatchIfEligible( + "req-1", + "run-1", + HANDLERS[0], + (dispatcher as any).buildGraph(HANDLERS), + HANDLERS, + ); + + expect(queueClient.sendMessage).toHaveBeenCalledTimes(1); + expect(collection.updateOne).toHaveBeenCalledWith( + expect.objectContaining({ + _id: "req-1", + "run._id": "run-1", + "run.status": "done", + "run.handlerStatus.pp-atif.status": { $nin: ["queued", "processing"] }, + $or: [ + { "run.handlerStatus.pp-atif.version": { $exists: false } }, + { "run.handlerStatus.pp-atif.version": { $lt: 1 } }, + ], + }), + { + $set: { + "run.handlerStatus.pp-atif.status": "queued", + "run.handlerStatus.pp-atif.updatedAt": expect.any(Date), + }, + }, + ); + }); + + it("enqueues before marking the handler queued", async () => { + collection.findOne.mockResolvedValue({ + _id: "req-1", + run: { _id: "run-1", status: "done", handlerStatus: {} }, + }); + collection.updateOne.mockResolvedValue({ matchedCount: 1 }); + + const { dispatcher, queueClient } = makeDispatcherWithQueue(collection); + await (dispatcher as any).dispatchIfEligible( + "req-1", + "run-1", + HANDLERS[0], + (dispatcher as any).buildGraph(HANDLERS), + HANDLERS, + ); + + expect(queueClient.sendMessage.mock.invocationCallOrder[0]).toBeLessThan( + collection.updateOne.mock.invocationCallOrder[0], + ); + }); + }); +}); diff --git a/apps/scheduler/src/handler-dispatcher.ts b/apps/scheduler/src/handler-dispatcher.ts new file mode 100644 index 000000000..4ff1f8275 --- /dev/null +++ b/apps/scheduler/src/handler-dispatcher.ts @@ -0,0 +1,538 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +import { Collection, Db } from "mongodb"; +import { QueueClient } from "@azure/storage-queue"; +import type { RequestDocument, HandlerServiceDocument, HandlerRunStatus } from "shared"; +import { DependencyGraph, withRetry } from "shared"; +import type { RetryOptions } from "shared"; +import type { NotifyHandler } from "./notify-routes.js"; + +/** + * HandlerDispatcher — DAG-aware post-processing orchestrator. + * + * Reads handler topology from the `services` collection (type: "post-process-handler"), + * builds a DependencyGraph, and dispatches handlers whose dependencies are satisfied. + * + * Two trigger paths: + * 1. `onRunTerminal` — a coder worker signals run completion. Dispatches root handlers. + * 2. `onHandlerComplete` — a handler signals completion. Dispatches downstream handlers. + * + * Also runs a poll safety net to catch missed notifications. + */ +export class HandlerDispatcher implements NotifyHandler { + private interval: ReturnType | null = null; + private dispatching = false; + private queueClients = new Map(); + + /** Cached handler topology. The `services` collection changes only on deploys + * (via registerHandler), yet loadHandlers() sits on the notify/poll hot path + * (2–3 reads per notify, plus every poll cycle). A TTL cache eliminates nearly + * all of those reads — significant on CosmosDB where each query costs RUs. + * + * Coherence comes primarily from `registerHandler` invalidating the cache the + * moment topology changes, not from expiry: the scheduler is a singleton + * (`replicas: 1`, `Recreate`) and the sole writer of the `services` collection, + * so a same-process registration is reflected immediately regardless of TTL. + * The TTL is therefore a coarse defensive backstop (out-of-band writes, or a + * future scale-out to >1 replica) and is intentionally set well above the poll + * interval so the poll loop doesn't force a fresh read every cycle. */ + private handlerCache: { docs: HandlerServiceDocument[]; expiresAt: number } | null = null; + + /** Retry policy for the report-trigger POST. The 30s poll net is the durable + * backstop, so in-call retries are modest. Overridable in tests. */ + private reportRetryOptions: RetryOptions = { + maxRetries: 4, + baseDelayMs: 250, + maxDelayMs: 2_000, + isRetryable: () => true, + }; + + constructor( + private readonly collection: Collection, + private readonly db: Db, + private readonly createQueueClient: (queueName: string) => QueueClient, + private readonly pollIntervalMs: number = 30_000, + private readonly batchSize: number = 30, + private readonly apiUrl: string = process.env.API_URL || "http://api:80", + private readonly handlerCacheTtlMs: number = 300_000, + ) {} + + start(): void { + if (this.interval) return; + this.interval = setInterval(() => this.pollDispatch(), this.pollIntervalMs); + void this.pollDispatch(); + } + + async stop(): Promise { + if (this.interval) { + clearInterval(this.interval); + this.interval = null; + } + const deadline = Date.now() + 5_000; + while (this.dispatching && Date.now() < deadline) { + await new Promise((r) => setTimeout(r, 100)); + } + } + + // ── NotifyHandler implementation ──────────────────────────────────── + + async onRunTerminal(requestId: string, runId: string): Promise { + console.log(`[HandlerDispatcher] Notify: run-terminal ${requestId} (run=${runId})`); + const handlers = await this.loadHandlers(); + + if (handlers.length > 0) { + const graph = this.buildGraph(handlers); + const roots = handlers.filter((h) => !h.dependsOn || h.dependsOn.length === 0); + + for (const handler of roots) { + await this.dispatchIfEligible(requestId, runId, handler, graph, handlers); + } + } + + // A run with no handlers (or an already-drained DAG) triggers reports now. + await this.maybeTriggerReports(requestId, runId); + } + + async onHandlerComplete( + requestId: string, + runId: string, + handlerId: string, + status: "done" | "failed", + ): Promise { + console.log( + `[HandlerDispatcher] Notify: handler-complete ${handlerId} for ${requestId} (status=${status})`, + ); + + // Only apply terminal status from an active in-flight execution. + // This avoids stale/duplicate notifies regressing a newly re-dispatched + // handler back to "done" or "failed". + const updateResult = await this.collection.updateOne( + { _id: requestId, [`run.handlerStatus.${handlerId}.status`]: "processing" } as any, + { + $set: { + [`run.handlerStatus.${handlerId}.status`]: status, + [`run.handlerStatus.${handlerId}.updatedAt`]: new Date(), + }, + } as any, + ); + + if (updateResult.matchedCount === 0) { + console.log( + `[HandlerDispatcher] Ignored stale/duplicate completion for ${handlerId} on ${requestId}`, + ); + return; + } + + if (status === "failed") { + // Don't dispatch downstream on failure, but the failure may have drained + // the DAG (descendants are now permanently blocked) — check for reports. + await this.maybeTriggerReports(requestId, runId); + return; + } + + // Find and dispatch eligible downstream handlers + const handlers = await this.loadHandlers(); + if (handlers.length === 0) return; + + const graph = this.buildGraph(handlers); + const descendants = graph.getDescendants(handlerId); + + // Only check immediate children (not all descendants) + const children = handlers.filter( + (h) => h.dependsOn?.includes(handlerId) && descendants.has(h._id), + ); + + for (const child of children) { + await this.dispatchIfEligible(requestId, runId, child, graph, handlers); + } + + // Completing this handler may have drained the DAG (e.g. it was the leaf). + await this.maybeTriggerReports(requestId, runId); + } + + async registerHandler(doc: HandlerServiceDocument): Promise { + // Validate the resulting topology has no cycles before persisting. Read + // fresh (bypass the cache) so validation reflects any concurrent registrations. + const existing = await this.loadHandlers(true); + const merged = [ + ...existing.filter((h) => h._id !== doc._id), + doc, + ]; + // Throws on cycle — rejects invalid registrations. + this.buildGraph(merged); + + const { _id, ...fields } = doc; + await this.db.collection("services").updateOne( + { _id } as any, + { + $set: { + ...fields, + dependsOn: [...fields.dependsOn], + updatedAt: new Date(), + }, + } as any, + { upsert: true }, + ); + // Topology changed — drop the cache so the next read picks up the new handler. + this.invalidateHandlerCache(); + console.log( + `[HandlerDispatcher] Registered handler ${_id} (version=${doc.version}, queue=${doc.queue}, dependsOn=[${doc.dependsOn.join(", ")}])`, + ); + } + + // ── Report trigger on DAG drain ───────────────────────────────────── + + /** + * A handler is "terminal" for a run when it can produce no further work: + * it has succeeded (`done`), failed (`failed`), or is permanently blocked + * because one of its transitive dependencies failed (so it can never run). + */ + private isHandlerTerminal( + handlerId: string, + graph: DependencyGraph, + handlerStatus: Record | undefined, + ): boolean { + const status = handlerStatus?.[handlerId]?.status; + if (status === "done" || status === "failed") return true; + + // Not yet dispatched/processed — terminal only if blocked by a failed + // ancestor (any transitive dependency that failed). + if (status === undefined) { + for (const ancestorId of graph.getAncestors(handlerId)) { + if (handlerStatus?.[ancestorId]?.status === "failed") return true; + } + } + + // "queued" or "processing" — still active. + return false; + } + + /** + * The DAG is drained for a run when every registered handler is terminal. + * Vacuously true when there are no handlers. + */ + private isDagDrained( + handlers: HandlerServiceDocument[], + graph: DependencyGraph, + handlerStatus: Record | undefined, + ): boolean { + return handlers.every((h) => + this.isHandlerTerminal(h._id, graph, handlerStatus), + ); + } + + /** + * If the run's handler DAG has drained and reports have not yet been + * triggered, atomically claim the exactly-once guard and POST to the API's + * report-trigger endpoint. Rolls the guard back on failure so the poll + * safety net retries. + * + * All reads/writes are scoped to the specific attempt (`run._id === runId`) + * so a stale notify or a concurrent retry can never trigger reports for, or + * roll back the guard of, a different attempt. + */ + private async maybeTriggerReports(requestId: string, runId: string): Promise { + const doc = await this.collection.findOne({ _id: requestId } as any); + const run = doc?.run; + if (!run || run._id !== runId || run.status !== "done") return; + if ((run as any).reportsTriggeredAt) return; + + const handlers = await this.loadHandlers(); + const graph = this.buildGraph(handlers); + const handlerStatus = (run as any).handlerStatus as + | Record + | undefined; + + if (!this.isDagDrained(handlers, graph, handlerStatus)) return; + + // Atomically claim the trigger so concurrent notify/poll paths fire once. + const claimed = await this.collection.findOneAndUpdate( + { + _id: requestId, + "run._id": runId, + "run.status": "done", + "run.reportsTriggeredAt": { $exists: false }, + } as any, + { $set: { "run.reportsTriggeredAt": new Date() } } as any, + { returnDocument: "after" }, + ); + if (!claimed) return; // Another path already claimed it. + + try { + await this.triggerReports(requestId); + console.log( + `[HandlerDispatcher] Triggered report generation for ${requestId} (DAG drained)`, + ); + } catch (err) { + // Roll back the guard so a later poll retries — scoped to this attempt so + // a concurrent retry's valid guard is never cleared. + await this.collection.updateOne( + { _id: requestId, "run._id": runId } as any, + { $unset: { "run.reportsTriggeredAt": "" } } as any, + ); + console.error( + `[HandlerDispatcher] Report trigger failed for ${requestId}, rolled back guard:`, + err, + ); + } + } + + /** POSTs to the API's report-trigger endpoint with retry on transient errors. */ + private async triggerReports(requestId: string): Promise { + const url = `${this.apiUrl}/api/v1/reports/trigger`; + await withRetry(async () => { + const res = await fetch(url, { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ requestId }), + }); + if (!res.ok) { + const text = await res.text().catch(() => ""); + throw new Error(`POST ${url} -> ${res.status} ${text}`.trim()); + } + }, this.reportRetryOptions); + } + + // ── Poll safety net ───────────────────────────────────────────────── + + private async pollDispatch(): Promise { + if (this.dispatching) return; + this.dispatching = true; + try { + const handlers = await this.loadHandlers(); + + if (handlers.length > 0) { + const graph = this.buildGraph(handlers); + + for (const handler of handlers) { + const filter = this.buildDispatchFilter(handler, handlers, graph); + if (!filter) continue; + + const candidates = await this.collection + .find(filter as any, { + projection: { _id: 1, "run._id": 1 }, + sort: { updatedAt: -1 }, + limit: this.batchSize, + }) + .toArray(); + + for (const candidate of candidates) { + const runId = (candidate as any).run?._id as string | undefined; + if (!runId) continue; + await this.dispatchIfEligible(candidate._id, runId, handler, graph, handlers); + } + } + } + + // Safety net: trigger reports for any done run whose DAG has drained but + // whose report trigger was missed (e.g. a dropped notify). + await this.pollTriggerReports(); + } catch (err) { + console.error("[HandlerDispatcher] Error during poll dispatch:", err); + } finally { + this.dispatching = false; + } + } + + /** + * Scans for `done` runs whose reports have not yet been triggered and runs + * the drain check on each. Bounded by `batchSize` per poll. + */ + private async pollTriggerReports(): Promise { + const candidates = await this.collection + .find({ + "run.status": "done", + "run.reportsTriggeredAt": { $exists: false }, + deletedAt: { $exists: false }, + } as any) + .limit(this.batchSize) + .toArray(); + + for (const doc of candidates) { + const runId = (doc as any).run?._id as string | undefined; + if (!runId) continue; + await this.maybeTriggerReports(doc._id, runId); + } + } + + // ── Dispatch logic ────────────────────────────────────────────────── + + private async dispatchIfEligible( + requestId: string, + runId: string, + handler: HandlerServiceDocument, + graph: DependencyGraph, + allHandlers: HandlerServiceDocument[], + ): Promise { + // Check if all dependencies are satisfied for this request + const doc = await this.collection.findOne({ _id: requestId } as any); + if (!doc?.run || doc.run.status !== "done") return; + if (doc.run._id !== runId) return; + + const handlerStatus = (doc.run as any).handlerStatus as + | Record + | undefined; + + // Check dependencies are all "done" + if (handler.dependsOn && handler.dependsOn.length > 0) { + for (const depId of handler.dependsOn) { + const depStatus = handlerStatus?.[depId]?.status; + if (depStatus !== "done") return; // Dependency not satisfied + } + } + + // Check whether this handler still needs work. + const currentStatus = handlerStatus?.[handler._id]?.status; + const currentVersion = handlerStatus?.[handler._id]?.version; + const statusExclusions = handler.autoBackfill + ? ["queued", "processing"] + : ["queued", "processing", "done"]; + + if (currentStatus && statusExclusions.includes(currentStatus)) { + return; + } + + if (handler.autoBackfill) { + if (currentVersion !== undefined && currentVersion >= handler.version) { + return; + } + } else if (currentVersion !== undefined) { + return; + } + + // Queue FIRST, then write the best-effort "queued" marker back to Mongo. + // The queue message is the durable work signal; the DB status is only an + // optimization / observability hint. If we crash after sendMessage() but + // before this update, the worker entrypoint claim still makes the message + // safe and prevents a permanent "queued" wedge. + const enqueued = await this.enqueueToHandler(handler, requestId, runId); + if (!enqueued) return; + + // This is NOT the authoritative processing claim. It only records that a + // message was successfully placed on the queue and is expected to be picked + // up soon. The worker performs the real claim when it atomically flips the + // handler from eligible -> processing for this specific run attempt. + const queueMarkFilter = this.buildDispatchFilter(handler, allHandlers, graph); + const queueMarkResult = await this.collection.updateOne( + { + _id: requestId, + "run._id": runId, + ...(queueMarkFilter ?? {}), + } as any, + { + $set: { + [`run.handlerStatus.${handler._id}.status`]: "queued", + [`run.handlerStatus.${handler._id}.updatedAt`]: new Date(), + }, + } as any, + ); + + if ((queueMarkResult.matchedCount ?? 0) === 0) { + console.log( + `[HandlerDispatcher] ${handler._id} for ${requestId} advanced concurrently after enqueue; leaving queue message to worker idempotency guard`, + ); + return; + } + + console.log(`[HandlerDispatcher] Dispatched ${handler._id} for ${requestId}`); + } + + // ── Helpers ───────────────────────────────────────────────────────── + + private buildDispatchFilter( + handler: HandlerServiceDocument, + allHandlers: HandlerServiceDocument[], + _graph: DependencyGraph, + ): Record | null { + const statusExclusions = handler.autoBackfill + ? ["queued", "processing"] + : ["queued", "processing", "done"]; + + const filter: Record = { + "run.status": "done", + deletedAt: { $exists: false }, + [`run.handlerStatus.${handler._id}.status`]: { $nin: statusExclusions }, + }; + + // All dependencies must be "done" + if (handler.dependsOn && handler.dependsOn.length > 0) { + for (const depId of handler.dependsOn) { + filter[`run.handlerStatus.${depId}.status`] = "done"; + } + } + + // Backfill: version check + if (handler.autoBackfill) { + filter.$or = [ + { [`run.handlerStatus.${handler._id}.version`]: { $exists: false } }, + { [`run.handlerStatus.${handler._id}.version`]: { $lt: handler.version } }, + ]; + } else { + // Only dispatch if never processed + filter[`run.handlerStatus.${handler._id}.version`] = { $exists: false }; + } + + return filter; + } + + private async enqueueToHandler( + handler: HandlerServiceDocument, + requestId: string, + runId: string, + ): Promise { + const queueClient = await this.getOrCreateQueueClient(handler.queue); + const message = Buffer.from( + JSON.stringify({ + type: handler.selector, + requestId, + runId, + }), + ).toString("base64"); + + try { + await queueClient.sendMessage(message); + return true; + } catch (err) { + console.error( + `[HandlerDispatcher] Queue send failed for ${handler._id}/${requestId}:`, + err, + ); + return false; + } + } + + private async getOrCreateQueueClient(queueName: string): Promise { + let client = this.queueClients.get(queueName); + if (!client) { + client = this.createQueueClient(queueName); + await client.createIfNotExists(); + this.queueClients.set(queueName, client); + } + return client; + } + + private async loadHandlers(forceFresh = false): Promise { + const now = Date.now(); + if (!forceFresh && this.handlerCache && this.handlerCache.expiresAt > now) { + return this.handlerCache.docs; + } + const docs = (await this.db + .collection("services") + .find({ type: "post-process-handler" }) + .toArray()) as unknown as HandlerServiceDocument[]; + this.handlerCache = { docs, expiresAt: now + this.handlerCacheTtlMs }; + return docs; + } + + /** Drops the cached topology so the next loadHandlers() re-reads from Mongo. + * Called after registerHandler persists a change. */ + private invalidateHandlerCache(): void { + this.handlerCache = null; + } + + private buildGraph(handlers: HandlerServiceDocument[]): DependencyGraph { + const nodes = handlers.map((h) => ({ id: h._id, dependsOn: h.dependsOn })); + return new DependencyGraph(nodes); + } +} diff --git a/apps/scheduler/src/index.ts b/apps/scheduler/src/index.ts index fcaebd232..101750a30 100644 --- a/apps/scheduler/src/index.ts +++ b/apps/scheduler/src/index.ts @@ -2,12 +2,13 @@ // Licensed under the MIT License. import "dotenv/config"; -import http from "node:http"; import { MongoClient } from "mongodb"; import { QueueClient } from "@azure/storage-queue"; import { DefaultAzureCredential } from "@azure/identity"; import { RequestScheduler } from "./request-scheduler.js"; import { PostProcessorDispatcher } from "./post-processor-dispatcher.js"; +import { HandlerDispatcher } from "./handler-dispatcher.js"; +import { createHttpServer, type NotifyHandler } from "./notify-routes.js"; import { StuckRunReaper } from "./stuck-run-reaper.js"; import { RedisHeartbeatStore, type HeartbeatStore } from "shared"; import type { CodingAgentDocument, RequestDocument } from "shared"; @@ -44,6 +45,7 @@ const QUEUE_RECONCILIATION_INTERVAL_MS = parsePositiveInt( "SCHEDULER_QUEUE_RECONCILIATION_INTERVAL_MS", ); const HEALTH_PORT = parseInt(process.env.PORT || "8080", 10); +const API_URL = process.env.API_URL || "http://api:80"; /** * Parse a positive-integer env var with validation. Returns `fallback` (with a @@ -108,16 +110,6 @@ function createQueueClient(queueName: string): QueueClient { return new QueueClient(queueUrl, credential); } -// ── Health Check Server ────────────────────────────────────────────── - -function createHealthServer(): http.Server { - const server = http.createServer((_req, res) => { - res.writeHead(200, { "Content-Type": "application/json" }); - res.end(JSON.stringify({ status: "ok", service: "scheduler" })); - }); - return server; -} - // ── Main ───────────────────────────────────────────────────────────── async function main(): Promise { @@ -167,6 +159,14 @@ async function main(): Promise { 10, ); + // Handler topology is deploy-static and invalidated on registration, so this + // TTL only backstops out-of-band/multi-replica changes — keep it well above + // the poll interval so the poll loop doesn't force a fresh read every cycle. + const handlerCacheTtlMs = parseInt( + process.env.SCHEDULER_HANDLER_CACHE_TTL_MS || "300000", + 10, + ); + const postProcessorDispatcher = new PostProcessorDispatcher( collection, db, @@ -176,6 +176,21 @@ async function main(): Promise { postProcessorDispatcher.start(); console.log("[Scheduler] Post-processor dispatch loop started"); + // Start the DAG-aware handler dispatcher (replaces legacy PostProcessorDispatcher) + const handlerDispatcher = new HandlerDispatcher( + collection, + db, + createQueueClient, + ppPollIntervalMs, + 30, + API_URL, + handlerCacheTtlMs, + ); + handlerDispatcher.start(); + console.log( + `[Scheduler] Handler dispatcher (DAG) started (poll=${ppPollIntervalMs}ms, topologyCacheTtl=${handlerCacheTtlMs}ms)`, + ); + // Start the stuck-run reaper (authoritative backstop). Redis is non-fatal: // on misconfig or connection failure the reaper self-disables (its sweeps // skip when ping() fails) while the dispatch loop above keeps running. @@ -212,10 +227,10 @@ async function main(): Promise { ); } - // Start health check server - const healthServer = createHealthServer(); - healthServer.listen(HEALTH_PORT, () => { - console.log(`[Scheduler] Health check listening on :${HEALTH_PORT}`); + // Wire notify handler to the HandlerDispatcher + const httpServer = createHttpServer(handlerDispatcher); + httpServer.listen(HEALTH_PORT, () => { + console.log(`[Scheduler] HTTP server listening on :${HEALTH_PORT}`); }); // Graceful shutdown @@ -223,9 +238,10 @@ async function main(): Promise { console.log("[Scheduler] Shutting down..."); await scheduler.stop(); await postProcessorDispatcher.stop(); + await handlerDispatcher.stop(); if (stuckRunReaper) await stuckRunReaper.stop(); if (heartbeatStore) await heartbeatStore.close(); - healthServer.close(); + httpServer.close(); await mongoClient.close(); await shutdownTelemetry(); console.log("[Scheduler] Shutdown complete"); diff --git a/apps/scheduler/src/notify-routes.test.ts b/apps/scheduler/src/notify-routes.test.ts new file mode 100644 index 000000000..2d7ef741b --- /dev/null +++ b/apps/scheduler/src/notify-routes.test.ts @@ -0,0 +1,286 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +import type { AddressInfo } from "node:net"; +import type { IncomingMessage } from "node:http"; +import { PassThrough } from "node:stream"; +import { afterEach, describe, expect, it, vi } from "vitest"; +import { + createHttpServer, + readJsonBody, + MAX_BODY_BYTES, + type HttpServerOptions, + type NotifyHandler, +} from "./notify-routes.js"; + +async function startServer( + handler: NotifyHandler, + options?: HttpServerOptions, +): Promise<{ server: ReturnType; baseUrl: string }> { + const server = createHttpServer(handler, options); + + await new Promise((resolve, reject) => { + server.once("error", reject); + server.listen(0, "127.0.0.1", () => resolve()); + }); + + const address = server.address(); + if (!address || typeof address === "string") { + throw new Error("Server did not bind to an address"); + } + + return { + server, + baseUrl: `http://127.0.0.1:${(address as AddressInfo).port}`, + }; +} + +async function stopServer(server: ReturnType): Promise { + await new Promise((resolve, reject) => { + server.close((error) => { + if (error) { + reject(error); + return; + } + resolve(); + }); + }); +} + +describe("createHttpServer", () => { + const servers: Array> = []; + + afterEach(async () => { + while (servers.length > 0) { + const server = servers.pop(); + if (server) { + await stopServer(server); + } + } + }); + + it("returns 200 for the health endpoint", async () => { + const handler: NotifyHandler = { + onRunTerminal: vi.fn().mockResolvedValue(undefined), + onHandlerComplete: vi.fn().mockResolvedValue(undefined), + registerHandler: vi.fn().mockResolvedValue(undefined), + }; + const { server, baseUrl } = await startServer(handler); + servers.push(server); + + const response = await fetch(`${baseUrl}/health`); + + expect(response.status).toBe(200); + expect(await response.json()).toEqual({ status: "ok", service: "scheduler" }); + }); + + it("calls onRunTerminal for valid run-terminal notifications", async () => { + const handler: NotifyHandler = { + onRunTerminal: vi.fn().mockResolvedValue(undefined), + onHandlerComplete: vi.fn().mockResolvedValue(undefined), + registerHandler: vi.fn().mockResolvedValue(undefined), + }; + const { server, baseUrl } = await startServer(handler); + servers.push(server); + + const response = await fetch(`${baseUrl}/notify/run-terminal`, { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ requestId: "req-1", runId: "run-1" }), + }); + + expect(response.status).toBe(200); + expect(await response.json()).toEqual({ ok: true }); + expect(handler.onRunTerminal).toHaveBeenCalledWith("req-1", "run-1"); + expect(handler.onHandlerComplete).not.toHaveBeenCalled(); + }); + + it("calls onHandlerComplete for valid handler-complete notifications", async () => { + const handler: NotifyHandler = { + onRunTerminal: vi.fn().mockResolvedValue(undefined), + onHandlerComplete: vi.fn().mockResolvedValue(undefined), + registerHandler: vi.fn().mockResolvedValue(undefined), + }; + const { server, baseUrl } = await startServer(handler); + servers.push(server); + + const response = await fetch(`${baseUrl}/notify/handler-complete`, { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ + requestId: "req-1", + runId: "run-1", + handlerId: "atif", + status: "done", + }), + }); + + expect(response.status).toBe(200); + expect(await response.json()).toEqual({ ok: true }); + expect(handler.onHandlerComplete).toHaveBeenCalledWith("req-1", "run-1", "atif", "done"); + expect(handler.onRunTerminal).not.toHaveBeenCalled(); + }); + + it("returns 400 when required fields are missing", async () => { + const handler: NotifyHandler = { + onRunTerminal: vi.fn().mockResolvedValue(undefined), + onHandlerComplete: vi.fn().mockResolvedValue(undefined), + registerHandler: vi.fn().mockResolvedValue(undefined), + }; + const { server, baseUrl } = await startServer(handler); + servers.push(server); + + const response = await fetch(`${baseUrl}/notify/handler-complete`, { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ requestId: "req-1", runId: "run-1", status: "done" }), + }); + + expect(response.status).toBe(400); + expect(await response.json()).toEqual({ + error: "requestId, runId, handlerId, and status (done|failed) are required", + }); + expect(handler.onRunTerminal).not.toHaveBeenCalled(); + expect(handler.onHandlerComplete).not.toHaveBeenCalled(); + }); + + it("returns 404 for unknown routes", async () => { + const handler: NotifyHandler = { + onRunTerminal: vi.fn().mockResolvedValue(undefined), + onHandlerComplete: vi.fn().mockResolvedValue(undefined), + registerHandler: vi.fn().mockResolvedValue(undefined), + }; + const { server, baseUrl } = await startServer(handler); + servers.push(server); + + const response = await fetch(`${baseUrl}/notify/unknown`, { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({}), + }); + + expect(response.status).toBe(404); + expect(await response.json()).toEqual({ error: "Not found" }); + }); + + it("calls registerHandler for valid handler registrations", async () => { + const handler: NotifyHandler = { + onRunTerminal: vi.fn().mockResolvedValue(undefined), + onHandlerComplete: vi.fn().mockResolvedValue(undefined), + registerHandler: vi.fn().mockResolvedValue(undefined), + }; + const { server, baseUrl } = await startServer(handler); + servers.push(server); + + const doc = { + _id: "pp-example", + type: "post-process-handler", + version: 1, + queue: "pp-example-queue", + selector: "example", + autoBackfill: false, + dependsOn: ["pp-atif"], + }; + + const response = await fetch(`${baseUrl}/handlers/register`, { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify(doc), + }); + + expect(response.status).toBe(200); + expect(await response.json()).toEqual({ ok: true, handlerId: "pp-example" }); + expect(handler.registerHandler).toHaveBeenCalledWith(doc); + }); + + it("returns 400 for invalid handler registrations", async () => { + const handler: NotifyHandler = { + onRunTerminal: vi.fn().mockResolvedValue(undefined), + onHandlerComplete: vi.fn().mockResolvedValue(undefined), + registerHandler: vi.fn().mockResolvedValue(undefined), + }; + const { server, baseUrl } = await startServer(handler); + servers.push(server); + + const response = await fetch(`${baseUrl}/handlers/register`, { + method: "POST", + headers: { "Content-Type": "application/json" }, + // Missing queue/selector, wrong type, dependsOn not array + body: JSON.stringify({ _id: "bad", type: "wrong", version: "x" }), + }); + + expect(response.status).toBe(400); + expect(handler.registerHandler).not.toHaveBeenCalled(); + }); + + it("returns 413 and does not invoke the handler when the body exceeds the cap", async () => { + const handler: NotifyHandler = { + onRunTerminal: vi.fn().mockResolvedValue(undefined), + onHandlerComplete: vi.fn().mockResolvedValue(undefined), + registerHandler: vi.fn().mockResolvedValue(undefined), + }; + // Tiny cap so a small, deterministic payload trips it (no socket-buffer race). + const { server, baseUrl } = await startServer(handler, { maxBodyBytes: 32 }); + servers.push(server); + + const response = await fetch(`${baseUrl}/notify/run-terminal`, { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ requestId: "r".repeat(100), runId: "run-1" }), + }); + + expect(response.status).toBe(413); + expect(await response.json()).toEqual({ error: "Request body too large" }); + expect(handler.onRunTerminal).not.toHaveBeenCalled(); + }); +}); + +describe("readJsonBody", () => { + /** Build a mock IncomingMessage backed by a PassThrough with the given headers. */ + function mockReq(headers: Record = {}): IncomingMessage & PassThrough { + const req = new PassThrough() as PassThrough & { headers: Record }; + req.headers = headers; + return req as unknown as IncomingMessage & PassThrough; + } + + it("parses a valid JSON body under the cap", async () => { + const req = mockReq(); + const promise = readJsonBody(req, 1024); + req.end(JSON.stringify({ hello: "world" })); + + await expect(promise).resolves.toEqual({ hello: "world" }); + }); + + it("rejects a declared Content-Length over the cap without reading the body", async () => { + const req = mockReq({ "content-length": "5000" }); + const onData = vi.fn(); + req.on("data", onData); + + await expect(readJsonBody(req, 1024)).rejects.toThrow(/limit/); + expect(onData).not.toHaveBeenCalled(); // fast-path rejects before consuming + }); + + it("rejects a chunked body that streams past the cap (no Content-Length)", async () => { + const req = mockReq(); // no content-length → exercises the streaming guard + const promise = readJsonBody(req, 16); + req.write("x".repeat(10)); + req.write("x".repeat(10)); // total 20 > 16 + req.end(); + + await expect(promise).rejects.toThrow(/limit/); + }); + + it("rejects an empty body", async () => { + const req = mockReq(); + const promise = readJsonBody(req, 1024); + req.end(""); + + await expect(promise).rejects.toThrow("Request body is required"); + }); + + it("defaults the cap to MAX_BODY_BYTES (64 KB)", async () => { + const req = mockReq({ "content-length": String(MAX_BODY_BYTES + 1) }); + + await expect(readJsonBody(req)).rejects.toThrow(/limit/); + }); +}); diff --git a/apps/scheduler/src/notify-routes.ts b/apps/scheduler/src/notify-routes.ts new file mode 100644 index 000000000..b63d13dd9 --- /dev/null +++ b/apps/scheduler/src/notify-routes.ts @@ -0,0 +1,297 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +import http from "node:http"; +import type { HandlerServiceDocument } from "shared"; + +/** + * Maximum accepted request body size for the scheduler's HTTP endpoints. + * Notify/register payloads are a few hundred bytes; 64 KB is a generous cap + * that protects this singleton control-plane process from unbounded buffering + * (and OOM) if a buggy caller streams a large body. + */ +export const MAX_BODY_BYTES = 64 * 1024; + +/** Thrown by readJsonBody when the request body exceeds MAX_BODY_BYTES. */ +class PayloadTooLargeError extends Error { + constructor(maxBytes: number) { + super(`Request body exceeds ${maxBytes}-byte limit`); + this.name = "PayloadTooLargeError"; + } +} + +export interface NotifyHandler { + onRunTerminal(requestId: string, runId: string): Promise; + onHandlerComplete( + requestId: string, + runId: string, + handlerId: string, + status: "done" | "failed", + ): Promise; + registerHandler(doc: HandlerServiceDocument): Promise; +} + +interface RunTerminalRequest { + requestId: string; + runId: string; +} + +interface HandlerCompleteRequest { + requestId: string; + runId: string; + handlerId: string; + status: "done" | "failed"; +} + +function sendJson( + res: http.ServerResponse, + statusCode: number, + body: object, + extraHeaders?: http.OutgoingHttpHeaders, +): void { + res.writeHead(statusCode, { "Content-Type": "application/json", ...extraHeaders }); + res.end(JSON.stringify(body)); +} + +function isRecord(value: unknown): value is Record { + return typeof value === "object" && value !== null; +} + +function isNonEmptyString(value: unknown): value is string { + return typeof value === "string" && value.trim().length > 0; +} + +function parseRunTerminalRequest(body: unknown): RunTerminalRequest | null { + if (!isRecord(body)) { + return null; + } + + if (!isNonEmptyString(body.requestId) || !isNonEmptyString(body.runId)) { + return null; + } + + return { + requestId: body.requestId, + runId: body.runId, + }; +} + +function parseHandlerCompleteRequest(body: unknown): HandlerCompleteRequest | null { + if (!isRecord(body)) { + return null; + } + + if ( + !isNonEmptyString(body.requestId) || + !isNonEmptyString(body.runId) || + !isNonEmptyString(body.handlerId) || + (body.status !== "done" && body.status !== "failed") + ) { + return null; + } + + return { + requestId: body.requestId, + runId: body.runId, + handlerId: body.handlerId, + status: body.status, + }; +} + +function parseHandlerRegisterRequest(body: unknown): HandlerServiceDocument | null { + if (!isRecord(body)) { + return null; + } + + const { + _id, + type, + version, + queue, + selector, + autoBackfill, + dependsOn, + } = body; + + if ( + !isNonEmptyString(_id) || + type !== "post-process-handler" || + typeof version !== "number" || + !Number.isFinite(version) || + !isNonEmptyString(queue) || + !isNonEmptyString(selector) || + typeof autoBackfill !== "boolean" || + !Array.isArray(dependsOn) || + !dependsOn.every((d) => isNonEmptyString(d)) + ) { + return null; + } + + return { + _id, + type, + version, + queue, + selector, + autoBackfill, + dependsOn: dependsOn as string[], + }; +} + +export function readJsonBody( + req: http.IncomingMessage, + maxBytes: number = MAX_BODY_BYTES, +): Promise { + return new Promise((resolve, reject) => { + // Fast path: reject an oversized declared Content-Length before reading any + // body. Well-behaved clients (fetch/axios/http.request with a body) always + // send it, so this catches the common case without buffering a single byte. + const declaredLength = Number(req.headers["content-length"]); + if (Number.isFinite(declaredLength) && declaredLength > maxBytes) { + reject(new PayloadTooLargeError(maxBytes)); + return; + } + + const chunks: Buffer[] = []; + let total = 0; + let settled = false; + + const cleanup = (): void => { + req.removeListener("data", onData); + req.removeListener("end", onEnd); + req.removeListener("error", onError); + }; + + // Streaming guard: Content-Length may be absent (chunked) or lie, so also + // enforce the cap as bytes arrive and bail the moment it is exceeded. + const onData = (chunk: Buffer | string): void => { + if (settled) return; + const buf = typeof chunk === "string" ? Buffer.from(chunk) : chunk; + total += buf.length; + if (total > maxBytes) { + settled = true; + cleanup(); + reject(new PayloadTooLargeError(maxBytes)); + return; + } + chunks.push(buf); + }; + + const onEnd = (): void => { + if (settled) return; + settled = true; + cleanup(); + const rawBody = Buffer.concat(chunks).toString("utf-8").trim(); + if (rawBody.length === 0) { + reject(new Error("Request body is required")); + return; + } + try { + resolve(JSON.parse(rawBody) as unknown); + } catch (err) { + reject(err); + } + }; + + const onError = (err: Error): void => { + if (settled) return; + settled = true; + cleanup(); + reject(err); + }; + + req.on("data", onData); + req.on("end", onEnd); + req.on("error", onError); + }); +} + +export interface HttpServerOptions { + /** Max request body size in bytes. Defaults to MAX_BODY_BYTES (64 KB). */ + maxBodyBytes?: number; +} + +export function createHttpServer( + handler: NotifyHandler, + options: HttpServerOptions = {}, +): http.Server { + const maxBodyBytes = options.maxBodyBytes ?? MAX_BODY_BYTES; + return http.createServer(async (req, res) => { + const method = req.method ?? "GET"; + const url = new URL(req.url ?? "/", "http://localhost"); + + try { + if (method === "GET" && url.pathname === "/health") { + sendJson(res, 200, { status: "ok", service: "scheduler" }); + return; + } + + if (method === "POST" && url.pathname === "/notify/run-terminal") { + const parsedBody = parseRunTerminalRequest(await readJsonBody(req, maxBodyBytes)); + if (!parsedBody) { + sendJson(res, 400, { error: "requestId and runId are required" }); + return; + } + + await handler.onRunTerminal(parsedBody.requestId, parsedBody.runId); + sendJson(res, 200, { ok: true }); + return; + } + + if (method === "POST" && url.pathname === "/notify/handler-complete") { + const parsedBody = parseHandlerCompleteRequest(await readJsonBody(req, maxBodyBytes)); + if (!parsedBody) { + sendJson(res, 400, { + error: "requestId, runId, handlerId, and status (done|failed) are required", + }); + return; + } + + await handler.onHandlerComplete( + parsedBody.requestId, + parsedBody.runId, + parsedBody.handlerId, + parsedBody.status, + ); + sendJson(res, 200, { ok: true }); + return; + } + + if (method === "POST" && url.pathname === "/handlers/register") { + const parsedBody = parseHandlerRegisterRequest(await readJsonBody(req, maxBodyBytes)); + if (!parsedBody) { + sendJson(res, 400, { + error: + "Invalid handler registration: _id, type='post-process-handler', version (number), queue, selector, autoBackfill (boolean), and dependsOn (string[]) are required", + }); + return; + } + + await handler.registerHandler(parsedBody); + sendJson(res, 200, { ok: true, handlerId: parsedBody._id }); + return; + } + } catch (error) { + // Body was not fully consumed — close the connection so leftover bytes + // can't corrupt the next request on a keep-alive socket. + if (error instanceof PayloadTooLargeError) { + sendJson(res, 413, { error: "Request body too large" }, { Connection: "close" }); + return; + } + + const message = error instanceof Error ? error.message : "Internal server error"; + const statusCode = + message === "Request body is required" || error instanceof SyntaxError ? 400 : 500; + + sendJson(res, statusCode, { + error: + statusCode === 400 + ? "Invalid JSON request body" + : "Internal server error", + }); + return; + } + + sendJson(res, 404, { error: "Not found" }); + }); +} diff --git a/apps/workers/post-processor/Dockerfile b/apps/workers/post-processor/Dockerfile index f339e9b81..5966c79ec 100644 --- a/apps/workers/post-processor/Dockerfile +++ b/apps/workers/post-processor/Dockerfile @@ -23,6 +23,7 @@ COPY packages/telemetry/tsconfig.json ./packages/telemetry/ COPY packages/telemetry/src ./packages/telemetry/src/ COPY apps/workers/post-processor/tsconfig.json ./apps/workers/post-processor/ COPY apps/workers/post-processor/src ./apps/workers/post-processor/src/ +COPY apps/workers/post-processor/handler.yaml ./apps/workers/post-processor/ RUN pnpm --filter shared build && pnpm --filter telemetry build && pnpm --filter post-processor build @@ -58,6 +59,7 @@ RUN pnpm install --frozen-lockfile --prod COPY --from=builder /app/packages/shared/dist ./packages/shared/dist/ COPY --from=builder /app/packages/telemetry/dist ./packages/telemetry/dist/ COPY --from=builder /app/apps/workers/post-processor/dist ./apps/workers/post-processor/dist/ +COPY --from=builder /app/apps/workers/post-processor/handler.yaml ./apps/workers/post-processor/handler.yaml WORKDIR /app/apps/workers/post-processor diff --git a/apps/workers/post-processor/handler.yaml b/apps/workers/post-processor/handler.yaml new file mode 100644 index 000000000..dc3c1f348 --- /dev/null +++ b/apps/workers/post-processor/handler.yaml @@ -0,0 +1,10 @@ +# Post-process handler registration metadata. +# Read by the generic register-handler script and POSTed to the +# scheduler's POST /handlers/register endpoint on deploy. +_id: pp-atif +type: post-process-handler +version: 1 +queue: post-processor-queue +selector: atif +autoBackfill: true +dependsOn: [] diff --git a/apps/workers/post-processor/src/handlers/atif-handler.ts b/apps/workers/post-processor/src/handlers/atif-handler.ts index 5b51e208c..9a28cb604 100644 --- a/apps/workers/post-processor/src/handlers/atif-handler.ts +++ b/apps/workers/post-processor/src/handlers/atif-handler.ts @@ -5,6 +5,7 @@ import { writeFile, mkdtemp, rm } from "node:fs/promises"; import { tmpdir } from "node:os"; import { join } from "node:path"; import type { PostProcessHandler, PostProcessorMessage, HandlerContext } from "../types.js"; +import { POST_PROCESSOR_VERSION } from "../version.js"; /** * Extracts the blob name (path within the container) from a full blob storage URL. @@ -30,6 +31,8 @@ function extractBlobName(url: string, container = "snapshots"): string | null { */ export class AtifHandler implements PostProcessHandler { readonly type = "atif"; + readonly version = POST_PROCESSOR_VERSION; + readonly autoBackfill = true; async process(message: PostProcessorMessage, ctx: HandlerContext): Promise { const { requestId, runId, iteration } = message; 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..a0d0d3460 --- /dev/null +++ b/apps/workers/post-processor/src/post-processor.test.ts @@ -0,0 +1,239 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +import { beforeEach, describe, expect, it, vi } from "vitest"; +import type { RequestDocument } from "shared"; +import { PostProcessor } from "./post-processor.js"; +import type { PostProcessHandler, PostProcessorMessage } from "./types.js"; + +const CONFIG = { + mongoUri: "mongodb://localhost:27017", + mongoDatabase: "requests-db", + mongoCollection: "requests", + storageAccountName: "devstoreaccount1", + storageConnectionString: "UseDevelopmentStorage=true", + queueName: "post-processor-queue", + batchSize: 1, + pollIntervalMs: 5000, + redisHost: "localhost", + redisPort: 6379, + redisPassword: "", +}; + +function makeDoc(overrides: Partial = {}): RequestDocument { + return { + _id: "req-1", + projectId: "project", + scenario: { task: "task", criteria: [] }, + workerType: "post-processor", + createdAt: new Date(), + priority: 0, + run: { + _id: "run-1", + attemptNumber: 1, + status: "done", + handlerStatus: {}, + ...overrides, + }, + }; +} + +function makeProcessor() { + const processor = new PostProcessor(CONFIG); + const collection = { + findOneAndUpdate: vi.fn(), + updateOne: vi.fn(), + }; + const blobStorage = {}; + + (processor as any).collection = collection; + (processor as any).blobStorage = blobStorage; + + const safeDeleteMessage = vi + .spyOn(processor as any, "safeDeleteMessage") + .mockResolvedValue(undefined); + const notifyHandlerComplete = vi + .spyOn(processor as any, "notifyHandlerComplete") + .mockResolvedValue(undefined); + + return { processor, collection, blobStorage, safeDeleteMessage, notifyHandlerComplete }; +} + +function makeMessage(messageId = "message-1") { + return { + messageId, + popReceipt: "pop-0", + } as any; +} + +function makeHeartbeat() { + return { + popReceipt: "pop-0", + stop: vi.fn().mockReturnValue("pop-1"), + } as any; +} + +describe("PostProcessor", () => { + let handler: PostProcessHandler; + let log: ReturnType; + + beforeEach(() => { + handler = { + type: "atif", + version: 1, + autoBackfill: true, + process: vi.fn().mockResolvedValue(undefined), + }; + log = vi.fn().mockResolvedValue(undefined); + }); + + it("discards stale messages for an older run attempt", async () => { + const { processor, collection, safeDeleteMessage } = makeProcessor(); + processor.registerHandler(handler); + + const doc = makeDoc({ _id: "run-2" }); + const heartbeat = makeHeartbeat(); + const payload: PostProcessorMessage = { type: "atif", requestId: "req-1", runId: "run-1" }; + + await (processor as any).handleRequest(doc, makeMessage(), heartbeat, log, payload); + + expect(collection.findOneAndUpdate).not.toHaveBeenCalled(); + expect(handler.process).not.toHaveBeenCalled(); + expect(safeDeleteMessage).toHaveBeenCalledWith("message-1", "pop-1"); + }); + + it("no-ops duplicate messages when the handler is already current", async () => { + const { processor, collection, safeDeleteMessage, notifyHandlerComplete } = makeProcessor(); + processor.registerHandler(handler); + + const doc = makeDoc({ + handlerStatus: { + "pp-atif": { status: "done", version: 1, updatedAt: new Date() }, + }, + }); + const heartbeat = makeHeartbeat(); + const payload: PostProcessorMessage = { type: "atif", requestId: "req-1", runId: "run-1" }; + + collection.findOneAndUpdate.mockResolvedValue(null); + + await (processor as any).handleRequest(doc, makeMessage(), heartbeat, log, payload); + + expect(collection.findOneAndUpdate).toHaveBeenCalledWith( + expect.objectContaining({ + _id: "req-1", + "run._id": "run-1", + "run.status": "done", + "run.handlerStatus.pp-atif.status": { $nin: ["processing"] }, + $or: [ + { "run.handlerStatus.pp-atif.version": { $exists: false } }, + { "run.handlerStatus.pp-atif.version": { $lt: 1 } }, + ], + }), + expect.any(Object), + expect.objectContaining({ returnDocument: "after" }), + ); + expect(handler.process).not.toHaveBeenCalled(); + expect(collection.updateOne).not.toHaveBeenCalled(); + expect(notifyHandlerComplete).not.toHaveBeenCalled(); + expect(safeDeleteMessage).toHaveBeenCalledWith("message-1", "pop-1"); + }); + + it("re-processes stale versions and stamps success for the same run attempt", async () => { + const { processor, collection, safeDeleteMessage, notifyHandlerComplete } = makeProcessor(); + processor.registerHandler(handler); + + const doc = makeDoc({ + handlerStatus: { + "pp-atif": { status: "done", version: 0, updatedAt: new Date() }, + }, + }); + const heartbeat = makeHeartbeat(); + const payload: PostProcessorMessage = { type: "atif", requestId: "req-1", runId: "run-1" }; + + collection.findOneAndUpdate.mockResolvedValue({ _id: "req-1" }); + collection.updateOne.mockResolvedValue({ matchedCount: 1 }); + + await (processor as any).handleRequest(doc, makeMessage(), heartbeat, log, payload); + + expect(handler.process).toHaveBeenCalledWith( + payload, + expect.objectContaining({ + blobStorage: expect.any(Object), + collection, + log, + }), + ); + expect(collection.updateOne).toHaveBeenCalledWith( + { + _id: "req-1", + "run._id": "run-1", + "run.handlerStatus.pp-atif.status": "processing", + }, + { + $set: expect.objectContaining({ + "run.postProcessorVersion": 1, + "run.postProcessorStatus": "done", + "run.handlerStatus.pp-atif.status": "done", + "run.handlerStatus.pp-atif.version": 1, + "run.handlerStatus.pp-atif.updatedAt": expect.any(Date), + }), + }, + ); + expect(notifyHandlerComplete).toHaveBeenCalledWith("req-1", "run-1", "pp-atif", "done"); + expect(safeDeleteMessage).toHaveBeenCalledWith("message-1", "pop-1"); + }); + + it("claims a run left in the scheduler's 'queued' state and drives it to done", async () => { + // Regression: the scheduler writes a best-effort "queued" marker right after + // enqueue-before-claim, so a dequeued message almost always sees status + // "queued". The worker MUST be able to claim it (queued -> processing). + // Previously "queued" was in the claim exclusion list, so the worker + // discarded its own message and stranded the run at "queued" forever. + const { processor, collection, safeDeleteMessage, notifyHandlerComplete } = makeProcessor(); + processor.registerHandler(handler); + + const doc = makeDoc({ + handlerStatus: { + "pp-atif": { status: "queued", updatedAt: new Date() }, + }, + }); + const heartbeat = makeHeartbeat(); + const payload: PostProcessorMessage = { type: "atif", requestId: "req-1", runId: "run-1" }; + + collection.findOneAndUpdate.mockResolvedValue({ _id: "req-1" }); + collection.updateOne.mockResolvedValue({ matchedCount: 1 }); + + await (processor as any).handleRequest(doc, makeMessage(), heartbeat, log, payload); + + // The claim filter must NOT exclude "queued" — only "processing". + expect(collection.findOneAndUpdate).toHaveBeenCalledWith( + expect.objectContaining({ + _id: "req-1", + "run._id": "run-1", + "run.status": "done", + "run.handlerStatus.pp-atif.status": { $nin: ["processing"] }, + }), + expect.any(Object), + expect.objectContaining({ returnDocument: "after" }), + ); + expect(handler.process).toHaveBeenCalledWith( + payload, + expect.objectContaining({ collection, log }), + ); + expect(collection.updateOne).toHaveBeenCalledWith( + { + _id: "req-1", + "run._id": "run-1", + "run.handlerStatus.pp-atif.status": "processing", + }, + { + $set: expect.objectContaining({ + "run.handlerStatus.pp-atif.status": "done", + "run.handlerStatus.pp-atif.version": 1, + }), + }, + ); + expect(notifyHandlerComplete).toHaveBeenCalledWith("req-1", "run-1", "pp-atif", "done"); + expect(safeDeleteMessage).toHaveBeenCalledWith("message-1", "pop-1"); + }); +}); diff --git a/apps/workers/post-processor/src/post-processor.ts b/apps/workers/post-processor/src/post-processor.ts index 03b276db8..2b818cd8d 100644 --- a/apps/workers/post-processor/src/post-processor.ts +++ b/apps/workers/post-processor/src/post-processor.ts @@ -2,26 +2,15 @@ // Licensed under the MIT License. import type { DequeuedMessageItem } from "@azure/storage-queue"; -import { BaseQueueProcessor, BlobStorage, Retry, type BaseQueueProcessorConfig, type LogEvent, type VisibilityHeartbeat } from "shared"; +import { BaseQueueProcessor, BlobStorage, type BaseQueueProcessorConfig, type LogEvent, type VisibilityHeartbeat } from "shared"; +import type { RequestDocument, HandlerRunStatus } from "shared"; import { trackMetric, trackEvent } from "telemetry"; -import { POST_PROCESSOR_VERSION } from "./version.js"; import type { PostProcessHandler, PostProcessorMessage, HandlerContext } from "./types.js"; export interface PostProcessorConfig extends BaseQueueProcessorConfig { apiBaseUrl?: string; } -interface RequestDocument { - _id: string; - run?: { - _id: string; - status: string; - turns?: Array<{ iteration: number; harUrl?: string; atifUrl?: string }>; - postProcessorVersion?: number; - postProcessorStatus?: string; - }; -} - /** * Post-processor worker: extensible queue processor with handler registry. * Dispatched by the PostProcessorDispatcher in the scheduler when runs @@ -47,6 +36,53 @@ export class PostProcessor extends BaseQueueProcessor { console.log(`[post-processor] Registered handler: ${handler.type}`); } + private isPostProcessorMessage(payload: unknown): payload is PostProcessorMessage { + return typeof payload === "object" && + payload !== null && + typeof (payload as Record).type === "string" && + typeof (payload as Record).requestId === "string" && + typeof (payload as Record).runId === "string"; + } + + private buildHandlerClaimFilter( + requestId: string, + runId: string, + handlerId: string, + handler: PostProcessHandler, + ): Record { + // "queued" is deliberately NOT excluded: the scheduler writes a best-effort + // "queued" marker immediately after enqueueing (enqueue-before-claim), so by + // the time this message is dequeued the status is almost always "queued". + // The worker is exactly what transitions queued -> processing, so it must be + // able to claim a "queued" attempt. Excluding it would let the worker discard + // its own message and strand the run at "queued" (the scheduler poll also + // excludes "queued", so nothing would ever retry it). We only exclude + // "processing" (another worker is mid-flight) and — for non-backfill handlers + // — "done" (already complete). For autoBackfill, "done" is gated by the + // version check below instead, so a version bump can re-process it. + const statusExclusions = handler.autoBackfill + ? ["processing"] + : ["processing", "done"]; + + const filter: Record = { + _id: requestId, + "run._id": runId, + "run.status": "done", + [`run.handlerStatus.${handlerId}.status`]: { $nin: statusExclusions }, + }; + + if (handler.autoBackfill) { + filter.$or = [ + { [`run.handlerStatus.${handlerId}.version`]: { $exists: false } }, + { [`run.handlerStatus.${handlerId}.version`]: { $lt: handler.version } }, + ]; + } else { + filter[`run.handlerStatus.${handlerId}.version`] = { $exists: false }; + } + + return filter; + } + protected async handleRequest( doc: RequestDocument, message: DequeuedMessageItem, @@ -55,19 +91,24 @@ export class PostProcessor extends BaseQueueProcessor { payload?: Record, ): Promise { const processStart = Date.now(); - const msg = payload as unknown as PostProcessorMessage; if (!PostProcessor.coldStartTracked) { PostProcessor.coldStartTracked = true; - trackMetric({ name: "post_processor.cold_start_ms", value: process.uptime() * 1000, properties: { service: "post-processor" } }); + trackMetric({ + name: "post_processor.cold_start_ms", + value: process.uptime() * 1000, + properties: { service: "post-processor" }, + }); } - if (!msg?.type) { + if (!this.isPostProcessorMessage(payload)) { await log("warn", "Message missing 'type' field, discarding"); await this.safeDeleteMessage(message.messageId, heartbeat.stop()); return; } + const msg: PostProcessorMessage = payload; + const handler = this.handlers.get(msg.type); if (!handler) { await log("warn", `No handler registered for type '${msg.type}', discarding`); @@ -75,17 +116,58 @@ 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, - { $set: { "run.postProcessorStatus": "processing" } } as any, + if (doc.run?._id !== msg.runId) { + await log("info", "Stale post-processing message discarded (runId mismatch)", { + messageRunId: msg.runId, + currentRunId: doc.run?._id, + }); + await this.safeDeleteMessage(message.messageId, heartbeat.stop()); + return; + } + + const handlerId = `pp-${msg.type}`; + const currentStatus = doc.run?.handlerStatus?.[handlerId] as HandlerRunStatus | undefined; + + // This is the authoritative claim. A queued message alone does not mean the + // handler still needs work: the scheduler may have enqueued optimistically, + // another worker may already be processing it, or a fresher backfill may + // have completed it. Only the worker can safely decide that by atomically + // transitioning the specific run attempt into "processing". + const claim = await this.collection.findOneAndUpdate( + this.buildHandlerClaimFilter(doc._id, msg.runId, handlerId, handler) as any, + { + $set: { + "run.postProcessorStatus": "processing", + [`run.handlerStatus.${handlerId}.status`]: "processing", + [`run.handlerStatus.${handlerId}.updatedAt`]: new Date(), + }, + $unset: { + [`run.handlerStatus.${handlerId}.error`]: "", + }, + } as any, + { returnDocument: "after" }, ); + if (!claim) { + // No claim means the message is stale, duplicate, or already satisfied. + // Ack and drop it — at-least-once delivery is expected here. + await log("info", "Post-processing message no longer needed, discarding", { + handlerId, + runId: msg.runId, + currentStatus: currentStatus?.status, + currentVersion: currentStatus?.version, + targetVersion: handler.version, + }); + await this.safeDeleteMessage(message.messageId, heartbeat.stop()); + return; + } + try { await log("info", `Running handler: ${msg.type}`); - trackEvent({ name: "post_processor.processing_started", properties: { handlerType: msg.type, requestId: doc._id } }); + trackEvent({ + name: "post_processor.processing_started", + properties: { handlerType: msg.type, requestId: doc._id }, + }); const ctx: HandlerContext = { blobStorage: this.blobStorage, @@ -95,37 +177,72 @@ 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, + // Stamp version and status on success (legacy fields + new handlerStatus) + const successResult = await this.collection.updateOne( + { + _id: doc._id, + "run._id": msg.runId, + [`run.handlerStatus.${handlerId}.status`]: "processing", + } as any, { $set: { - "run.postProcessorVersion": POST_PROCESSOR_VERSION, + "run.postProcessorVersion": handler.version, "run.postProcessorStatus": "done", + [`run.handlerStatus.${handlerId}.status`]: "done", + [`run.handlerStatus.${handlerId}.version`]: handler.version, + [`run.handlerStatus.${handlerId}.updatedAt`]: new Date(), }, } 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 ((successResult.matchedCount ?? 0) === 0) { + await log("warn", `Handler '${msg.type}' completed after state changed; skipping terminal write`, { + handlerId, + runId: msg.runId, + }); + await this.safeDeleteMessage(message.messageId, heartbeat.stop()); + return; } + + await log("info", `Post-processing complete (v${handler.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 }, + }); + + // Notify the scheduler that this handler is done (best-effort) + await this.notifyHandlerComplete(doc._id, doc.run?._id ?? "", handlerId, "done"); } catch (err) { const errMsg = err instanceof Error ? err.message : String(err); await log("error", `Handler '${msg.type}' failed: ${errMsg}`); // Mark as failed so scheduler doesn't immediately re-dispatch - await this.collection.updateOne( - { _id: doc._id } as any, - { $set: { "run.postProcessorStatus": "failed" } } as any, + const failureResult = await this.collection.updateOne( + { + _id: doc._id, + "run._id": msg.runId, + [`run.handlerStatus.${handlerId}.status`]: "processing", + } as any, + { + $set: { + "run.postProcessorStatus": "failed", + [`run.handlerStatus.${handlerId}.status`]: "failed", + [`run.handlerStatus.${handlerId}.updatedAt`]: new Date(), + [`run.handlerStatus.${handlerId}.error`]: errMsg, + }, + } as any, ); + if ((failureResult.matchedCount ?? 0) > 0) { + // Notify scheduler of failure (best-effort — don't dispatch downstream) + await this.notifyHandlerComplete(doc._id, doc.run?._id ?? "", handlerId, "failed"); + } + throw err; } @@ -135,27 +252,31 @@ export class PostProcessor extends BaseQueueProcessor { } /** - * Trigger report generation via REST API. - * Called after post-processing succeeds so reports can use enriched data. + * Notify the scheduler that a handler completed (or failed). + * Best-effort — the scheduler's poll safety net will catch up if this fails. */ - @Retry({ maxRetries: 3, baseDelayMs: 1000, isRetryable: () => true }) - private async triggerReportGeneration(requestId: string, log: (level: LogEvent["level"], message: string) => Promise): Promise { - if (!this.apiBaseUrl) return; - - const response = await fetch(`${this.apiBaseUrl}/api/v1/reports/trigger`, { - method: "POST", - headers: { "Content-Type": "application/json" }, - body: JSON.stringify({ requestId }), - }); - if (response.ok) { - const result = await response.json() as { triggered: number }; - await log("info", `Triggered report generation: ${result.triggered} report(s) created`); - return; - } - if (response.status >= 400 && response.status < 500) { - await log("warn", `Failed to trigger report generation (${response.status}), not retrying`); - return; + private async notifyHandlerComplete( + requestId: string, + runId: string, + handlerId: string, + status: "done" | "failed", + ): Promise { + const schedulerUrl = process.env.SCHEDULER_URL; + if (!schedulerUrl) return; + + try { + const response = await fetch(`${schedulerUrl}/notify/handler-complete`, { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ requestId, runId, handlerId, status }), + }); + if (response.ok) { + console.log(`[post-processor] Notified scheduler: ${handlerId} ${status} for ${requestId}`); + } else { + console.warn(`[post-processor] Scheduler notify failed: ${response.status}`); + } + } catch (err) { + console.warn(`[post-processor] Scheduler notify error for ${handlerId}/${requestId}:`, err); } - throw new Error(`Report trigger returned ${response.status} ${response.statusText}`); } } diff --git a/apps/workers/post-processor/src/register-version.ts b/apps/workers/post-processor/src/register-version.ts deleted file mode 100644 index 6eb0ca889..000000000 --- a/apps/workers/post-processor/src/register-version.ts +++ /dev/null @@ -1,39 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// Licensed under the MIT License. - -/** - * Standalone script run by the K8s Job on each deployment. - * Writes the current post-processor version to MongoDB so the scheduler - * can detect runs needing (re-)processing. - */ -import "dotenv/config"; -import { MongoClient } from "mongodb"; -import { POST_PROCESSOR_VERSION } from "./version.js"; - -const MONGO_URI = - process.env.MONGO_CONNECTION_STRING || - process.env.AZURE_COSMOS_CONNECTION_STRING || - "mongodb://localhost:27017"; -const MONGO_DATABASE = process.env.MONGO_DATABASE || "requests-db"; - -async function main(): Promise { - const mongo = new MongoClient(MONGO_URI); - await mongo.connect(); - - await mongo - .db(MONGO_DATABASE) - .collection("services") - .updateOne( - { _id: "post-processor" } as any, - { $set: { version: POST_PROCESSOR_VERSION, updatedAt: new Date() } }, - { upsert: true }, - ); - - await mongo.close(); - console.log(`[register-version] Registered post-processor version: ${POST_PROCESSOR_VERSION}`); -} - -main().catch((err) => { - console.error("[register-version] Fatal error:", err); - process.exit(1); -}); diff --git a/apps/workers/post-processor/src/types.ts b/apps/workers/post-processor/src/types.ts index 8765dd926..e1ee29669 100644 --- a/apps/workers/post-processor/src/types.ts +++ b/apps/workers/post-processor/src/types.ts @@ -22,5 +22,7 @@ export interface HandlerContext { /** Interface for post-processing handlers (extensibility point). */ export interface PostProcessHandler { readonly type: string; + readonly version: number; + readonly autoBackfill: boolean; process(message: PostProcessorMessage, ctx: HandlerContext): Promise; } diff --git a/docker-compose.dev.yml b/docker-compose.dev.yml index 50410ab84..1050fb826 100644 --- a/docker-compose.dev.yml +++ b/docker-compose.dev.yml @@ -217,7 +217,7 @@ services: - action: rebuild path: ./pnpm-lock.yaml - register-version-post-processor: + register-handler-post-processor: build: target: dev diff --git a/docker-compose.yml b/docker-compose.yml index 31b6b191a..42fe2145e 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -35,6 +35,7 @@ x-worker-env: &worker-env <<: *common-env JUDGE_SERVICE_URL: http://judge:80 SCOPE_MT_API_URL: http://api:80 + SCHEDULER_URL: http://scheduler:8080 TOKEN_MANAGER_URL: http://token-manager:80 x-worker-deps: &worker-deps @@ -148,7 +149,7 @@ services: timeout: 5s retries: 5 - # Creates the 'har' blob container in Azurite for the gateway HAR plugin. + # Creates the 'har' blob container in Azurite. azurite-init: image: mcr.microsoft.com/azure-cli:2.68.0 restart: "no" @@ -555,6 +556,7 @@ services: SCHEDULER_TARGET_QUEUE_DEPTH: "2" SCHEDULER_POLL_INTERVAL_MS: "2000" PORT: "8080" + API_URL: http://api:80 depends_on: mongodb: condition: service_healthy @@ -867,7 +869,7 @@ services: api: condition: service_started - register-version-post-processor: + register-handler-post-processor: build: context: . dockerfile: apps/workers/post-processor/Dockerfile @@ -876,12 +878,13 @@ services: profiles: [post-processor] environment: <<: *common-env + HANDLER_YAML_PATH: apps/workers/post-processor/handler.yaml + SCHEDULER_URL: http://scheduler:8080 depends_on: - mongodb: - condition: service_healthy - required: false + scheduler: + condition: service_started restart: "no" - entrypoint: ["npx", "tsx", "apps/workers/post-processor/src/register-version.ts"] + entrypoint: ["npx", "tsx", "packages/shared/src/scripts/register-handler.ts"] post-processor: # Resource limits mirror prod K8s main container @@ -897,6 +900,7 @@ services: <<: *common-env AZURE_STORAGE_QUEUE_POSTPROCESSOR: post-processor-queue SCOPE_MT_API_URL: http://api:80 + SCHEDULER_URL: http://scheduler:8080 depends_on: mongodb: condition: service_healthy @@ -907,7 +911,7 @@ services: condition: service_healthy api: condition: service_started - register-version-post-processor: + register-handler-post-processor: condition: service_completed_successfully volumes: diff --git a/docs/architecture/db-migrations.md b/docs/architecture/db-migrations.md index 8195a2ac5..d01bcc5aa 100644 --- a/docs/architecture/db-migrations.md +++ b/docs/architecture/db-migrations.md @@ -114,6 +114,10 @@ When you run a command: | `027-uuid-keys-mcp-profileversions` | Opaque UUID `_id` + reference key for `mcp-servers` (`slug`) and `profile-versions` (`ref`) with `{projectId,slug}` / `{projectId,ref}` indexes; drops dead `prompt-feature-extractions` (see [db.md](db.md#per-project-entity-keying-migration-027)) | | `028-isolate-mcp-secrets-per-project` | Reconciles the token-manager `mcp-secrets` unique index — drops the legacy global-unique `{mcpId,name}` and (re)creates the per-project `{projectId,mcpId,name}` (see [token-manager.md](token-manager.md#mcp-secrets)) | | `029-create-users-collection` | Creates the database-enforced unique identity index for authentication; uses Cosmos collection-creation extensions for continuous-backup accounts and preserves native MongoDB index creation (see below) | +| `030-create-resource-indexes` | Creates project-scoped resource indexes | +| `031-migrate-services-to-handler-dag` | Replaces the legacy single `post-processor` service row with per-handler service registrations and dependencies | +| `032-add-handler-status-index` | Adds a handler-specific dispatch index for `pp-atif` polling | +| `033-backfill-handler-status` | Backfills `run.handlerStatus.pp-atif` from deprecated `run.postProcessorStatus` / `run.postProcessorVersion` fields | ### Migration 029: users identity uniqueness diff --git a/docs/architecture/overview.md b/docs/architecture/overview.md index 1858dfba2..1d4b6cdee 100644 --- a/docs/architecture/overview.md +++ b/docs/architecture/overview.md @@ -78,6 +78,7 @@ flowchart TB | `shared` | Shared types and utilities | | `workers/coder-acp-claude-code` | Claude Code agent via Agent Client Protocol (ACP) | | `workers/coder-acp-copilot` | GitHub Copilot agent via Agent Client Protocol (ACP) | +| `workers/report-generator` | Copilot SDK worker that generates post-run markdown reports | | `gateway` | AI Gateway — shared Rust TLS-intercepting proxy with plugin architecture (HAR capture, future: token refresh, rate limiting) | | `evaluations/static-prompts` | Developer-run static prompt quality and user-controlled prompt red-team suites | diff --git a/docs/architecture/post-processing.md b/docs/architecture/post-processing.md index 9774769e4..1370e5b7c 100644 --- a/docs/architecture/post-processing.md +++ b/docs/architecture/post-processing.md @@ -14,8 +14,9 @@ flowchart LR Two paths trigger post-processing: -1. **Event-driven (primary):** Coder workers enqueue a post-processing message immediately when a run completes, setting `run.postProcessorStatus: "queued"` atomically in the same DB write. -2. **Scheduler backfill:** The `PostProcessorDispatcher` polls every 30 seconds for runs that have `status: "done"` but no `postProcessorStatus` set (or a `postProcessorVersion` below the current target). This catches runs that completed before the event-driven dispatch was deployed, or that need re-processing after a version bump. +1. **Event-driven (primary):** Coder workers enqueue a post-processing message immediately when a run completes, currently setting the deprecated `run.postProcessorStatus: "queued"` compatibility field atomically in the same DB write while the system migrates to per-handler status. +2. **Scheduler backfill:** The `PostProcessorDispatcher` polls every 30 seconds for runs that have `status: "done"` but no post-processing marker set (or a stale handler version). During the migration window this may consult the deprecated flat fields; the long-term source of truth is `run.handlerStatus[handlerId]`. This catches runs that completed before the event-driven dispatch was deployed, or that need re-processing after a version bump. +3. **Scheduler notify & registration server:** The scheduler also exposes `GET /health`, `POST /notify/run-terminal`, `POST /notify/handler-complete`, and `POST /handlers/register` on a lightweight `node:http` server. Workers and post-process handlers use the notify endpoints to signal terminal states so the scheduler can re-evaluate which root or downstream handlers are ready to dispatch. The `/handlers/register` endpoint is the single writer to the `services` collection — see [Handler Registration](#handler-registration). All POST endpoints cap the request body at 64 KB (`MAX_BODY_BYTES` in `notify-routes.ts`); oversized payloads are rejected with `413 Payload Too Large` and `Connection: close` before the body is fully buffered, guarding the singleton scheduler against unbounded-memory requests. ## ATIF Generation @@ -32,8 +33,8 @@ Two paths trigger post-processing: - Uploads `atif.trajectory.json` to blob storage at `{requestId}/runs/{runId}/iteration-{N}/atif.trajectory.json` - Updates `run.turns.$.atifUrl` in the request document -4. On success, stamps `run.postProcessorVersion` and sets `run.postProcessorStatus: "done"` -5. Triggers report generation via the API (best-effort, with retry) +4. On success, stamps `run.handlerStatus["pp-atif"] = { status: "done", version: N }` (while legacy `run.postProcessorVersion` / `run.postProcessorStatus` may still be mirrored during the transition) +5. Triggers any downstream handlers once the scheduler sees the dependency graph is satisfied. (Reports are **not** a DAG handler — see [Report Triggering](#report-triggering-on-dag-drain).) ### Blob Storage Layout @@ -53,7 +54,7 @@ Two paths trigger post-processing: | `GET /api/v1/requests/:id/atif?iteration=N` | Download ATIF for the latest run | | `GET /api/v1/requests/:id/runs/:runId/atif?iteration=N` | Download ATIF for a specific run | -The `iteration` query parameter is mandatory. ATIF files are also included in archive exports as `iteration-{N}.atif.trajectory.json`. +The `iteration` query parameter is mandatory for the ATIF endpoints. ATIF files are also included in archive exports as `iteration-{N}.atif.trajectory.json`. ## Handler Interface @@ -62,6 +63,8 @@ The post-processor uses a registry pattern for extensibility: ```typescript interface PostProcessHandler { readonly type: string; + readonly version: number; + readonly autoBackfill: boolean; process(message: PostProcessorMessage, ctx: HandlerContext): Promise; } @@ -87,6 +90,23 @@ processor.registerHandler(new AtifHandler()); processor.start(); ``` +## Idempotency and Dispatch Semantics + +Post-processing delivery is treated as **at-least-once**. Queue messages may be duplicated or redelivered, and the scheduler may enqueue a handler message before the corresponding `queued` marker is durably written back to MongoDB. + +To keep this safe, idempotency is enforced at the **post-processor worker entrypoint**, not inside each individual handler: + +1. The scheduler decides whether work is needed from `run.handlerStatus[handlerId]` + handler metadata (`version`, `autoBackfill`). +2. The worker re-checks that same eligibility when a message is dequeued, scoped to the specific `runId`. +3. The worker atomically claims the handler by transitioning its status to `processing`. The claim filter accepts both the *absent* state and the `queued` marker as valid source states — it only excludes `processing` (another worker is mid-flight) and, for non-backfill handlers, `done` (already complete; for `autoBackfill` handlers a stale `done` is re-claimable via the version gate). Because the scheduler writes the best-effort `queued` marker in-process immediately after enqueueing (enqueue-before-claim), a dequeued message almost always sees `queued`; excluding `queued` from the claim filter would make the worker discard its own message and strand the run at `queued` forever (the 30s backfill poll also skips `queued`), so `queued → processing` **must** be a permitted worker transition. +4. If the message is stale, duplicated, or the handler is already current, the worker deletes the queue message and no-ops. + +This means: + +- A crash after queue send but before the scheduler writes `queued` no longer wedges the run forever. +- Duplicate queue messages are harmless. +- Handler implementations can stay simple; correctness does not depend on each handler re-implementing duplicate guards. + ### Adding a New Handler 1. Create a new file in `apps/workers/post-processor/src/handlers/` @@ -105,13 +125,98 @@ The `AtifHandler` has two complementary test suites: The post-processor uses a version-based re-processing scheme: -- The worker declares its version in `src/version.ts` (e.g., `export const POST_PROCESSOR_VERSION = 1`) -- On each deploy, a registration Job writes `{ _id: "post-processor", version: N }` to the `services` collection -- The scheduler's `PostProcessorDispatcher` compares each request's `run.postProcessorVersion` against the registered version -- Requests with a lower version are re-dispatched for processing with the new handler logic +- Each handler is registered independently in the `services` collection using a `post-process-handler` document keyed by handler ID (for example `pp-atif`) +- Each registration stores the handler `version`, queue name, selector, `autoBackfill` behavior, and `dependsOn` DAG edges +- The scheduler compares each request's per-handler status/version state (for example `run.handlerStatus["pp-atif"]`) against the registered handler entry for that selector +- Requests with missing or outdated handler state are re-dispatched for that specific handler once its dependencies are satisfied This allows deploying handler improvements and having them automatically applied to historical data. +## Handler Registration + +Each post-process handler declares its DAG metadata in a `handler.yaml` file colocated with the worker (e.g. `apps/workers/post-processor/handler.yaml`). This mirrors the coding-agent `agent.yaml` registration pattern. + +```yaml +# apps/workers/post-processor/handler.yaml +_id: pp-atif +type: post-process-handler +version: 1 +queue: post-processor-queue +selector: atif +autoBackfill: true +dependsOn: [] +``` + +On deploy, a generic registration script — `packages/shared/src/scripts/register-handler.ts` — reads the `handler.yaml` pointed to by `HANDLER_YAML_PATH`, waits for the scheduler's `/health`, and POSTs the parsed document to `${SCHEDULER_URL}/handlers/register` (with retry). The scheduler validates the document, runs a cycle check against the existing topology, and upserts it into the `services` collection via `HandlerDispatcher.registerHandler`. **Workers never write to the `services` collection directly** — the scheduler is the sole owner. + +Because handler topology is read on the dispatch hot path (every notify reads it 2–3 times, and every poll cycle once) but only changes on deploys, the scheduler caches it in-memory with a TTL (default 5 min, `SCHEDULER_HANDLER_CACHE_TTL_MS`). Coherence comes primarily from invalidation, not expiry: `registerHandler` invalidates the cache immediately after upserting, and since the scheduler is a singleton (`replicas: 1`, `Recreate`) and the sole writer of the `services` collection, a change is reflected at once regardless of TTL. The TTL is therefore only a defensive backstop for out-of-band writes or a future scale-out to more than one replica, and is deliberately kept well above the poll interval so the poll loop doesn't force a fresh read every cycle. Cycle-check validation always reads fresh (bypassing the cache). + +> **Limitation — the topology cache assumes a single scheduler replica.** In-memory cache invalidation is process-local, so it is only correct while the scheduler runs as one replica. This holds today (`deploy/base/scheduler.yaml`: `replicas: 1`, `strategy: Recreate`, so pods never overlap). **If the scheduler is ever scaled beyond one replica, this cache becomes stale-prone:** a `registerHandler` on one replica invalidates only that replica's cache, so the others keep serving old topology for up to the TTL. The blast radius is bounded — with idempotent dispatch a newly added handler is still dispatched by whichever replica has fresh topology, and a removed/re-wired handler self-heals within the TTL — but it is real. Before scaling out, replace process-local invalidation with a cross-replica signal (e.g. Redis pub/sub on `registerHandler`, or a change-stream watcher on the `services` collection), or lower `SCHEDULER_HANDLER_CACHE_TTL_MS` to shrink the staleness window. + +```mermaid +flowchart LR + YAML[handler.yaml] --> SCRIPT[register-handler.ts] + SCRIPT -->|"POST /handlers/register"| SCH[Scheduler] + SCH -->|"validate + cycle check + upsert"| SVC[(services collection)] +``` + +**Where it runs:** + +- **Docker Compose:** one `register-handler-` init service per handler (e.g. `pp-atif`) runs the generic script via `tsx`, gated on the scheduler being up. The worker waits for its registration service to complete. +- **Kubernetes:** a `register-handler-post-processor` Job runs the compiled script (`node packages/shared/dist/scripts/register-handler.js`) on each deploy. Only `pp-atif` is registered in K8s today because it is the only post-process handler deployed to the cluster. The scheduler exposes a `Service` (`scheduler.scoped.svc.cluster.local:8080`) so Jobs and workers can reach `/handlers/register` and the notify endpoints. + +To **add or change a handler**, edit its `handler.yaml` (bump `version`, adjust `dependsOn`, etc.) — no code or migration changes are needed for registration. + +## Report Triggering on DAG Drain + +Reports are intentionally **not** modeled as a DAG handler. The report-generator +worker is keyed by `reportId` and requires a `templateId`, whereas the DAG +dispatches a generic `{ type, requestId, runId }` message — so a report node +could never process a DAG message. Reports are also fundamentally different from +post-process handlers: they are parameterized (per `templateId`), N-per-run (one +per matching template), and trigger-gated (each template has its own +`ReportTrigger`). + +Instead, the scheduler watches for a run's handler DAG to **drain** and then +calls the existing `POST /api/v1/reports/trigger` endpoint, which evaluates every +active template's trigger and creates the properly-templated `ReportDocument`(s). + +```mermaid +flowchart LR + RT["run done"] --> DAG["scheduler dispatches DAG
(pp-atif → …)"] + DAG -->|"all handlers terminal"| DR{DAG drained?} + DR -->|yes| TRIG["POST /api/v1/reports/trigger
{ requestId }"] + TRIG --> API["API evaluates each template trigger
→ ReportDocument + report-queue {reportId}"] + API --> RG["report-generator (unchanged)"] +``` + +**Drain definition.** A handler is _terminal_ when its status is `done` or +`failed`, **or** when its status is absent but a transitive dependency `failed` +(so it can never run — effectively blocked). The DAG is drained when every +registered handler is terminal. A run with **zero** registered handlers is +vacuously drained, so reports trigger immediately on run completion. + +**Trigger regardless of handler success.** A flaky handler must not block report +generation (reports analyze the run/ATIF), so the scheduler triggers reports +whenever the DAG drains — not only when every handler succeeded. + +**Exactly-once, attempt-scoped.** Before POSTing, the scheduler atomically claims +a per-attempt guard (`run.reportsTriggeredAt`) via `findOneAndUpdate` filtered by +`{ _id: requestId, "run._id": runId, "run.status": "done", "run.reportsTriggeredAt": { $exists: false } }`. +Scoping by `run._id` ensures a stale notify or a concurrent retry can never fire +reports for — or roll back the guard of — a different attempt. If the POST fails, +the guard is rolled back (also `run._id`-scoped) so the 30s poll safety net +retries. + +**Where it fires.** `maybeTriggerReports(requestId, runId)` runs at the end of +`onRunTerminal` (handles zero-handler / already-drained runs), both branches of +`onHandlerComplete` (success **and** the failure early-return, since a failure can +drain the DAG by blocking descendants), and a `pollTriggerReports()` scan in the +scheduler's 30s poll loop (durable backstop for dropped notifications). + +The scheduler reaches the API via the `API_URL` env var (`http://api:80` in +Compose, `http://api.scoped.svc.cluster.local:80` in K8s). + ## Status Lifecycle ```mermaid @@ -125,7 +230,7 @@ stateDiagram-v2 | Status | Meaning | | --- | --- | | *(absent)* | Run hasn't been dispatched for post-processing yet | -| `queued` | Message sent to queue, awaiting pickup | +| `queued` | Message sent to queue, awaiting pickup (best-effort marker; the worker entrypoint is authoritative) | | `processing` | Worker is actively processing | | `done` | Post-processing completed successfully | | `failed` | Handler threw an error (may be retried by scheduler on next poll) | @@ -134,16 +239,20 @@ stateDiagram-v2 | Environment Variable | Default | Description | | --- | --- | --- | -| `AZURE_STORAGE_QUEUE_POSTPROCESSOR` | `post-processor-queue` | Queue name for post-processor messages | +| `AZURE_STORAGE_QUEUE_POSTPROCESSOR` | `post-processor-queue` | Queue name for ATIF post-processor messages | | `SCHEDULER_PP_POLL_INTERVAL_MS` | `30000` | Scheduler backfill polling interval | | `BATCH_SIZE` | `1` | Messages to process per poll (worker-side) | | `POLL_INTERVAL_MS` | `5000` | Worker queue polling interval | -| `SCOPE_MT_API_URL` / `API_BASE_URL` | `http://localhost:3001` | API URL for report generation trigger | +| `SCOPE_MT_API_URL` / `API_BASE_URL` | `http://localhost:3001` | API URL used by report Copilot tools | +| `SESSION_TIMEOUT_MS` | `300000` | Timeout for report Copilot SDK sessions | +| `SCHEDULER_URL` | _(unset)_ | Scheduler base URL — used by handlers to notify completion and by the register-handler script to reach `/handlers/register` | +| `API_URL` | `http://api:80` | API base URL the **scheduler** uses to call `POST /api/v1/reports/trigger` when a run's handler DAG drains | +| `HANDLER_YAML_PATH` | _(unset)_ | Path to a handler's `handler.yaml`, read by the generic `register-handler` script | ## Infrastructure -- **Queue:** Azure Storage Queue (`post-processor-queue`) +- **Queues:** Azure Storage Queues (`post-processor-queue` for ATIF, `report-queue` for reports) - **Blob Storage:** `snapshots` container (shared with HAR, tool-calls, and other artifacts) -- **Database:** MongoDB `requests` collection (`run.postProcessorStatus`, `run.postProcessorVersion` fields) +- **Database:** MongoDB `requests` collection (`run.handlerStatus` map keyed by handler ID, with deprecated `run.postProcessorStatus` / `run.postProcessorVersion` compatibility fields during migration) - **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 +- **Handler registration:** declarative `handler.yaml` per worker → generic `register-handler` script → scheduler `POST /handlers/register` (sole writer to the `services` collection). A Kubernetes Job runs the script on each deploy; Docker Compose uses per-handler init services. See [Handler Registration](#handler-registration). \ No newline at end of file diff --git a/packages/db-migrations/src/check-migrations.test.ts b/packages/db-migrations/src/check-migrations.test.ts index 56e238ece..7a7a07b63 100644 --- a/packages/db-migrations/src/check-migrations.test.ts +++ b/packages/db-migrations/src/check-migrations.test.ts @@ -7,6 +7,7 @@ import { resetMigrationCheckCache, type MigrationCheckResult, } from "./check-migrations.js"; +import { REQUIRED_MIGRATIONS } from "./required-migrations.js"; // Mock Db function makeMockDb(appliedFiles: string[]) { @@ -30,81 +31,19 @@ describe("checkMigrations", () => { }); it("returns ready when all required migrations are applied", async () => { - const db = makeMockDb([ - "001-backfill-task-prompts.ts", - "002-create-indexes.ts", - "003-create-skill-indexes.ts", - "004-add-submission-id-index.ts", - "005-backfill-iteration-durations.ts", - "006-split-status-outcome.ts", - "007-rename-exhausted-to-finished.ts", - "008-backfill-ai-call-count.ts", - "009-add-requests-filter-indexes.ts", - "010-add-requests-pagination-index.ts", - "011-add-profile-indexes.ts", - "012-add-profile-name-index.ts", - "013-remove-logs-from-docs.ts", - "014-introduce-runs-and-run.ts", - "015-add-priority-and-scheduler-index.ts", - "016-fix-scheduler-sort-index.ts", - "017-add-post-processor-dispatch-index.ts", - "018-backfill-criteria-gates.ts", - "019-backfill-task-prompt-type.ts", - "020-create-codebase-indexes.ts", - "021-add-runs-filter-indexes.ts", - "022-add-runs-sort-indexes.ts", - "023-add-runs-search-task-index.ts", - "024-add-criteria-sort-index.ts", - "025-create-projects.ts", - "026-isolate-catalogs-per-project.ts", - "027-uuid-keys-mcp-profileversions.ts", - "028-isolate-mcp-secrets-per-project.ts", - "029-create-users-collection.ts", - "030-create-resource-indexes.ts", - ]); + const db = makeMockDb([...REQUIRED_MIGRATIONS]); const result = await checkMigrations(db); expect(result.ready).toBe(true); expect(result.pending).toEqual([]); - expect(result.totalApplied).toBe(30); - expect(result.applied).toEqual([ - "001-backfill-task-prompts.ts", - "002-create-indexes.ts", - "003-create-skill-indexes.ts", - "004-add-submission-id-index.ts", - "005-backfill-iteration-durations.ts", - "006-split-status-outcome.ts", - "007-rename-exhausted-to-finished.ts", - "008-backfill-ai-call-count.ts", - "009-add-requests-filter-indexes.ts", - "010-add-requests-pagination-index.ts", - "011-add-profile-indexes.ts", - "012-add-profile-name-index.ts", - "013-remove-logs-from-docs.ts", - "014-introduce-runs-and-run.ts", - "015-add-priority-and-scheduler-index.ts", - "016-fix-scheduler-sort-index.ts", - "017-add-post-processor-dispatch-index.ts", - "018-backfill-criteria-gates.ts", - "019-backfill-task-prompt-type.ts", - "020-create-codebase-indexes.ts", - "021-add-runs-filter-indexes.ts", - "022-add-runs-sort-indexes.ts", - "023-add-runs-search-task-index.ts", - "024-add-criteria-sort-index.ts", - "025-create-projects.ts", - "026-isolate-catalogs-per-project.ts", - "027-uuid-keys-mcp-profileversions.ts", - "028-isolate-mcp-secrets-per-project.ts", - "029-create-users-collection.ts", - "030-create-resource-indexes.ts", - ]); + expect(result.totalApplied).toBe(REQUIRED_MIGRATIONS.length); + expect(result.applied).toEqual(REQUIRED_MIGRATIONS); }); it("returns not ready when migrations are missing", async () => { const db = makeMockDb(["001-backfill-task-prompts.ts"]); const result = await checkMigrations(db); expect(result.ready).toBe(false); - expect(result.pending).toEqual(["002-create-indexes.ts", "003-create-skill-indexes.ts", "004-add-submission-id-index.ts", "005-backfill-iteration-durations.ts", "006-split-status-outcome.ts", "007-rename-exhausted-to-finished.ts", "008-backfill-ai-call-count.ts", "009-add-requests-filter-indexes.ts", "010-add-requests-pagination-index.ts", "011-add-profile-indexes.ts", "012-add-profile-name-index.ts", "013-remove-logs-from-docs.ts", "014-introduce-runs-and-run.ts", "015-add-priority-and-scheduler-index.ts", "016-fix-scheduler-sort-index.ts", "017-add-post-processor-dispatch-index.ts", "018-backfill-criteria-gates.ts", "019-backfill-task-prompt-type.ts", "020-create-codebase-indexes.ts", "021-add-runs-filter-indexes.ts", "022-add-runs-sort-indexes.ts", "023-add-runs-search-task-index.ts", "024-add-criteria-sort-index.ts", "025-create-projects.ts", "026-isolate-catalogs-per-project.ts", "027-uuid-keys-mcp-profileversions.ts", "028-isolate-mcp-secrets-per-project.ts", "029-create-users-collection.ts", "030-create-resource-indexes.ts"]); + expect(result.pending).toEqual(REQUIRED_MIGRATIONS.slice(1)); expect(result.applied).toEqual(["001-backfill-task-prompts.ts"]); expect(result.totalApplied).toBe(1); }); @@ -113,80 +52,20 @@ describe("checkMigrations", () => { const db = makeMockDb([]); const result = await checkMigrations(db); expect(result.ready).toBe(false); - expect(result.pending).toEqual([ - "001-backfill-task-prompts.ts", - "002-create-indexes.ts", - "003-create-skill-indexes.ts", - "004-add-submission-id-index.ts", - "005-backfill-iteration-durations.ts", - "006-split-status-outcome.ts", - "007-rename-exhausted-to-finished.ts", - "008-backfill-ai-call-count.ts", - "009-add-requests-filter-indexes.ts", - "010-add-requests-pagination-index.ts", - "011-add-profile-indexes.ts", - "012-add-profile-name-index.ts", - "013-remove-logs-from-docs.ts", - "014-introduce-runs-and-run.ts", - "015-add-priority-and-scheduler-index.ts", - "016-fix-scheduler-sort-index.ts", - "017-add-post-processor-dispatch-index.ts", - "018-backfill-criteria-gates.ts", - "019-backfill-task-prompt-type.ts", - "020-create-codebase-indexes.ts", - "021-add-runs-filter-indexes.ts", - "022-add-runs-sort-indexes.ts", - "023-add-runs-search-task-index.ts", - "024-add-criteria-sort-index.ts", - "025-create-projects.ts", - "026-isolate-catalogs-per-project.ts", - "027-uuid-keys-mcp-profileversions.ts", - "028-isolate-mcp-secrets-per-project.ts", - "029-create-users-collection.ts", - "030-create-resource-indexes.ts", - ]); + expect(result.pending).toEqual(REQUIRED_MIGRATIONS); expect(result.applied).toEqual([]); expect(result.totalApplied).toBe(0); }); it("ignores extra applied migrations not in the required list", async () => { const db = makeMockDb([ - "001-backfill-task-prompts.ts", - "002-create-indexes.ts", - "003-create-skill-indexes.ts", - "004-add-submission-id-index.ts", - "005-backfill-iteration-durations.ts", - "006-split-status-outcome.ts", - "007-rename-exhausted-to-finished.ts", - "008-backfill-ai-call-count.ts", - "009-add-requests-filter-indexes.ts", - "010-add-requests-pagination-index.ts", - "011-add-profile-indexes.ts", - "012-add-profile-name-index.ts", - "013-remove-logs-from-docs.ts", - "014-introduce-runs-and-run.ts", - "015-add-priority-and-scheduler-index.ts", - "016-fix-scheduler-sort-index.ts", - "017-add-post-processor-dispatch-index.ts", - "018-backfill-criteria-gates.ts", - "019-backfill-task-prompt-type.ts", - "020-create-codebase-indexes.ts", - "021-add-runs-filter-indexes.ts", - "022-add-runs-sort-indexes.ts", - "023-add-runs-search-task-index.ts", - "024-add-criteria-sort-index.ts", - "025-create-projects.ts", - "026-isolate-catalogs-per-project.ts", - "027-uuid-keys-mcp-profileversions.ts", - "028-isolate-mcp-secrets-per-project.ts", - "029-create-users-collection.ts", - "030-create-resource-indexes.ts", + ...REQUIRED_MIGRATIONS, "999-future-migration.ts", ]); const result = await checkMigrations(db); expect(result.ready).toBe(true); expect(result.pending).toEqual([]); - expect(result.totalApplied).toBe(31); + expect(result.totalApplied).toBe(REQUIRED_MIGRATIONS.length + 1); }); it("caches results within TTL", async () => { diff --git a/packages/db-migrations/src/migrations/031-migrate-services-to-handler-dag.ts b/packages/db-migrations/src/migrations/031-migrate-services-to-handler-dag.ts new file mode 100644 index 000000000..570b12e40 --- /dev/null +++ b/packages/db-migrations/src/migrations/031-migrate-services-to-handler-dag.ts @@ -0,0 +1,42 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +import type { Db } from "mongodb"; +import type { MigrationInterface } from "mongo-migrate-ts"; + +/** + * Migration 031: Remove the legacy "post-processor" service entry. + * + * Handler service registration (e.g. pp-atif) is now owned + * by each worker's handler.yaml + the generic register-handler script, + * which POSTs to the scheduler's POST /handlers/register endpoint on deploy. + * Reports are not a DAG handler — the scheduler triggers report generation + * (POST /api/v1/reports/trigger) once the DAG drains. + * This migration only cleans up the old single-entry format. + */ +export class MigrateServicesToHandlerDag implements MigrationInterface { + async up(db: Db): Promise { + const services = db.collection("services"); + await services.deleteOne({ _id: "post-processor" } as any); + console.log("[031] Deleted legacy post-processor service entry (handler registration now owned by workers)"); + } + + async down(db: Db): Promise { + const services = db.collection("services"); + await services.updateOne( + { _id: "post-processor" } as any, + { + $set: { version: 2 }, + $unset: { + type: "", + queue: "", + selector: "", + autoBackfill: "", + dependsOn: "", + }, + }, + { upsert: true }, + ); + console.log("[031-down] Restored legacy post-processor service entry"); + } +} diff --git a/packages/db-migrations/src/migrations/032-add-handler-status-index.ts b/packages/db-migrations/src/migrations/032-add-handler-status-index.ts new file mode 100644 index 000000000..44d7a448a --- /dev/null +++ b/packages/db-migrations/src/migrations/032-add-handler-status-index.ts @@ -0,0 +1,56 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +import type { Db } from "mongodb"; +import type { MigrationInterface } from "mongo-migrate-ts"; + +const HANDLER_STATUS_INDEXES = [ + { + name: "idx_handler_dispatch_atif", + key: { + "run.handlerStatus.pp-atif.status": 1 as const, + "run.status": 1 as const, + updatedAt: -1 as const, + }, + }, +] as const; + +export class AddHandlerStatusIndex implements MigrationInterface { + async up(db: Db): Promise { + const requests = db.collection("requests"); + + for (const index of HANDLER_STATUS_INDEXES) { + try { + const createdName = await requests.createIndex(index.key, { name: index.name }); + const exists = await requests.indexExists(createdName); + console.log( + exists + ? `[032] Created ${index.name} index on requests` + : `[032] ${index.name} was requested but is not visible in getIndexes()`, + ); + } catch (err: unknown) { + const message = err instanceof Error ? err.message : String(err); + console.log(`[032] ${index.name} already exists or couldn't be created: ${message}`); + } + } + } + + async down(db: Db): Promise { + const requests = db.collection("requests"); + + for (const index of HANDLER_STATUS_INDEXES) { + try { + await requests.dropIndex(index.name); + console.log(`[032-down] Dropped ${index.name} index`); + } catch (err: unknown) { + const codeName = typeof err === "object" && err !== null && "codeName" in err + ? String((err as { codeName?: unknown }).codeName) + : undefined; + if (codeName !== "IndexNotFound") { + const message = err instanceof Error ? err.message : String(err); + console.log(`[032-down] ${index.name} not found: ${message}`); + } + } + } + } +} diff --git a/packages/db-migrations/src/migrations/033-backfill-handler-status.ts b/packages/db-migrations/src/migrations/033-backfill-handler-status.ts new file mode 100644 index 000000000..ac5326afd --- /dev/null +++ b/packages/db-migrations/src/migrations/033-backfill-handler-status.ts @@ -0,0 +1,76 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +import type { AnyBulkWriteOperation, Db } from "mongodb"; +import type { MigrationInterface } from "mongo-migrate-ts"; + +const BATCH_SIZE = 100; + +type RequestWithPostProcessorState = { + _id: string; + run?: { + postProcessorStatus?: "queued" | "processing" | "done" | "failed"; + postProcessorVersion?: number; + }; +}; + +export class BackfillHandlerStatus implements MigrationInterface { + async up(db: Db): Promise { + const requests = db.collection("requests"); + const cursor = requests.find( + { "run.postProcessorVersion": { $exists: true } }, + { + projection: { + _id: 1, + "run.postProcessorStatus": 1, + "run.postProcessorVersion": 1, + }, + batchSize: BATCH_SIZE, + }, + ); + + let processed = 0; + let modified = 0; + let batch: AnyBulkWriteOperation[] = []; + + for await (const doc of cursor) { + const version = doc.run?.postProcessorVersion; + if (version === undefined) { + continue; + } + + batch.push({ + updateOne: { + filter: { _id: doc._id, "run.postProcessorVersion": { $exists: true } }, + update: { + $set: { + "run.handlerStatus.pp-atif": { + status: doc.run?.postProcessorStatus ?? "done", + version, + }, + }, + }, + }, + }); + + if (batch.length >= BATCH_SIZE) { + const result = await requests.bulkWrite(batch, { ordered: false }); + processed += batch.length; + modified += result.modifiedCount; + batch = []; + } + } + + if (batch.length > 0) { + const result = await requests.bulkWrite(batch, { ordered: false }); + processed += batch.length; + modified += result.modifiedCount; + } + + console.log(` Backfilled handlerStatus for ${modified} runs (processed ${processed})`); + } + + async down(_db: Db): Promise { + // No-op: deprecated flat fields remain intact for rollback compatibility. + } +} diff --git a/packages/db-migrations/src/required-migrations.ts b/packages/db-migrations/src/required-migrations.ts index 3fd0899ed..b5e4d29ce 100644 --- a/packages/db-migrations/src/required-migrations.ts +++ b/packages/db-migrations/src/required-migrations.ts @@ -43,4 +43,7 @@ export const REQUIRED_MIGRATIONS: readonly string[] = [ "028-isolate-mcp-secrets-per-project.ts", "029-create-users-collection.ts", "030-create-resource-indexes.ts", + "031-migrate-services-to-handler-dag.ts", + "032-add-handler-status-index.ts", + "033-backfill-handler-status.ts", ]; diff --git a/packages/shared/src/queue/queue-processor.test.ts b/packages/shared/src/queue/queue-processor.test.ts index 6a7ed97da..43da3385c 100644 --- a/packages/shared/src/queue/queue-processor.test.ts +++ b/packages/shared/src/queue/queue-processor.test.ts @@ -554,8 +554,9 @@ describe("CodingAgentQueueProcessor pre-processing failures", () => { }); // ─── Post-processor event-driven dispatch ──────────────────────────────────── -describe("CodingAgentQueueProcessor.enqueuePostProcessing", () => { - it("sends a post-processor queue message when postProcessorQueueName is configured", async () => { +describe("CodingAgentQueueProcessor.notifyRunTerminal", () => { + it("sends a post-processor queue message via legacy fallback when SCHEDULER_URL is not set", async () => { + delete process.env.SCHEDULER_URL; const configWithPP: QueueProcessorConfig = { ...testConfig, postProcessorQueueName: "post-processor-queue", @@ -565,7 +566,7 @@ describe("CodingAgentQueueProcessor.enqueuePostProcessing", () => { const sendMessage = vi.fn().mockResolvedValue({}); (qp as any).postProcessorQueueClient = { sendMessage }; - await (qp as any).enqueuePostProcessing("req-123", "run-456"); + await (qp as any).notifyRunTerminal("req-123", "run-456"); expect(sendMessage).toHaveBeenCalledTimes(1); const decoded = JSON.parse( @@ -574,13 +575,15 @@ describe("CodingAgentQueueProcessor.enqueuePostProcessing", () => { expect(decoded).toEqual({ type: "atif", requestId: "req-123", runId: "run-456" }); }); - it("does nothing when postProcessorQueueClient is null", async () => { + it("does nothing when postProcessorQueueClient is null and SCHEDULER_URL not set", async () => { + delete process.env.SCHEDULER_URL; const qp = new CodingAgentQueueProcessor(testConfig, stubProcessor); expect((qp as any).postProcessorQueueClient).toBeNull(); - await (qp as any).enqueuePostProcessing("req-123", "run-456"); + await (qp as any).notifyRunTerminal("req-123", "run-456"); }); it("logs a warning but does not throw on queue send failure", async () => { + delete process.env.SCHEDULER_URL; const configWithPP: QueueProcessorConfig = { ...testConfig, postProcessorQueueName: "post-processor-queue", @@ -591,7 +594,7 @@ describe("CodingAgentQueueProcessor.enqueuePostProcessing", () => { (qp as any).postProcessorQueueClient = { sendMessage }; const warnSpy = vi.spyOn(console, "warn").mockImplementation(() => {}); - await (qp as any).enqueuePostProcessing("req-123", "run-456"); + await (qp as any).notifyRunTerminal("req-123", "run-456"); expect(warnSpy).toHaveBeenCalledWith( expect.stringContaining("Failed to enqueue post-processing"), expect.any(Error), diff --git a/packages/shared/src/queue/queue-processor.ts b/packages/shared/src/queue/queue-processor.ts index 2be51bf18..a37f9a4b5 100644 --- a/packages/shared/src/queue/queue-processor.ts +++ b/packages/shared/src/queue/queue-processor.ts @@ -95,7 +95,7 @@ export class CodingAgentQueueProcessor extends BaseQueueProcessor { + private async notifyRunTerminal(requestId: string, runId: string): Promise { + const schedulerUrl = process.env.SCHEDULER_URL; + if (schedulerUrl) { + try { + const response = await fetch(`${schedulerUrl}/notify/run-terminal`, { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ requestId, runId }), + }); + if (response.ok) { + console.log(`[${this.workerName}] Notified scheduler: run-terminal for ${requestId}`); + } else { + console.warn(`[${this.workerName}] Scheduler notify failed: ${response.status} — poll will catch up`); + } + return; + } catch (err) { + console.warn(`[${this.workerName}] Scheduler notify error for ${requestId} — poll will catch up:`, err); + return; + } + } + + // Legacy fallback: direct queue enqueue + await this.enqueuePostProcessingLegacy(requestId, runId); + } + + /** + * @deprecated Legacy direct queue enqueue. Use notifyRunTerminal instead. + * Kept as fallback when SCHEDULER_URL is not configured. + */ + private async enqueuePostProcessingLegacy(requestId: string, runId: string): Promise { if (!this.postProcessorQueueClient) return; try { const message = Buffer.from( @@ -130,9 +160,9 @@ export class CodingAgentQueueProcessor extends BaseQueueProcessor { - await this.enqueuePostProcessing(requestId, runId); + await this.notifyRunTerminal(requestId, runId); } /** Build workerVersion and OS fields for stamping on request documents. @@ -357,8 +387,8 @@ export class CodingAgentQueueProcessor extends BaseQueueProcessor { - const apiBaseUrl = (this.config as QueueProcessorConfig).apiBaseUrl; - if (!apiBaseUrl) return; - - try { - const response = await fetch(`${apiBaseUrl}/api/v1/reports/trigger`, { - method: "POST", - headers: { "Content-Type": "application/json" }, - body: JSON.stringify({ requestId }), - }); - if (response.ok) { - const result = await response.json() as { triggered: number }; - console.log(`[${this.workerName}] Triggered report generation for request ${requestId}: ${result.triggered} report(s) created`); - } else { - console.warn(`[${this.workerName}] Failed to trigger report generation: ${response.status} ${response.statusText}`); - } - } catch (error) { - console.warn(`[${this.workerName}] Failed to trigger report generation: ${error}`); - } - } - /** * Extract skill archives to the workspace filesystem so agents can discover them. * Must be called after setup() so that processor.workspacePath points to the @@ -971,8 +975,8 @@ export class CodingAgentQueueProcessor extends BaseQueueProcessor { + const valid = { + _id: "pp-example", + type: "post-process-handler", + version: 1, + queue: "pp-example-queue", + selector: "example", + autoBackfill: false, + dependsOn: ["pp-atif"], + }; + + it("accepts a valid handler document", () => { + expect(parseHandlerDocument(valid)).toEqual(valid); + }); + + it("defaults dependsOn to an empty array when omitted", () => { + const { dependsOn, ...withoutDeps } = valid; + expect(parseHandlerDocument(withoutDeps).dependsOn).toEqual([]); + }); + + it("rejects a missing _id", () => { + expect(() => parseHandlerDocument({ ...valid, _id: "" })).toThrow(/_id/); + }); + + it("rejects a wrong type", () => { + expect(() => parseHandlerDocument({ ...valid, type: "agent" })).toThrow(/type/); + }); + + it("rejects a non-numeric version", () => { + expect(() => parseHandlerDocument({ ...valid, version: "1" })).toThrow(/version/); + }); + + it("rejects a non-boolean autoBackfill", () => { + expect(() => parseHandlerDocument({ ...valid, autoBackfill: "yes" })).toThrow( + /autoBackfill/, + ); + }); + + it("rejects a dependsOn that is not an array of strings", () => { + expect(() => parseHandlerDocument({ ...valid, dependsOn: [1, 2] })).toThrow( + /dependsOn/, + ); + }); + + it("rejects a non-object input", () => { + expect(() => parseHandlerDocument("nope")).toThrow(/mapping/); + }); +}); diff --git a/packages/shared/src/scripts/register-handler.ts b/packages/shared/src/scripts/register-handler.ts new file mode 100644 index 000000000..c585e9eb1 --- /dev/null +++ b/packages/shared/src/scripts/register-handler.ts @@ -0,0 +1,170 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +/** + * Generic post-process handler registration script. + * + * Replaces the per-worker, inlined `register-version.ts` scripts that wrote + * directly to MongoDB. Reads a declarative `handler.yaml` (path from + * HANDLER_YAML_PATH) and POSTs it to the scheduler's + * `POST /handlers/register` endpoint (the scheduler owns the `services` + * collection). Used identically by docker-compose init services and the K8s + * registration Jobs. + * + * Mirrors the coding-agent registration pattern (`agent.yaml` → + * `POST /api/v1/agents`). + * + * Environment: + * HANDLER_YAML_PATH Path to the handler.yaml file (required) + * SCHEDULER_URL Base URL of the scheduler HTTP server, e.g. + * http://scheduler:8080 (required) + */ +import { readFile } from "node:fs/promises"; +import { parse as parseYaml } from "yaml"; +import type { HandlerServiceDocument } from "../types/types.js"; + +const HEALTH_RETRIES = 30; +const HEALTH_DELAY_MS = 3_000; +const POST_RETRIES = 10; +const POST_DELAY_MS = 3_000; + +function isNonEmptyString(value: unknown): value is string { + return typeof value === "string" && value.trim().length > 0; +} + +/** Validate the parsed YAML conforms to the HandlerServiceDocument shape. */ +export function parseHandlerDocument(raw: unknown): HandlerServiceDocument { + if (typeof raw !== "object" || raw === null) { + throw new Error("handler.yaml must contain a mapping"); + } + const doc = raw as Record; + + if (!isNonEmptyString(doc._id)) { + throw new Error("handler.yaml: '_id' is required"); + } + if (doc.type !== "post-process-handler") { + throw new Error("handler.yaml: 'type' must be 'post-process-handler'"); + } + if (typeof doc.version !== "number" || !Number.isFinite(doc.version)) { + throw new Error("handler.yaml: 'version' must be a number"); + } + if (!isNonEmptyString(doc.queue)) { + throw new Error("handler.yaml: 'queue' is required"); + } + if (!isNonEmptyString(doc.selector)) { + throw new Error("handler.yaml: 'selector' is required"); + } + if (typeof doc.autoBackfill !== "boolean") { + throw new Error("handler.yaml: 'autoBackfill' must be a boolean"); + } + const dependsOn = doc.dependsOn ?? []; + if (!Array.isArray(dependsOn) || !dependsOn.every((d) => isNonEmptyString(d))) { + throw new Error("handler.yaml: 'dependsOn' must be an array of strings"); + } + + return { + _id: doc._id, + type: "post-process-handler", + version: doc.version, + queue: doc.queue, + selector: doc.selector, + autoBackfill: doc.autoBackfill, + dependsOn: dependsOn as string[], + }; +} + +async function delay(ms: number): Promise { + await new Promise((resolve) => setTimeout(resolve, ms)); +} + +async function waitForScheduler(schedulerUrl: string): Promise { + for (let attempt = 1; attempt <= HEALTH_RETRIES; attempt++) { + try { + const res = await fetch(`${schedulerUrl}/health`); + if (res.ok) { + return; + } + } catch { + // not ready yet + } + console.log( + `[register-handler] Scheduler not ready (attempt ${attempt}/${HEALTH_RETRIES}), retrying in ${HEALTH_DELAY_MS}ms...`, + ); + await delay(HEALTH_DELAY_MS); + } + throw new Error(`Scheduler did not become healthy at ${schedulerUrl}`); +} + +async function postRegistration( + schedulerUrl: string, + doc: HandlerServiceDocument, +): Promise { + for (let attempt = 1; attempt <= POST_RETRIES; attempt++) { + try { + const res = await fetch(`${schedulerUrl}/handlers/register`, { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify(doc), + }); + if (res.ok) { + return; + } + const body = await res.text(); + console.log( + `[register-handler] Register failed (attempt ${attempt}/${POST_RETRIES}): HTTP ${res.status} ${body}`, + ); + } catch (err) { + console.log( + `[register-handler] Register error (attempt ${attempt}/${POST_RETRIES}): ${ + err instanceof Error ? err.message : String(err) + }`, + ); + } + await delay(POST_DELAY_MS); + } + throw new Error(`Failed to register handler ${doc._id} after ${POST_RETRIES} attempts`); +} + +export async function registerHandlerFromYaml( + yamlPath: string, + schedulerUrl: string, +): Promise { + const content = await readFile(yamlPath, "utf-8"); + const doc = parseHandlerDocument(parseYaml(content)); + + console.log( + `[register-handler] Registering ${doc._id} (version=${doc.version}, queue=${doc.queue}) at ${schedulerUrl}`, + ); + + await waitForScheduler(schedulerUrl); + await postRegistration(schedulerUrl, doc); + + console.log(`[register-handler] Registered ${doc._id} successfully`); + return doc; +} + +async function main(): Promise { + const yamlPath = process.env.HANDLER_YAML_PATH; + const schedulerUrl = process.env.SCHEDULER_URL; + + if (!yamlPath) { + throw new Error("HANDLER_YAML_PATH is required"); + } + if (!schedulerUrl) { + throw new Error("SCHEDULER_URL is required"); + } + + await registerHandlerFromYaml(yamlPath, schedulerUrl.replace(/\/+$/, "")); +} + +// Only run when executed directly (not when imported by tests). +const isDirectRun = + process.argv[1] !== undefined && + import.meta.url === new URL(`file://${process.argv[1]}`).href; + +if (isDirectRun) { + main().catch((err) => { + console.error("[register-handler] Fatal error:", err); + process.exit(1); + }); +} diff --git a/packages/shared/src/types/types.ts b/packages/shared/src/types/types.ts index d30d5cb29..b01dd0b8b 100644 --- a/packages/shared/src/types/types.ts +++ b/packages/shared/src/types/types.ts @@ -370,6 +370,14 @@ export interface QueueMessagePayload { runId?: string; } +/** Status of a single post-process handler for a specific run. */ +export interface HandlerRunStatus { + status: "queued" | "processing" | "done" | "failed"; + version?: number; + updatedAt?: Date; + error?: string; +} + /** * RunState — represents one execution attempt of a request. * @@ -452,12 +460,16 @@ export interface RunState { aiCallCount?: number; rawChatUrl?: string; rawChatFormat?: string; - /** Version of the post-processor that last successfully processed this run. - * Used by the scheduler to detect runs needing (re-)processing when the - * post-processor version advances. */ + /** Per-handler post-processing status. Keys are handler IDs (e.g. "pp-atif"). */ + handlerStatus?: Record; + /** Set once the scheduler has triggered report generation for this run after + * the handler DAG drained (every handler reached a terminal state). Acts as + * an exactly-once guard so the scheduler POSTs /api/v1/reports/trigger at + * most once per run, even under notify/poll races. */ + reportsTriggeredAt?: Date; + /** @deprecated Use handlerStatus["pp-atif"].version instead. Kept for migration compatibility. */ postProcessorVersion?: number; - /** Lifecycle status of post-processing for this run. Prevents duplicate - * enqueuing and tracks processing progress. */ + /** @deprecated Use handlerStatus["pp-atif"].status instead. Kept for migration compatibility. */ postProcessorStatus?: "queued" | "processing" | "done" | "failed"; } @@ -470,6 +482,17 @@ export interface RunHistoryDocument extends RunState { projectId: string; // FK → ProjectDocument._id (denormalized from request) } +/** A post-process handler registered in the `services` collection. */ +export interface HandlerServiceDocument { + _id: string; + type: "post-process-handler"; + version: number; + queue: string; + selector: string; + autoBackfill: boolean; + dependsOn: string[]; +} + /** * Field names whose values move from `RequestDocument` (top-level, legacy * shape) into `RequestDocument.run: RunState` (new shape introduced for