Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -3478,8 +3478,12 @@ function JobRunBudgetDisplay({ message }: JobRunBudgetDisplayProps) {
)
}

interface PlanUsageLimitDisplayProps {
data: UsageUpgradeTagData
}

/** The payer's plan limit: the card links to where the plan or the limit can change. */
function PlanUsageLimitDisplay({ data }: { data: UsageUpgradeTagData }) {
function PlanUsageLimitDisplay({ data }: PlanUsageLimitDisplayProps) {
const { data: session } = useSession()
const hostContext = useOptionalWorkspaceHostContext()
const organizationContext = useOptionalOrganizationContext()
Expand Down
10 changes: 5 additions & 5 deletions apps/sim/lib/mothership/jobs/brief.ts
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,7 @@ const BRIEF_SUMMARY_CHARS = 400

type JobBriefSchedule = Pick<
JobRun,
'scheduleId' | 'jobTitle' | 'prompt' | 'successCondition' | 'lifecycle' | 'maxRuns' | 'runCount'
'scheduleId' | 'jobTitle' | 'prompt' | 'successCondition' | 'lifecycle' | 'maxRuns' | 'runNumber'
>

/**
Expand All @@ -20,7 +20,7 @@ type JobBriefSchedule = Pick<
export async function buildJobBrief(
schedule: JobBriefSchedule,
occurrenceAt: Date,
runBudgetUsd: number
runBudgetUsd: number | null
): Promise<string> {
const recent = await db
.select({
Expand All @@ -40,8 +40,8 @@ export async function buildJobBrief(
.limit(BRIEF_RECENT_RUNS)

const run = schedule.maxRuns
? `Run ${schedule.runCount} of ${schedule.maxRuns}`
: `Run ${schedule.runCount}`
? `Run ${schedule.runNumber} of ${schedule.maxRuns}`
: `Run ${schedule.runNumber}`
return [
`Job: ${schedule.jobTitle ?? 'Untitled job'} (${run}, scheduled for ${occurrenceAt.toISOString()})`,
`Task:\n${schedule.prompt ?? ''}`,
Expand All @@ -59,7 +59,7 @@ export async function buildJobBrief(
)
.join('\n')}`
: 'This is the first run.',
`This run may spend up to $${runBudgetUsd.toFixed(2)}.`,
runBudgetUsd === null ? '' : `This run may spend up to $${runBudgetUsd.toFixed(2)}.`,
]
.filter((part) => part.length > 0)
.join('\n\n')
Expand Down
2 changes: 2 additions & 0 deletions apps/sim/lib/mothership/jobs/claim.ts
Original file line number Diff line number Diff line change
Expand Up @@ -84,6 +84,8 @@ export async function claimDueJobOccurrences(
status: 'queued',
trigger: JOB_RUN_TRIGGER,
startedAt: now,
/** Fixed at claim: the schedule's count moves on as later occurrences are claimed. */
executionData: { runNumber: runCount },
})
.onConflictDoNothing()
await tx
Expand Down
31 changes: 29 additions & 2 deletions apps/sim/lib/mothership/jobs/dispatch.ts
Original file line number Diff line number Diff line change
@@ -1,10 +1,12 @@
import { createLogger } from '@sim/logger'
import { getErrorMessage } from '@sim/utils/errors'
import { getErrorMessage, toError } from '@sim/utils/errors'
import { generateShortId } from '@sim/utils/id'
import {
checkAttributedUsageLimits,
resolveBillingAttribution,
} from '@/lib/billing/core/attribution'
import { isHosted } from '@/lib/core/config/env-flags'
import { acquireLock, releaseLock } from '@/lib/core/config/redis'
import { asOrchestrationError } from '@/lib/core/orchestration/types'
import { buildJobBrief } from '@/lib/mothership/jobs/brief'
import { jobRunBudgetUsd } from '@/lib/mothership/jobs/budget'
Expand All @@ -26,6 +28,8 @@ const logger = createLogger('MothershipJobDispatch')

