fix(mothership): retry a job dispatch whose usage read failed instead of skipping the run

apps/sim/lib/mothership/jobs/dispatch.ts

@@ -106,7 +106,13 @@ async function dispatchQueuedOccurrence(executionId: string): Promise<JobDispatc
     if (isHosted) {
       const usage = await checkAttributedUsageLimits(billingAttribution)
       if (usage.isExceeded) {
-        await endUnstartedJobRun(executionId, 'skipped_budget', 'The usage limit was reached.')
+        // An outage, not a verdict.
+        if (usage.reason === 'usage_unavailable') return retryDispatch(run, usage.message)
+        await endUnstartedJobRun(
+          executionId,
+          'skipped_budget',
+          usage.message ?? 'The usage limit was reached.'
+        )
         return 'skipped_budget'
       }
     }
@@ -138,14 +144,19 @@ async function dispatchQueuedOccurrence(executionId: string): Promise<JobDispatc
     if (code === 'forbidden' || code === 'not_found') {
       return block(run, 'The job owner can no longer run it in this workspace.')
     }
-    const failures = await countJobDispatchFailure(executionId)
-    logger.warn('Job occurrence could not start; the reconciler retries it', {
-      executionId,
-      failures,
-      error: getErrorMessage(error),
-    })
-    if (failures < MAX_JOB_DISPATCH_ATTEMPTS) return 'retry'
-    await failUnstartedJobRun(executionId, run.scheduleId, 'The run could not be started.')
-    return 'failed'
+    return retryDispatch(run, getErrorMessage(error))
   }
 }
+
+/** Counts a failed attempt; the reconciler retries it, and the last allowed one ends the run as failed. */
+async function retryDispatch(run: JobRun, reason?: string): Promise<JobDispatchOutcome> {
+  const failures = await countJobDispatchFailure(run.executionId)
+  logger.warn('Job occurrence could not start; the reconciler retries it', {
+    executionId: run.executionId,
+    failures,
+    error: reason,
+  })
+  if (failures < MAX_JOB_DISPATCH_ATTEMPTS) return 'retry'
+  await failUnstartedJobRun(run.executionId, run.scheduleId, 'The run could not be started.')
+  return 'failed'
+}

apps/sim/lib/mothership/jobs/lane.integration.ts

@@ -6,7 +6,7 @@ import { deleteWorkspaceFixture, insertWorkspaceFixture } from '@sim/db/testing/
  * and finalizing each occurrence once. A local HTTP server stands in for the worker; each POST
  * for a turn is one leg, which ends the turn, fails it, or stays open until the test acts.
  */
-import { envFlagsMock, resetEnvFlagsMock } from '@sim/testing/mocks/env-flags.mock'
+import { envFlagsMock, resetEnvFlagsMock, setEnvFlags } from '@sim/testing/mocks/env-flags.mock'
 import { featureFlagsMock, featureFlagsMockFns } from '@sim/testing/mocks/feature-flags.mock'
 import { afterAll, beforeAll, beforeEach, describe, expect, it, vi } from 'vitest'
 
@@ -108,6 +108,11 @@ vi.mock('@/lib/mothership/jobs/reconcile', async (importOriginal) => {
   return { ...actual, reconcileJobRuns: vi.fn(actual.reconcileJobRuns) }
 })
 vi.mock('@/lib/core/config/feature-flags', () => featureFlagsMock)
