fix(codeapi-service): poll job results when queue events lag

service/src/queue-wait.test.ts

@@ -0,0 +1,118 @@
+import { describe, expect, test } from 'bun:test';
+import type { Job, Queue, QueueEvents } from 'bullmq';
+import { EventEmitter } from 'events';
+import { waitForJobFinished } from './queue-wait';
+
+type TestResult = { ok: true; source: string };
+
+function pending<T>(): Promise<T> {
+  return new Promise(() => undefined);
+}
+
+function wait(ms: number): Promise<void> {
+  return new Promise(resolve => setTimeout(resolve, ms));
+}
+
+function fakeJob(overrides: Partial<Job<unknown, TestResult, string>>): Job<unknown, TestResult, string> {
+  return {
+    id: 'job-1',
+    waitUntilFinished: () => pending<TestResult>(),
+    ...overrides,
+  } as Job<unknown, TestResult, string>;
+}
+
+function fakeQueue(getJob: Queue<unknown, TestResult, string>['getJob']): Queue<unknown, TestResult, string> {
+  return { getJob } as Queue<unknown, TestResult, string>;
+}
+
+function fakeQueueEvents(): QueueEvents & EventEmitter {
+  return new EventEmitter() as QueueEvents & EventEmitter;
+}
+
+describe('waitForJobFinished', () => {
+  test('uses the QueueEvents result when it arrives first', async () => {
+    const result = { ok: true, source: 'events' } satisfies TestResult;
+    const job = fakeJob({});
+    const queue = fakeQueue(async () => ({
+      returnvalue: result,
+    } as Job<unknown, TestResult, string>));
+    const queueEvents = fakeQueueEvents();
+    const waitForResult = waitForJobFinished(job, queue, queueEvents, 1000);
+    queueEvents.emit('completed', { jobId: 'job-1', returnvalue: result });
+
+    await expect(waitForResult).resolves.toEqual(result);
+  });
+
+  test('stops polling after the QueueEvents result wins', async () => {
+    const result = { ok: true, source: 'events' } satisfies TestResult;
+    let getJobCalls = 0;
+    const job = fakeJob({});
+    const queue = fakeQueue(async () => {
+      getJobCalls += 1;
+      if (getJobCalls === 1) {
+        return {
+          getState: async () => 'active',
+        } as Job<unknown, TestResult, string>;
+      }
+      return {
+        returnvalue: result,
+      } as Job<unknown, TestResult, string>;
+    });
+    const queueEvents = fakeQueueEvents();
+    const waitForResult = waitForJobFinished(job, queue, queueEvents, 1000);
+    setTimeout(() => {
+      queueEvents.emit('completed', { jobId: 'job-1', returnvalue: result });
+    }, 5);
+
+    await expect(waitForResult).resolves.toEqual(result);
+    await wait(300);
+    expect(getJobCalls).toBe(2);
+  });
+
+  test('polls completed jobs when QueueEvents do not arrive', async () => {
+    const result = { ok: true, source: 'poll' } satisfies TestResult;
+    const job = fakeJob({});
+    const queue = fakeQueue(async () => ({
+      getState: async () => 'completed',
+      returnvalue: result,
+    } as Job<unknown, TestResult, string>));
+    const queueEvents = fakeQueueEvents();
+
+    await expect(waitForJobFinished(job, queue, queueEvents, 1000)).resolves.toEqual(result);
+  });
+
+  test('rejects failed jobs found by polling', async () => {
+    const job = fakeJob({});
+    const queue = fakeQueue(async () => ({
+      getState: async () => 'failed',
+      failedReason: 'sandbox failed',
+    } as Job<unknown, TestResult, string>));
+    const queueEvents = fakeQueueEvents();
+
+    await expect(waitForJobFinished(job, queue, queueEvents, 1000)).rejects.toThrow('sandbox failed');
+  });
+
+  test('rejects missing jobs instead of polling until timeout', async () => {
+    const job = fakeJob({});
+    const queue = fakeQueue(async () => undefined);
+    const queueEvents = fakeQueueEvents();
+
+    await expect(waitForJobFinished(job, queue, queueEvents, 1000)).rejects.toThrow(
+      'Job job-1 no longer exists before reaching a terminal state',
+    );
+  });
+
+  test('removes QueueEvents listeners after polling wins', async () => {
+    const result = { ok: true, source: 'poll' } satisfies TestResult;
+    const job = fakeJob({});
+    const queue = fakeQueue(async () => ({
+      getState: async () => 'completed',
+      returnvalue: result,
+    } as Job<unknown, TestResult, string>));
+    const queueEvents = fakeQueueEvents();
+
+    await expect(waitForJobFinished(job, queue, queueEvents, 1000)).resolves.toEqual(result);
+    expect(queueEvents.listenerCount('completed')).toBe(0);
+    expect(queueEvents.listenerCount('failed')).toBe(0);
+  });
+});