/** A queued occurrence whose start failed this many times ends as failed instead of retrying. */
const MAX_JOB_DISPATCH_ATTEMPTS = 5
/** Outlasts one dispatch (its reads and the turn's admission), so an attempt never overlaps another. */
const JOB_DISPATCH_LOCK_SECONDS = 120

type JobDispatchOutcome =
| 'started'
Expand Down Expand Up @@ -67,6 +71,27 @@ async function block(run: JobRun, summary: string): Promise<JobDispatchOutcome>
* skipped rather than queued behind it.
*/
export async function dispatchJobOccurrence(executionId: string): Promise<JobDispatchOutcome> {
/**
* One attempt at a time per occurrence (a tick's and the reconciler's can meet): a skip or
* block decided while another attempt admits the turn would close the record of a live run.
*/
const lock = `mothership:job-dispatch:${executionId}`
const holder = generateShortId()
if (!(await acquireLock(lock, holder, JOB_DISPATCH_LOCK_SECONDS, { reclaimOnFailure: true })))
return 'already_started'
try {
return await dispatchQueuedOccurrence(executionId)
} finally {
await releaseLock(lock, holder).catch((error) =>
logger.warn('Job dispatch lock release failed; it expires', {
executionId,
error: toError(error).message,
})
)
}
}

