Skip to content

Commit 2e98120

Browse files
authored
fix(mothership): take over runs whose controller died and re-run their interrupted tool (#9059)
* fix(mothership): take over runs whose controller died and re-run their interrupted tool A Chat run is driven by the Sim process that admitted it. When that process stops mid-run (a deploy's old task set is terminated, or a crash) and no browser is attached, nothing takes the run over: recovery only ran from a browser reconnect, the in-flight tool's result was never recorded, and the worker eventually ended the run as detached. - Every Sim process now looks every 20s for interactive runs whose chat lock is gone and takes each over through the same recovery a reconnect runs, acting as the run's user under a delegated principal that re-checks that user's current access. The chat lock and controller token still admit exactly one controller, so a racing reconnect or another process loses. - A user who lost access gets the run settled as an error and the worker is told to stop it. - A Sim tool whose owner died before recording a result (its execution lease lapsed while the call was still running) is claimed again under a new owner token and run again, by the server for a workflow tool. Each start is counted; after three, the call fails with a clear message instead. - A takeover waits while one of the run's Sim tools still holds its lease. * fix(mothership): import the takeover by absolute path and regenerate the schema mock * fix(mothership): page past live controllers, stop a scan on shutdown, dispose chat pubsub in tests * test(mothership): assert the claim admits a first execution or an interrupted re-run * test(mothership): add execution_attempts to the browser download fixture table * fix(mothership): bound the takeover scan, stop earlier attempts' commands, decline takeovers on shutdown - The abandoned-run scan resumes from a durable cursor and examines a bounded window per round, so runs whose controller is alive cannot hide later runs. - A re-run's earlier attempts' workbench commands are handed to the workbench recovery to stop before the re-run may use the workbench; each recorded command now names the attempt that started it. - Shutdown aborts a takeover still in flight at its controller claim. - Only an authorization denial settles a run as revoked; other takeover failures leave it recoverable, as a failed reconnect does. - The suite sets REDIS_URL only when a test Redis is configured. * fix(mothership): keep the takeover stop handler module-private * fix(mothership): let the server re-claim a workflow call its dead claimant never started * fix(mothership): harden abandoned-run takeover - Add the `mothership-abandoned-run-takeover-disabled` AppConfig kill switch, checked once per scan; off by default, so takeovers run unless it is on. - Stop a controller whose chat lock lapsed while every renewal threw, so a controller partitioned from Redis cannot keep driving a run taken over. - Never revoke another call awaiting its re-run from a sibling's wait or workbench access; the sibling's workbench access still stops that call's dead attempt's commands. - Treat a workspace access loss found inside recovery like an authorization denial, by error type; other recovery failures stay recoverable. - Cover organization Assistant takeover for a member and a removed member. * fix(mothership): negate the re-runnable predicate without an optional SQL operand * fix(mothership): time the lock lapse from its acquisition, type the re-run predicates as conditions, revoke access mid-recovery in the test * test(mothership): find the attempt cap inside the typed re-run conjunction
1 parent 02cd2f6 commit 2e98120

22 files changed

Lines changed: 33110 additions & 208 deletions

‎apps/sim/instrumentation-node.ts‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -414,4 +414,9 @@ export async function register() {
414414

415415
const { startServiceUsageReplay } = await import('./lib/mothership/billing/service-delivery')
416416
startServiceUsageReplay()
417+
418+
const { startAbandonedRunTakeover } = await import(
419+
'@/lib/mothership/request/application/abandoned-runs'
420+
)
421+
startAbandonedRunTakeover()
417422
}

‎apps/sim/lib/core/config/feature-flags.ts‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -120,6 +120,13 @@ const FEATURE_FLAGS = {
120120
'targeting. Off-AppConfig falls back to KNOWLEDGE_MEMBER_ACCESS.',
121121
fallback: 'KNOWLEDGE_MEMBER_ACCESS',
122122
},
123+
'mothership-abandoned-run-takeover-disabled': {
124+
description:
125+
'Kill switch for Sim taking over Chat runs whose controller died (a deploy or a crash) and ' +
126+
're-running their interrupted tool. Global on/off only; on stops every takeover scan. ' +
127+
'Off by default, so takeovers run unless this is turned on.',
128+
fallback: null,
129+
},
123130
} satisfies Record<string, FeatureFlagDefinition>
124131

125132
/**

‎apps/sim/lib/desktop/executor/bound-turn.integration.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -511,7 +511,7 @@ describe.runIf(Boolean(redisUrl))("a turn bound to a desktop's background execut
511511

512512
await lapse(toolCallId, 'lease')
513513
/** A Sim tool waiting on the same run sweeps its expired executions; this one is not its. */
514-
await revokeExpiredSimToolExecutions({ runId: run.runId, userId })
514+
await revokeExpiredSimToolExecutions({ runId: run.runId, userId, toolCallId: 'waiting-call' })
515515
expect((await storedCall(toolCallId)).status).toBe('running')
516516
await answer
517517
const late = await desktop.complete(toolCallId, executionToken, { output: 'built' })

‎apps/sim/lib/mothership/async-runs/browser-download-claim.integration.ts‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -52,6 +52,7 @@ describe('browser download admission with PostgreSQL', () => {
5252
claimed_at timestamp, claimed_by text, execution_started_at timestamp,
5353
execution_settled_at timestamp, execution_owner_token text,
5454
execution_lease_expires_at timestamptz, execution_revoked_at timestamptz,
55+
execution_attempts integer NOT NULL DEFAULT 0,
5556
client_workflow_execution_id text, sandbox_processes jsonb NOT NULL DEFAULT '{}',
5657
completed_at timestamp, created_at timestamp NOT NULL DEFAULT now(),
5758
updated_at timestamp NOT NULL DEFAULT now()

‎apps/sim/lib/mothership/async-runs/execution-lease.ts‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,9 +9,18 @@ export interface SimToolExecutionOwner {
99
ownerToken: string
1010
}
1111

12+
/**
13+
* Sim executions one call may start: the first, and a re-run each time a takeover finds the
14+
* previous owner died before recording a result. A tool that kills its executor every time
15+
* ends with {@link EXHAUSTED_SIM_TOOL_MESSAGE} instead of re-running forever.
16+
*/
17+
export const SIM_TOOL_EXECUTION_MAX_ATTEMPTS = 3
18+
1219
export const INTERRUPTED_SIM_TOOL_MESSAGE =
1320
'Tool execution lost its owner before a result was recorded. Its outcome is unknown; inspect current state before retrying a mutation.'
1421

22+
export const EXHAUSTED_SIM_TOOL_MESSAGE = `Tool execution was interrupted ${SIM_TOOL_EXECUTION_MAX_ATTEMPTS} times because the Sim server running it stopped, so it was not run again. Its outcome is unknown; inspect current state before retrying a mutation.`
23+
1524
export class SimToolExecutionLeaseLostError extends Error {
1625
constructor() {
1726
super(INTERRUPTED_SIM_TOOL_MESSAGE)

‎apps/sim/lib/mothership/async-runs/orphaned-runs.ts‎

Lines changed: 112 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -73,6 +73,8 @@ export const ORPHANED_RUN_GRACE_MS = 60 * 60 * 1000
7373
export const LEGACY_RUN_GRACE_MS = 24 * 60 * 60 * 1000
7474

7575
export const ORPHANED_RUN_ERROR = 'This response was interrupted before it finished.'
76+
export const REVOKED_RUN_ERROR =
77+
'This response was stopped because you no longer have access to continue it.'
7678
export const LEGACY_RUN_ERROR = 'Run was never finalized (pre-lease run).'
7779

7880
const SWEEP_BATCH_SIZE = 500
@@ -124,7 +126,7 @@ const stopRequested = sql`(EXISTS (SELECT 1 FROM ${copilotRequestStops} s
124126
WHERE s.user_id = ${copilotRuns.userId} AND s.organization_id = ${copilotRuns.organizationId}
125127
AND s.stream_id = ${copilotRuns.streamId}))`
126128

127-
interface UnownedRun {
129+
export interface UnownedRun {
128130
id: string
129131
chatId: string
130132
streamId: string
@@ -144,9 +146,15 @@ const unownedRunColumns = {
144146
controllerToken,
145147
}
146148

149+
const SETTLED_RUN_ERRORS = {
150+
orphaned: ORPHANED_RUN_ERROR,
151+
revoked: REVOKED_RUN_ERROR,
152+
legacy: LEGACY_RUN_ERROR,
153+
} as const
154+
147155
/** A run ends cancelled only if its user pressed Stop, whoever settles it. */
148-
function terminalValues(reason: 'orphaned' | 'legacy') {
149-
const error = reason === 'orphaned' ? ORPHANED_RUN_ERROR : LEGACY_RUN_ERROR
156+
function terminalValues(reason: keyof typeof SETTLED_RUN_ERRORS) {
157+
const error = SETTLED_RUN_ERRORS[reason]
150158
return {
151159
status: sql`(CASE WHEN ${stopRequested} THEN 'cancelled' ELSE 'error' END)::copilot_run_status`,
152160
error: sql`CASE WHEN ${stopRequested} THEN NULL ELSE ${error}::text END`,
@@ -171,7 +179,8 @@ function terminalValues(reason: 'orphaned' | 'legacy') {
171179
async function settleRuns(
172180
tx: DbTransaction,
173181
runs: UnownedRun[],
174-
guard: SQL | undefined
182+
guard: SQL | undefined,
183+
reason: Exclude<keyof typeof SETTLED_RUN_ERRORS, 'legacy'> = 'orphaned'
175184
): Promise<{ settled: UnownedRun[]; released: UnownedRun[] }> {
176185
if (runs.length === 0) return { settled: [], released: [] }
177186
const chatIds = [...new Set(runs.map((run) => run.chatId))]
@@ -235,7 +244,7 @@ async function settleRuns(
235244
for (const run of runs) {
236245
if (run.controllerToken === null) continue
237246
await apply([run], eq(controllerToken, run.controllerToken), {
238-
...terminalValues('orphaned'),
247+
...terminalValues(reason),
239248
completedAt: sql`now()`,
240249
updatedAt: sql`now()`,
241250
})
@@ -337,9 +346,9 @@ async function releaseChatLocks(fences: ChatLockFence[]): Promise<void> {
337346
)
338347
}
339348

340-
async function readSweepCursor(): Promise<string | undefined> {
349+
async function readSweepCursor(key = SWEEP_CURSOR_KEY): Promise<string | undefined> {
341350
try {
342-
return (await getRedisClient()?.get(SWEEP_CURSOR_KEY)) ?? undefined
351+
return (await getRedisClient()?.get(key)) ?? undefined
343352
} catch (error) {
344353
logger.warn('Orphaned-run sweep cursor is unreadable; starting from the first run', {
345354
error: getErrorMessage(error),
@@ -348,12 +357,12 @@ async function readSweepCursor(): Promise<string | undefined> {
348357
}
349358
}
350359

351-
async function writeSweepCursor(cursor: string | undefined): Promise<void> {
360+
async function writeSweepCursor(cursor: string | undefined, key = SWEEP_CURSOR_KEY): Promise<void> {
352361
try {
353362
const redis = getRedisClient()
354363
if (!redis) return
355-
if (cursor) await redis.set(SWEEP_CURSOR_KEY, cursor, 'EX', SWEEP_CURSOR_TTL_SECONDS)
356-
else await redis.del(SWEEP_CURSOR_KEY)
364+
if (cursor) await redis.set(key, cursor, 'EX', SWEEP_CURSOR_TTL_SECONDS)
365+
else await redis.del(key)
357366
} catch (error) {
358367
logger.warn('Orphaned-run sweep cursor could not be saved', { error: getErrorMessage(error) })
359368
}
@@ -453,3 +462,96 @@ export async function settleStoppedRunWithoutController(runId: string): Promise<
453462
announceReleased(released)
454463
return settled.length > 0
455464
}
465+
466+
/**
467+
* Settles a run its controller abandoned as an error when its user may no longer continue
468+
* it. The run's chat lock is fenced as a takeover's would be, so a reconnect cannot claim
469+
* the run while it settles; a run whose lock is held, or whose controller changed, is left.
470+
*/
471+
export async function settleRunWithoutAccess(run: UnownedRun): Promise<boolean> {
472+
if (!run.controllerToken) return false
473+
const fences = await fenceChatLocks([run])
474+
if (fences.length === 0) return false
475+
try {
476+
const { settled, released } = await db.transaction((tx) =>
477+
settleRuns(tx, [run], undefined, 'revoked')
478+
)
479+
announceReleased(released)
480+
return settled.length > 0
481+
} finally {
482+
await releaseChatLocks(fences)
483+
}
484+
}
485+
486+
/**
487+
* A run its controller is still driving writes its status at least once per leg, and a dead
488+
* controller's chat lock lapses within this long, so a run written to more recently is not
489+
* worth checking yet.
490+
*/
491+
const ABANDONED_RUN_IDLE_MS = 60_000
492+
/** Where the last takeover scan stopped, so live runs it skipped never hide the runs after them. */
493+
export const ABANDONED_RUN_CURSOR_KEY = 'copilot:abandoned-runs:scan-cursor'
494+
495+
/**
496+
* Runs an interactive controller abandoned that a takeover can still resume: unfinished and
497+
* not stopped, saved with their recovery intent, idle for a minute but within the recovery
498+
* window, and with no controller holding their chat lock. Each scan examines at most
499+
* `maxExamined` candidates, resuming after the last one the previous scan examined and wrapping
500+
* to the first, so runs whose controller is alive cost one examination per pass and never
501+
* crowd out an abandoned run. Empty while the locks cannot be read, since then no run is
502+
* provably abandoned.
503+
*/
504+
export async function findAbandonedRecoverableRuns(
505+
limit: number,
506+
{ pageSize = SWEEP_BATCH_SIZE, maxExamined = SWEEP_BATCH_SIZE } = {}
507+
): Promise<UnownedRun[]> {
508+
const abandoned: UnownedRun[] = []
509+
const start = await readSweepCursor(ABANDONED_RUN_CURSOR_KEY)
510+
let cursor = start
511+
let wrapped = start === undefined
512+
let examined = 0
513+
while (abandoned.length < limit && examined < maxExamined) {
514+
const size = Math.min(pageSize, maxExamined - examined)
515+
const page: UnownedRun[] = await db
516+
.select(unownedRunColumns)
517+
.from(copilotRuns)
518+
.where(
519+
and(
520+
inArray(copilotRuns.status, UNFINISHED_RUN_STATUSES),
521+
isNotNull(controllerToken),
522+
sql`${copilotRuns.requestContext}->'recovery'->>'kind' = 'interactive_stream'`,
523+
idleFor(ABANDONED_RUN_IDLE_MS),
524+
not(idleFor(ORPHANED_RUN_GRACE_MS)),
525+
not(stopRequested),
526+
cursor ? gt(copilotRuns.id, cursor) : undefined,
527+
wrapped && start ? lte(copilotRuns.id, start) : undefined
528+
)
529+
)
530+
.orderBy(asc(copilotRuns.id))
531+
.limit(size)
532+
let locked: Set<string>
533+
try {
534+
locked = page.length > 0 ? await findStreamsHoldingChatLock(page) : new Set()
535+
} catch (error) {
536+
logger.warn('Chat stream ownership is unreadable; no run is taken over this round', {
537+
error: getErrorMessage(error),
538+
})
539+
return []
540+
}
541+
for (const run of page) {
542+
examined++
543+
cursor = run.id
544+
if (locked.has(run.streamId)) continue
545+
abandoned.push(run)
546+
if (abandoned.length === limit) break
547+
}
548+
if (abandoned.length < limit && page.length < size) {
549+
/** The end of the table: wrap once to cover the runs before the starting point. */
550+
cursor = undefined
551+
if (wrapped) break
552+
wrapped = true
553+
}
554+
}
555+
await writeSweepCursor(cursor, ABANDONED_RUN_CURSOR_KEY)
556+
return abandoned
557+
}

‎apps/sim/lib/mothership/async-runs/repository.test.ts‎

Lines changed: 35 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,13 @@ import {
55
copilotRequestStops,
66
copilotRuns,
77
} from '@sim/db/schema'
8-
import { dbChainMockFns, hasMockCondition, queueTableRows, resetDbChainMock } from '@sim/testing'
8+
import {
9+
dbChainMockFns,
10+
flattenMockConditions,
11+
hasMockCondition,
12+
queueTableRows,
13+
resetDbChainMock,
14+
} from '@sim/testing'
915
import { beforeEach, describe, expect, it } from 'vitest'
1016
import {
1117
areStreamToolExecutionsSettled,
@@ -185,6 +191,21 @@ describe('run admission and early Stop', () => {
185191
})
186192
})
187193

194+
/** Finds a condition anywhere in a mocked expression, including inside `sql` and `sql.join` fragments. */
195+
function someNestedCondition(
196+
node: unknown,
197+
predicate: (condition: Record<string, unknown>) => boolean
198+
): boolean {
199+
if (!node || typeof node !== 'object') return false
200+
const record = node as Record<string, unknown>
201+
if (predicate(record)) return true
202+
return ['conditions', 'values', 'fragments'].some(
203+
(key) =>
204+
Array.isArray(record[key]) &&
205+
(record[key] as unknown[]).some((child) => someNestedCondition(child, predicate))
206+
)
207+
}
208+
188209
describe('durable Sim tool ownership', () => {
189210
beforeEach(() => {
190211
resetDbChainMock()
@@ -240,14 +261,26 @@ describe('durable Sim tool ownership', () => {
240261
(condition) => condition.type === 'eq' && condition.right === input.userId
241262
)
242263
).toBe(true)
264+
/** A first execution, or a re-run of one whose owner died before it recorded a result. */
265+
const claimable = flattenMockConditions(dbChainMockFns.where.mock.calls[1]?.[0]).find(
266+
(condition) => condition.type === 'or'
267+
)
268+
const [firstExecution, interruptedExecution] = (claimable?.conditions ?? []) as unknown[]
243269
expect(
244270
hasMockCondition(
245-
dbChainMockFns.where.mock.calls[1]?.[0],
271+
firstExecution,
246272
(condition) =>
247273
condition.type === 'isNull' &&
248274
condition.column === copilotAsyncToolCalls.executionStartedAt
249275
)
250276
).toBe(true)
277+
expect(
278+
someNestedCondition(
279+
interruptedExecution,
280+
(condition) =>
281+
condition.type === 'lt' && condition.left === copilotAsyncToolCalls.executionAttempts
282+
)
283+
).toBe(true)
251284
expect(dbChainMockFns.set).toHaveBeenCalledWith(
252285
expect.objectContaining({ executionStartedAt: expect.any(Date), status: 'running' })
253286
)

0 commit comments

Comments
 (0)