+/** Spied so a test can answer the hosted usage gate without a billing ledger. */
+vi.mock('@/lib/billing/core/attribution', async (importOriginal) => {
+  const actual = await importOriginal<typeof import('@/lib/billing/core/attribution')>()
+  return { ...actual, checkAttributedUsageLimits: vi.fn(actual.checkAttributedUsageLimits) }
+})
 /** Spied so a test can tell an automatic turn-off, which emails the job's people, from none. */
 vi.mock('@/lib/workflows/schedules/disable-notifications', () => ({
   notifyScheduleAutoDisabled: vi.fn(async () => {}),
@@ -139,7 +144,10 @@ import {
 import { sleep } from '@sim/utils/helpers'
 import { generateId } from '@sim/utils/id'
 import { eq, inArray, sql } from 'drizzle-orm'
-import { resolveBillingAttribution } from '@/lib/billing/core/attribution'
+import {
+  checkAttributedUsageLimits,
+  resolveBillingAttribution,
+} from '@/lib/billing/core/attribution'
 import {
   acquireLock,
   closeRedisConnection,
@@ -256,11 +264,21 @@ async function runsOf(scheduleId: string) {
       status: jobExecutionLogs.status,
       outcome: jobExecutionLogs.outcome,
       summary: jobExecutionLogs.summary,
+      dispatchFailures: sql<
+        number | null
+      >`(${jobExecutionLogs.executionData}->>'dispatchFailures')::int`,
     })
     .from(jobExecutionLogs)
     .where(eq(jobExecutionLogs.scheduleId, scheduleId))
 }
 
+/** The usage gate's answer when the payer's usage can't be read. */
+const usageOutage = {
+  isExceeded: true,
+  message: 'Usage could not be checked.',
+  reason: 'usage_unavailable' as const,
+}
+
 /** Makes a schedule due again now, as its next occurrence arriving would. */
 async function dueAgain(scheduleId: string) {
   const dueAt = new Date(Date.now() - 1_000)
@@ -290,6 +308,8 @@ async function reconcileUntilSettled(executionId: string) {
 }
 
 afterAll(async () => {
+  /** A test that failed before ending its stream must not hold the server open. */
+  worker.server.closeAllConnections()
   await new Promise<void>((resolve) => worker.server.close(() => resolve()))
   for (const [key, value] of Object.entries(inheritedEnv)) {
     if (value === undefined) delete process.env[key]
@@ -634,6 +654,55 @@ describe.runIf(Boolean(redisUrl))('the scheduled job lane', () => {
     await reconcileUntilSettled(occurrenceId)
   })
 
+  it('retries an occurrence whose usage could not be read, and runs it once usage reads again', async () => {
+    setEnvFlags({ isHosted: true })
+    const { scheduleId, occurrenceId } = await seedJob()
+    vi.mocked(checkAttributedUsageLimits).mockResolvedValueOnce(usageOutage)
+    await runJobLane(new Date())
+    expect(await runsOf(scheduleId)).toEqual([
+      expect.objectContaining({ status: 'queued', outcome: null, dispatchFailures: 1 }),
+    ])
+
+    vi.mocked(checkAttributedUsageLimits).mockResolvedValueOnce({ isExceeded: false })
+    await reconcileJobRuns(new Date(Date.now() + JOB_DISPATCH_GRACE_MS + 1_000))
+    await reconcileUntilSettled(occurrenceId)
+    expect(legsFor(occurrenceId)).toHaveLength(1)
+    expect(await runsOf(scheduleId)).toEqual([
+      expect.objectContaining({ status: 'complete', outcome: 'not_met', dispatchFailures: 1 }),
+    ])
+  })
+
+  it('ends an occurrence whose usage stays unreadable at its last attempt, and skips one whose billing is blocked', async () => {
+    setEnvFlags({ isHosted: true })
+    const outage = await seedJob()
+    await claimDueJobOccurrences(new Date(), 10)
+    for (let attempt = 1; attempt <= 5; attempt++) {
+      vi.mocked(checkAttributedUsageLimits).mockResolvedValueOnce(usageOutage)
+      expect(await dispatchJobOccurrence(outage.occurrenceId)).toBe(
+        attempt < 5 ? 'retry' : 'failed'
+      )
+    }
+    expect(await runsOf(outage.scheduleId)).toEqual([
+      expect.objectContaining({ outcome: 'failed', summary: 'The run could not be started.' }),
+    ])
+    expect(legsFor(outage.occurrenceId)).toHaveLength(0)
+
+    const blocked = await seedJob()
+    await claimDueJobOccurrences(new Date(), 10)
+    vi.mocked(checkAttributedUsageLimits).mockResolvedValueOnce({
+      isExceeded: true,
+      message: 'Billing is blocked for this account.',
+      reason: 'billing_blocked',
+    })
+    expect(await dispatchJobOccurrence(blocked.occurrenceId)).toBe('skipped_budget')
+    expect(await runsOf(blocked.scheduleId)).toEqual([
+      expect.objectContaining({
+        outcome: 'skipped_budget',
+        summary: 'Billing is blocked for this account.',
+      }),
+    ])
+  })
+
   it('runs a one-shot job once and completes it', async () => {
     const { scheduleId, occurrenceId } = await seedJob({ cronExpression: null })
     await runJobLane(new Date())