async function dispatchQueuedOccurrence(executionId: string): Promise<JobDispatchOutcome> {
const run = await loadJobRun(executionId)
if (!run || run.status !== 'queued' || !run.occurrenceAt) return 'not_queued'
const { userId, workspaceId, chatId } = run
Expand All @@ -85,7 +110,9 @@ export async function dispatchJobOccurrence(executionId: string): Promise<JobDis
return 'skipped_budget'
}
}
const message = await buildJobBrief(run, run.occurrenceAt, jobRunBudgetUsd(run.runBudgetUsd))
/** The run's spend cap is enforced through hosted billing callbacks, so only there is it stated. */
const runBudgetUsd = isHosted ? jobRunBudgetUsd(run.runBudgetUsd) : null
const message = await buildJobBrief(run, run.occurrenceAt, runBudgetUsd)
const start = await startUnattendedTurn({
userId,
workspaceId,
Expand Down
12 changes: 9 additions & 3 deletions apps/sim/lib/mothership/jobs/finalize.ts
Original file line number Diff line number Diff line change
Expand Up @@ -48,12 +48,18 @@ async function settleJobSchedule(
.set({
failedCount: sql`${workflowSchedule.failedCount} + 1`,
lastFailedAt: new Date(),
status: sql`CASE WHEN ${workflowSchedule.failedCount} + 1 >= ${MAX_CONSECUTIVE_JOB_FAILURES} AND ${workflowSchedule.status} = 'active' THEN 'disabled' ELSE ${workflowSchedule.status} END`,
updatedAt: new Date(),
})
.where(eq(workflowSchedule.id, scheduleId))
.returning({ status: workflowSchedule.status, failedCount: workflowSchedule.failedCount })
return schedule?.status === 'disabled' && schedule.failedCount === MAX_CONSECUTIVE_JOB_FAILURES
.returning({ failedCount: workflowSchedule.failedCount })
if (!schedule || schedule.failedCount < MAX_CONSECUTIVE_JOB_FAILURES) return false
/** Only a job still on is turned off here; one a person already turned off is not reported. */
const disabled = await tx
.update(workflowSchedule)
.set({ status: 'disabled', updatedAt: new Date() })
.where(and(eq(workflowSchedule.id, scheduleId), eq(workflowSchedule.status, 'active')))
.returning({ id: workflowSchedule.id })
return disabled.length > 0
}

/** The turn's last assistant text, as the summary of a run that reported no outcome. */
Expand Down
99 changes: 98 additions & 1 deletion apps/sim/lib/mothership/jobs/lane.integration.ts
Original file line number Diff line number Diff line change
Expand Up @@ -108,6 +108,10 @@ 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 tell an automatic turn-off, which emails the job's people, from none. */
vi.mock('@/lib/workflows/schedules/disable-notifications', () => ({
notifyScheduleAutoDisabled: vi.fn(async () => {}),
}))

const { controllers } = vi.hoisted(() => ({ controllers: new Map<string, AbortController>() }))
vi.mock('@/lib/mothership/request/session/abort', async (importOriginal) => {
Expand Down Expand Up @@ -136,7 +140,12 @@ 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 { closeRedisConnection, getRedisClient } from '@/lib/core/config/redis'
import {
acquireLock,
closeRedisConnection,
getRedisClient,
releaseLock,
} from '@/lib/core/config/redis'
import { ABANDONED_RUN_CURSOR_KEY } from '@/lib/mothership/async-runs/orphaned-runs'
import { chatPubSub } from '@/lib/mothership/chat-status'
import { recordJobRunSpend } from '@/lib/mothership/jobs/budget'
Expand All @@ -152,6 +161,7 @@ import {
chatStreamLockKey,
StreamControllerSupersededError,
} from '@/lib/mothership/request/session/controller-lease'
import { notifyScheduleAutoDisabled } from '@/lib/workflows/schedules/disable-notifications'

interface QueryPlan {
'Node Type': string
Expand Down Expand Up @@ -542,6 +552,88 @@ describe.runIf(Boolean(redisUrl))('the scheduled job lane', () => {
worker.open.delete(occurrenceId)
}, 60_000)

it('delivers a time-limit Stop an earlier pass recorded but died delivering, before finalizing', async () => {
const { scheduleId, occurrenceId } = await seedJob()
worker.legs.set(occurrenceId, [{ text: 'Still triaging. ', mode: 'hang' }])
await runJobLane(new Date())
await waitFor(async () => worker.open.has(occurrenceId), 'the first leg')
/** What a pass leaves when it dies after recording why the run stops, before delivering it. */
await db
.update(jobExecutionLogs)
.set({ executionData: sql`jsonb_build_object('deadlineStopRequestedAt', now())` })
.where(eq(jobExecutionLogs.executionId, occurrenceId))

await reconcileUntilSettled(occurrenceId)

expect(
worker.requests.filter(
(request) =>
request.path === '/api/streams/explicit-abort' && request.body.messageId === occurrenceId
)
).toHaveLength(1)
expect(await runsOf(scheduleId)).toEqual([
expect.objectContaining({
outcome: 'interrupted',
summary: 'Stopped at the run time limit.',
}),
])
worker.open.get(occurrenceId)?.response.end()
worker.open.delete(occurrenceId)
}, 60_000)

it('does not report a job a person turned off as turned off by its failures', async () => {
vi.mocked(notifyScheduleAutoDisabled).mockClear()
const { scheduleId, occurrenceId } = await seedJob()
worker.legs.set(occurrenceId, [{ text: 'Still triaging. ', mode: 'hang' }])
await runJobLane(new Date())
await waitFor(async () => worker.open.has(occurrenceId), 'the first leg')
/** Off by hand with two failures behind it, while this run is still live. */
await db
.update(workflowSchedule)
.set({ status: 'disabled', failedCount: 2 })
.where(eq(workflowSchedule.id, scheduleId))

await reconcileJobRuns(new Date(Date.now() + 6 * 60 * 60_000 + 60_000))
await reconcileUntilSettled(occurrenceId)

expect(await schedule(scheduleId)).toMatchObject({ status: 'disabled', failedCount: 3 })
expect(notifyScheduleAutoDisabled).not.toHaveBeenCalled()
worker.open.get(occurrenceId)?.response.end()
worker.open.delete(occurrenceId)
}, 60_000)

it('briefs a retried occurrence with its own run number, not the job count a later claim moved', async () => {
const { scheduleId, occurrenceId } = await seedJob()
await claimDueJobOccurrences(new Date(), 10)
const later = await dueAgain(scheduleId)
await claimDueJobOccurrences(new Date(), 10)
expect(await schedule(scheduleId)).toMatchObject({ runCount: 2 })

for (const id of [occurrenceId, later]) {
await dispatchJobOccurrence(id)
await reconcileUntilSettled(id)
}

expect(legsFor(occurrenceId)[0]?.body.message).toContain('(Run 1, scheduled for')
expect(legsFor(later)[0]?.body.message).toContain('(Run 2, scheduled for')
})

it('leaves an occurrence another attempt is dispatching to that attempt', async () => {
const { scheduleId, occurrenceId } = await seedJob()
await claimDueJobOccurrences(new Date(), 10)
const lock = `mothership:job-dispatch:${occurrenceId}`
expect(await acquireLock(lock, 'another-attempt', 60)).toBe(true)
try {
expect(await dispatchJobOccurrence(occurrenceId)).toBe('already_started')
expect(legsFor(occurrenceId)).toHaveLength(0)
expect(await runsOf(scheduleId)).toEqual([expect.objectContaining({ status: 'queued' })])
} finally {
await releaseLock(lock, 'another-attempt')
}
expect(await dispatchJobOccurrence(occurrenceId)).toBe('started')
await reconcileUntilSettled(occurrenceId)
})

it('runs a one-shot job once and completes it', async () => {
const { scheduleId, occurrenceId } = await seedJob({ cronExpression: null })
await runJobLane(new Date())
Expand Down Expand Up @@ -682,6 +774,7 @@ describe.runIf(Boolean(redisUrl))('the scheduled job lane', () => {
})

it('turns a job off after three failed runs in a row, and a verdict clears the streak', async () => {
vi.mocked(notifyScheduleAutoDisabled).mockClear()
const { scheduleId, occurrenceId } = await seedJob()
const fail = (id: string) =>
worker.legs.set(
Expand Down Expand Up @@ -710,6 +803,10 @@ describe.runIf(Boolean(redisUrl))('the scheduled job lane', () => {
}
const outcomes = (await runsOf(scheduleId)).map((run) => run.outcome).sort()
expect(outcomes).toEqual(['failed', 'failed', 'failed', 'failed', 'not_met'])
expect(notifyScheduleAutoDisabled).toHaveBeenCalledExactlyOnceWith({
scheduleId,
reason: 'consecutive_failures',
})

await dueAgain(scheduleId)
await runJobLane(new Date())
Expand Down
55 changes: 37 additions & 18 deletions apps/sim/lib/mothership/jobs/reconcile.ts
Original file line number Diff line number Diff line change
Expand Up @@ -46,8 +46,8 @@ const admittedBillingRequestId = sql<

/**
* Open records that need this pass, oldest first. Every queued record does; a running one only
* once its turn ended or its run row is gone, its time limit passed without a Stop, or its
* spend-cap link is missing, so a backlog of healthy long runs never hides the ones that finished.
* once its turn ended or its run row is gone, its time limit passed without a Stop, its Stop's
* delivery never finished, or its spend-cap link is missing, so a backlog of healthy long runs never hides the ones that finished.
*/
function openJobRuns(status: OpenJobRunStatus, now: Date) {
const needsAction =
Expand All @@ -60,6 +60,7 @@ function openJobRuns(status: OpenJobRunStatus, now: Date) {
lte(jobExecutionLogs.startedAt, new Date(now.getTime() - JOB_RUN_TIME_LIMIT_MS)),
sql`${jobExecutionLogs.executionData}->>'deadlineStopRequestedAt' IS NULL`
),
sql`${jobExecutionLogs.executionData}->>'deadlineStopRequestedAt' IS NOT NULL AND ${jobExecutionLogs.executionData}->>'deadlineStopDeliveredAt' IS NULL`,
and(isNull(jobExecutionLogs.billingRequestId), isNotNull(admittedBillingRequestId))
)
return db
Expand All @@ -70,6 +71,9 @@ function openJobRuns(status: OpenJobRunStatus, now: Date) {
deadlineStopRequestedAt: sql<
string | null
>`${jobExecutionLogs.executionData}->>'deadlineStopRequestedAt'`,
deadlineStopDeliveredAt: sql<
string | null
>`${jobExecutionLogs.executionData}->>'deadlineStopDeliveredAt'`,
billingRequestId: jobExecutionLogs.billingRequestId,
admittedBillingRequestId,
scheduleId: jobExecutionLogs.scheduleId,
Expand All @@ -89,8 +93,9 @@ function openJobRuns(status: OpenJobRunStatus, now: Date) {
* Drives every open occurrence to an end, so no crash leaves one queued or running forever:
* a queued record whose dispatch was lost is dispatched again under the same id (the turn is
* admitted at most once, so a retry of one that did start only finds it); a record whose turn
* ended is finalized; one past the time limit is stopped and finalized once it settles; one
* whose job was deleted before it started is ended. A run whose spend-cap link was lost (its
* ended is finalized; one past the time limit is stopped and finalized once it settles, its
* Stop delivered again first if an earlier pass died delivering it; one whose job was deleted
* before it started is ended. A run whose spend-cap link was lost (its
* process died between recording the billing admission and the link) gets it back.
*/
export async function reconcileJobRuns(now: Date): Promise<void> {
Expand Down Expand Up @@ -121,22 +126,23 @@ export async function reconcileJobRuns(now: Date): Promise<void> {
if (!row.billingRequestId && row.admittedBillingRequestId) {
await recordJobRunBillingRequest(row.executionId, row.admittedBillingRequestId)
}
if (TERMINAL_RUN_STATUSES.includes(row.runStatus)) {
await finalizeJobRun(row.executionId)
} else if (
const ended = TERMINAL_RUN_STATUSES.includes(row.runStatus)
/** A Stop whose delivery never finished is delivered before the run is finalized. */
const undelivered =
row.deadlineStopRequestedAt !== null && row.deadlineStopDeliveredAt === null
const overdue =
!ended &&
!row.deadlineStopRequestedAt &&
row.userId &&
row.workspaceId &&
row.chatId &&
now.getTime() - row.startedAt.getTime() >= JOB_RUN_TIME_LIMIT_MS
) {
if ((undelivered || overdue) && row.userId && row.workspaceId && row.chatId) {
await stopAtTimeLimit({
executionId: row.executionId,
userId: row.userId,
workspaceId: row.workspaceId,
chatId: row.chatId,
})
}
if (ended) await finalizeJobRun(row.executionId)
} catch (error) {
logger.warn('Job run could not be reconciled; the next pass retries', {
executionId: row.executionId,
Expand All @@ -161,13 +167,12 @@ async function stopAtTimeLimit(run: {
})
/** A run outside the job's scope records no Stop, so nothing is delivered to it. */
if (!stopped || stopped.workspaceId !== run.workspaceId || stopped.chatId !== run.chatId) return
/** Marked once the Stop is durable, so a pass that fails before it stops the run again. */
await db
.update(jobExecutionLogs)
.set({
executionData: sql`jsonb_set(${jobExecutionLogs.executionData}, '{deadlineStopRequestedAt}', to_jsonb(now()))`,
})
.where(eq(jobExecutionLogs.executionId, run.executionId))
/**
* Why the run stopped is recorded before delivery, so its summary survives a crash; delivery
* is marked separately once done, and a pass that finds it unmarked delivers again (requesting
* and delivering a Stop are idempotent), so the run's client workflows and sandbox work end.
*/
await markDeadlineStop(run.executionId, 'deadlineStopRequestedAt')
await deliverRunStop(
{
id: stopped.id,
Expand All @@ -178,5 +183,19 @@ async function stopAtTimeLimit(run: {
},
readStoppedRunWork(run.executionId, run.userId)
)
await markDeadlineStop(run.executionId, 'deadlineStopDeliveredAt')
logger.info('Stopped a job run at its time limit', { executionId: run.executionId })
}

/** Records a step of the time-limit Stop once; a repeated pass keeps the first time. */
async function markDeadlineStop(
executionId: string,
step: 'deadlineStopRequestedAt' | 'deadlineStopDeliveredAt'
): Promise<void> {
await db
Comment thread
waleedlatif1 marked this conversation as resolved.
.update(jobExecutionLogs)
.set({
executionData: sql`jsonb_set(${jobExecutionLogs.executionData}, ${`{${step}}`}::text[], COALESCE(${jobExecutionLogs.executionData}->${step}, to_jsonb(now())))`,
})
.where(eq(jobExecutionLogs.executionId, executionId))
}
Loading
Loading