service/src/queue-wait.ts

@@ -0,0 +1,187 @@
+import type { Job, Queue, QueueEvents } from 'bullmq';
+
+const JOB_RESULT_POLL_INTERVAL_MS = 250;
+
+class TerminalJobError extends Error {}
+
+function wait(ms: number, signal?: AbortSignal): Promise<void> {
+  if (signal?.aborted || ms <= 0) {
+    return Promise.resolve();
+  }
+
+  return new Promise(resolve => {
+    let timeout: ReturnType<typeof setTimeout> | undefined;
+    const done = () => {
+      if (timeout !== undefined) {
+        clearTimeout(timeout);
+      }
+      signal?.removeEventListener('abort', done);
+      resolve();
+    };
+    timeout = setTimeout(done, ms);
+    signal?.addEventListener('abort', done, { once: true });
+  });
+}
+
+function throwIfAborted(jobId: string, signal?: AbortSignal): void {
+  if (signal?.aborted) {
+    throw new Error(`Job wait aborted before finishing (id=${jobId})`);
+  }
+}
+
+function waitForJobEvent<TData, TReturn, TName extends string>(
+  jobId: string,
+  queue: Queue<TData, TReturn, TName>,
+  queueEvents: QueueEvents,
+  timeoutMs: number,
+  signal?: AbortSignal,
+): Promise<TReturn> {
+  return new Promise((resolve, reject) => {
+    let timeout: ReturnType<typeof setTimeout> | undefined;
+    let settled = false;
+
+    const cleanup = () => {
+      if (timeout !== undefined) {
+        clearTimeout(timeout);
+      }
+      queueEvents.off('completed', onCompleted);
+      queueEvents.off('failed', onFailed);
+      signal?.removeEventListener('abort', onAbort);
+    };
+
+    const settle = (callback: () => void) => {
+      if (settled) return;
+      settled = true;
+      cleanup();
+      callback();
+    };
+
+    const settleWith = (promise: Promise<TReturn>) => {
+      if (settled) return;
+      settled = true;
+      cleanup();
+      void promise.then(resolve, reject);
+    };
+
+    const onCompleted = (event: { jobId: string }) => {
+      if (event.jobId !== jobId) {
+        return;
+      }
+      settleWith((async () => {
+        const currentJob = await queue.getJob(jobId);
+        if (!currentJob) {
+          throw new Error(`Job ${jobId} no longer exists after completion event`);
+        }
+        return currentJob.returnvalue;
+      })());
+    };
+    const onFailed = (event: { jobId: string; failedReason?: string }) => {
+      if (event.jobId === jobId) {
+        settle(() => reject(new Error(event.failedReason || `Job ${jobId} failed`)));
+      }
+    };
+    const onAbort = () => {
+      settle(() => reject(new Error(`Job event wait aborted before finishing (id=${jobId})`)));
+    };
+
+    if (signal?.aborted) {
+      onAbort();
+      return;
+    }
+
+    timeout = setTimeout(() => {
+      settle(() => reject(new Error(`Job wait execute timed out before receiving a finish event after ${timeoutMs}ms (id=${jobId})`)));
+    }, timeoutMs);
+    queueEvents.on('completed', onCompleted);
+    queueEvents.on('failed', onFailed);
+    signal?.addEventListener('abort', onAbort, { once: true });
+  });
+}
+
+async function pollJobUntilFinished<TData, TReturn, TName extends string>(
+  job: Job<TData, TReturn, TName>,
+  queue: Queue<TData, TReturn, TName>,
+  timeoutMs: number,
+  signal?: AbortSignal,
+): Promise<TReturn> {
+  const jobId = job.id;
+  if (!jobId) {
+    throw new Error('Cannot poll BullMQ job without an id');
+  }
+
+  const deadline = Date.now() + timeoutMs;
+  let lastState = 'unknown';
+  let lastError: unknown;
+
+  while (Date.now() < deadline) {
+    throwIfAborted(jobId, signal);
+
+    try {
+      const currentJob = await queue.getJob(jobId);
+      if (!currentJob) {
+        lastState = 'missing';
+        throw new TerminalJobError(`Job ${jobId} no longer exists before reaching a terminal state`);
+      }
+      lastState = await currentJob.getState();
+      if (lastState === 'completed') {
+        return currentJob.returnvalue;
+      }
+      if (lastState === 'failed') {
+        throw new TerminalJobError(currentJob.failedReason || `Job ${jobId} failed`);
+      }
+    } catch (error) {
+      if (error instanceof TerminalJobError) {
+        throw error;
+      }
+      lastError = error;
+      await wait(Math.min(JOB_RESULT_POLL_INTERVAL_MS, Math.max(0, deadline - Date.now())), signal);
+      continue;
+    }
+
+    await wait(Math.min(JOB_RESULT_POLL_INTERVAL_MS, Math.max(0, deadline - Date.now())), signal);
+  }
+
+  throwIfAborted(jobId, signal);
+
+  const suffix = lastError instanceof Error ? `; last poll error: ${lastError.message}` : '';
+  throw new Error(
+    `Job wait execute timed out before finishing, no terminal state observed after ${timeoutMs}ms (id=${jobId}, lastState=${lastState})${suffix}`,
+  );
+}
+
+async function waitForJobFinished<TData, TReturn, TName extends string>(
+  job: Job<TData, TReturn, TName>,
+  queue: Queue<TData, TReturn, TName>,
+  queueEvents: QueueEvents,
+  timeoutMs: number,
+): Promise<TReturn> {
+  const pollAbortController = new AbortController();
+  const eventAbortController = new AbortController();
+  const jobId = job.id;
+  if (!jobId) {
+    throw new Error('Cannot wait for BullMQ job without an id');
+  }
+
+  const eventWait = waitForJobEvent<TData, TReturn, TName>(
+    jobId,
+    queue,
+    queueEvents,
+    timeoutMs,
+    eventAbortController.signal,
+  );
+  void eventWait.catch(() => undefined);
+  const pollWait = pollJobUntilFinished(job, queue, timeoutMs, pollAbortController.signal);
+  void pollWait.catch(() => undefined);
+
+  try {
+    return await Promise.race([
+      eventWait,
+      pollWait,
+    ]);
+  } finally {
+    pollAbortController.abort();
+    eventAbortController.abort();
+  }
+}
+
+export { waitForJobFinished };

