From 400673d25b99a0a01e414d74b2dd15eabf787596 Mon Sep 17 00:00:00 2001 From: Claude Date: Sun, 1 Feb 2026 11:56:23 +0000 Subject: [PATCH] feat(run-engine): add RunDataProvider for V3 format ack/nack operations Instead of creating message keys when dequeuing V3 messages, we now read run data from PostgreSQL via a RunDataProvider when needed for ack/nack operations. This approach: - Eliminates ALL message keys for V3 format (not just pending runs) - Uses PostgreSQL as the source of truth for run data - Falls back to Redis message key for legacy V2 messages - Adds RunDataProvider interface and RunData type The RunEngine creates a runDataProvider that queries the TaskRun table for queue, orgId, environmentId, etc. when readMessage is called and no Redis message key exists. https://claude.ai/code/session_01AyzQp6tbj7th5QRTCYjJR5 --- .../run-engine/src/engine/index.ts | 37 +++++++++++++ .../run-engine/src/run-queue/index.ts | 54 ++++++++++++++++--- .../run-engine/src/run-queue/types.ts | 28 ++++++++++ 3 files changed, 111 insertions(+), 8 deletions(-) diff --git a/internal-packages/run-engine/src/engine/index.ts b/internal-packages/run-engine/src/engine/index.ts index 9e81c9913..8a4e04a18 100644 --- a/internal-packages/run-engine/src/engine/index.ts +++ b/internal-packages/run-engine/src/engine/index.ts @@ -182,6 +182,43 @@ export class RunEngine { processWorkerQueueDebounceMs: options.queue?.processWorkerQueueDebounceMs, dequeueBlockingTimeoutSeconds: options.queue?.dequeueBlockingTimeoutSeconds, meter: options.meter, + // Run data provider for V3 optimized format - reads from PostgreSQL when no Redis message key exists + runDataProvider: { + getRunData: async (runId: string) => { + const run = await this.prisma.taskRun.findUnique({ + where: { id: runId }, + select: { + queue: true, + organizationId: true, + projectId: true, + runtimeEnvironmentId: true, + environmentType: true, + concurrencyKey: true, + attemptNumber: true, + queueTimestamp: true, + workerQueue: true, + taskIdentifier: true, + }, + }); + + if (!run || !run.organizationId || !run.environmentType) { + return undefined; + } + + return { + queue: run.queue, + orgId: run.organizationId, + projectId: run.projectId, + environmentId: run.runtimeEnvironmentId, + environmentType: run.environmentType, + concurrencyKey: run.concurrencyKey ?? undefined, + attempt: run.attemptNumber ?? 0, + timestamp: run.queueTimestamp?.getTime() ?? Date.now(), + workerQueue: run.workerQueue, + taskIdentifier: run.taskIdentifier, + }; + }, + }, }); this.worker = new Worker({ diff --git a/internal-packages/run-engine/src/run-queue/index.ts b/internal-packages/run-engine/src/run-queue/index.ts index 76421ba6b..7dd13d1b7 100644 --- a/internal-packages/run-engine/src/run-queue/index.ts +++ b/internal-packages/run-engine/src/run-queue/index.ts @@ -39,6 +39,7 @@ import { InputPayload, OutputPayload, OutputPayloadV2, + RunDataProvider, RunQueueKeyProducer, RunQueueKeyProducerEnvironment, RunQueueSelectionStrategy, @@ -110,6 +111,12 @@ export type RunQueueOptions = { * 3. Old messages drain naturally as they're processed */ useOptimizedMessageFormat?: boolean; + /** + * Provider for fetching run data from PostgreSQL. + * Required when using V3 optimized format for ack/nack operations. + * Falls back to Redis message key if not provided (legacy behavior). + */ + runDataProvider?: RunDataProvider; }; export interface ConcurrencySweeperCallback { @@ -194,9 +201,11 @@ export class RunQueue { private _meter: Meter; private _queueCooloffStates: Map = new Map(); private _useOptimizedMessageFormat: boolean; + private _runDataProvider?: RunDataProvider; constructor(public readonly options: RunQueueOptions) { this._useOptimizedMessageFormat = options.useOptimizedMessageFormat ?? false; + this._runDataProvider = options.runDataProvider; this.shardCount = options.shardCount ?? 2; this.retryOptions = options.retryOptions ?? defaultRetrySettings; this.redis = createRedisClient(options.redis, { @@ -575,8 +584,43 @@ export class RunQueue { return this.redis.exists(this.keys.messageKey(orgId, messageId)); } - public async readMessage(orgId: string, messageId: string) { - return this.readMessageFromKey(this.keys.messageKey(orgId, messageId)); + public async readMessage(orgId: string, messageId: string): Promise { + // First try to read from Redis (legacy V2 format) + const redisMessage = await this.readMessageFromKey(this.keys.messageKey(orgId, messageId)); + if (redisMessage) { + return redisMessage; + } + + // Fall back to runDataProvider (for V3 format where there's no message key) + if (this._runDataProvider) { + const runData = await this._runDataProvider.getRunData(messageId); + if (runData) { + // Convert RunData to OutputPayloadV2 + const queueKey = this.keys.queueKey( + runData.orgId, + runData.projectId, + runData.environmentId, + runData.taskIdentifier, + runData.concurrencyKey + ); + return { + version: "2" as const, + runId: messageId, + taskIdentifier: runData.taskIdentifier, + orgId: runData.orgId, + projectId: runData.projectId, + environmentId: runData.environmentId, + environmentType: runData.environmentType, + queue: queueKey, + concurrencyKey: runData.concurrencyKey, + timestamp: runData.timestamp, + attempt: runData.attempt, + workerQueue: runData.workerQueue, + }; + } + } + + return undefined; } public async readMessageFromKey(messageKey: string) { @@ -2373,12 +2417,6 @@ export class RunQueue { const descriptor = this.keys.descriptorFromQueue(decoded.queueKey); const message = reconstructMessageFromWorkerEntry(decoded, descriptor); - // For V3 format: create the message key now that the run is executing. - // This allows ack/nack to work (they read from message key). - // Storage savings come from not having message keys for PENDING runs (the backlog). - const messageKey = this.keys.messageKey(descriptor.orgId, message.runId); - await this.redis.set(messageKey, JSON.stringify(message)); - // Update the currentDequeued sets (this is done in the Lua script for legacy) const queueCurrentDequeuedKey = this.keys.queueCurrentDequeuedKeyFromQueue(decoded.queueKey); const envCurrentDequeuedKey = this.keys.envCurrentDequeuedKeyFromQueue(decoded.queueKey); diff --git a/internal-packages/run-engine/src/run-queue/types.ts b/internal-packages/run-engine/src/run-queue/types.ts index ee1ce41b7..903ae138e 100644 --- a/internal-packages/run-engine/src/run-queue/types.ts +++ b/internal-packages/run-engine/src/run-queue/types.ts @@ -133,3 +133,31 @@ export interface RunQueueSelectionStrategy { consumerId: string ): Promise>; } + +/** + * Provider for fetching run data from a persistent store (e.g., PostgreSQL). + * Used for V3 optimized format where message data is not stored in Redis. + */ +export interface RunDataProvider { + /** + * Fetch run data for ack/nack operations. + * Returns undefined if the run is not found. + */ + getRunData(runId: string): Promise; +} + +/** + * Run data needed for queue operations (ack, nack, release concurrency). + */ +export type RunData = { + queue: string; + orgId: string; + projectId: string; + environmentId: string; + environmentType: RuntimeEnvironmentType; + concurrencyKey?: string; + attempt: number; + timestamp: number; + workerQueue: string; + taskIdentifier: string; +};