service/src/queue.ts

@@ -10,6 +10,7 @@ import { env } from './config';
 import logger from './logger';
 import { redisKeepAliveOptions } from './redis-options';
 import { bullmqQueueJobs, registerBullmqQueueMetricsCollector } from './metrics';
+import { waitForJobFinished } from './queue-wait';
 
 const MAX_RECONNECT_ATTEMPTS = 5;
 const RECONNECT_DELAY = 2000;
@@ -109,4 +110,4 @@ registerBullmqQueueMetricsCollector(async () => {
  * BullMQ coordination objects. */
 setMaxListeners(0, pyQueue, otherQueue, pyQueueEvents, otherQueueEvents);
 
-export { pyQueue, otherQueue, pyQueueEvents, otherQueueEvents, connection };
+export { pyQueue, otherQueue, pyQueueEvents, otherQueueEvents, connection, waitForJobFinished };

service/src/service/programmatic-router.ts

@@ -6,7 +6,7 @@ import type { Queue, QueueEvents } from 'bullmq';
 import type * as t from '../types';
 import { checkServiceStartUp, checkServiceShutDown } from '../lifecycle';
 import { executionLimiter } from '../middleware/limits';
-import { pyQueue, pyQueueEvents, otherQueue, otherQueueEvents, connection } from '../queue';
+import { pyQueue, pyQueueEvents, otherQueue, otherQueueEvents, connection, waitForJobFinished } from '../queue';
 import { createProgrammaticPayload, extractPendingFromStdout } from '../preamble';
 import { findBashToolNameCollision } from '../preamble-bash';
 import type { LCTool } from '../preamble';
@@ -405,13 +405,13 @@ async function runReplayIteration(
     egressGrantClaims: sandboxSecurity.egressGrantClaims,
     egressGrantToken: sandboxSecurity.egressGrantToken,
   }, {
-    removeOnComplete: { age: 60, count: 1 },
+    removeOnComplete: { age: 60, count: 100 },
     removeOnFail: { age: 180, count: 1 },
     attempts: 1,
   });
   jobsSubmitted.inc({ language });
 
-  return job.waitUntilFinished(events, env.JOB_TIMEOUT);
+  return waitForJobFinished(job, queue, events, env.JOB_TIMEOUT);
 }
 
 function isSandboxRunSuccess(result: t.ExecuteResult): boolean {
@@ -1372,7 +1372,7 @@ async function handleBlocking(
       egressGrantClaims: sandboxSecurity.egressGrantClaims,
       egressGrantToken: sandboxSecurity.egressGrantToken,
     }, {
-      removeOnComplete: { age: 60, count: 1 },
+      removeOnComplete: { age: 60, count: 100 },
       removeOnFail: { age: 180, count: 1 },
       attempts: 1,
       jobId: session_id,
@@ -1394,7 +1394,7 @@ async function handleBlocking(
       }
     });
 
-    job.waitUntilFinished(pyQueueEvents, env.JOB_TIMEOUT)
+    waitForJobFinished(job, pyQueue, pyQueueEvents, env.JOB_TIMEOUT)
       .then(async (result) => {
         if (clientDisconnected) return;
         await setExecutionResult(execution_id, result);

service/src/service/router.ts

@@ -10,7 +10,7 @@ import { sessionAuth } from '../middleware/auth';
 import { executionLimiter, uploadLimiter, downloadLimiter, fetchLimiter } from '../middleware/limits';
 import { internalServiceHeaders } from '../internal-service-auth';
 import { resolveSessionKey, resolveOutputBucketSessionKey, SessionKeyResolutionError, parseUploadSessionKeyInput, type SessionKeyInput } from '../session-key';
-import { pyQueue, otherQueue, pyQueueEvents, otherQueueEvents, connection } from '../queue';
+import { pyQueue, otherQueue, pyQueueEvents, otherQueueEvents, connection, waitForJobFinished } from '../queue';
 import { sleep, getAxiosErrorDetails, publicExecutionFailure } from '../utils';
 import { env, planLimits, resolveLanguage } from '../config';
 import { createPayload } from '../payload';
@@ -226,7 +226,7 @@ router.post('/exec', executionLimiter, async (req: t.AuthenticatedRequest, res)
       }, {
         removeOnComplete: {
           age: 60,
-          count: 1,
+          count: 100,
         },
         removeOnFail: {
           age: 180,
@@ -251,7 +251,7 @@ router.post('/exec', executionLimiter, async (req: t.AuthenticatedRequest, res)
       'messaging.system': 'bullmq',
       'messaging.destination.name': queueName,
       'codeapi.language': language,
-    }, () => job.waitUntilFinished(queueEvents, env.JOB_TIMEOUT), 'CONSUMER');
+    }, () => waitForJobFinished(job, queue, queueEvents, env.JOB_TIMEOUT), 'CONSUMER');
 
     if (!isSyntheticRequest) {
       logger.info('Execution completed', { session_id, user_id });