Skip to content

Commit 56e9f79

Browse files
committed
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.
1 parent cb879b9 commit 56e9f79

6 files changed

Lines changed: 234 additions & 56 deletions

File tree

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

Lines changed: 44 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -346,9 +346,9 @@ async function releaseChatLocks(fences: ChatLockFence[]): Promise<void> {
346346
)
347347
}
348348

349-
async function readSweepCursor(): Promise<string | undefined> {
349+
async function readSweepCursor(key = SWEEP_CURSOR_KEY): Promise<string | undefined> {
350350
try {
351-
return (await getRedisClient()?.get(SWEEP_CURSOR_KEY)) ?? undefined
351+
return (await getRedisClient()?.get(key)) ?? undefined
352352
} catch (error) {
353353
logger.warn('Orphaned-run sweep cursor is unreadable; starting from the first run', {
354354
error: getErrorMessage(error),
@@ -357,12 +357,12 @@ async function readSweepCursor(): Promise<string | undefined> {
357357
}
358358
}
359359

360-
async function writeSweepCursor(cursor: string | undefined): Promise<void> {
360+
async function writeSweepCursor(cursor: string | undefined, key = SWEEP_CURSOR_KEY): Promise<void> {
361361
try {
362362
const redis = getRedisClient()
363363
if (!redis) return
364-
if (cursor) await redis.set(SWEEP_CURSOR_KEY, cursor, 'EX', SWEEP_CURSOR_TTL_SECONDS)
365-
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)
366366
} catch (error) {
367367
logger.warn('Orphaned-run sweep cursor could not be saved', { error: getErrorMessage(error) })
368368
}
@@ -489,19 +489,30 @@ export async function settleRunWithoutAccess(run: UnownedRun): Promise<boolean>
489489
* worth checking yet.
490490
*/
491491
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'
492494

493495
/**
494496
* Runs an interactive controller abandoned that a takeover can still resume: unfinished and
495-
* not stopped, saved with their recovery intent, idle for a minute but within the recovery window, and
496-
* with no controller holding their chat lock. Empty while the locks cannot be read, since then
497-
* no run is provably abandoned.
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.
498503
*/
499-
export async function findAbandonedRecoverableRuns(limit: number): Promise<UnownedRun[]> {
504+
export async function findAbandonedRecoverableRuns(
505+
limit: number,
506+
{ pageSize = SWEEP_BATCH_SIZE, maxExamined = SWEEP_BATCH_SIZE } = {}
507+
): Promise<UnownedRun[]> {
500508
const abandoned: UnownedRun[] = []
501-
let after: string | undefined
502-
/** Pages past runs whose controller is alive, so they never crowd out an abandoned one. */
503-
while (abandoned.length < limit) {
504-
const candidates: UnownedRun[] = await db
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
505516
.select(unownedRunColumns)
506517
.from(copilotRuns)
507518
.where(
@@ -512,23 +523,35 @@ export async function findAbandonedRecoverableRuns(limit: number): Promise<Unown
512523
idleFor(ABANDONED_RUN_IDLE_MS),
513524
not(idleFor(ORPHANED_RUN_GRACE_MS)),
514525
not(stopRequested),
515-
after ? gt(copilotRuns.id, after) : undefined
526+
cursor ? gt(copilotRuns.id, cursor) : undefined,
527+
wrapped && start ? lte(copilotRuns.id, start) : undefined
516528
)
517529
)
518530
.orderBy(asc(copilotRuns.id))
519-
.limit(SWEEP_BATCH_SIZE)
520-
if (candidates.length === 0) break
521-
after = candidates[candidates.length - 1].id
531+
.limit(size)
532+
let locked: Set<string>
522533
try {
523-
const locked = await findStreamsHoldingChatLock(candidates)
524-
abandoned.push(...candidates.filter((run) => !locked.has(run.streamId)))
534+
locked = page.length > 0 ? await findStreamsHoldingChatLock(page) : new Set()
525535
} catch (error) {
526536
logger.warn('Chat stream ownership is unreadable; no run is taken over this round', {
527537
error: getErrorMessage(error),
528538
})
529539
return []
530540
}
531-
if (candidates.length < SWEEP_BATCH_SIZE) break
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+
}
532554
}
533-
return abandoned.slice(0, limit)
555+
await writeSweepCursor(cursor, ABANDONED_RUN_CURSOR_KEY)
556+
return abandoned
534557
}

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

Lines changed: 17 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1144,7 +1144,9 @@ export async function recordSimSandboxProcess(
11441144
throw new Error('Sandbox command admission is closed or unavailable')
11451145
}
11461146
const { id, sandboxId, sessionKey } = input.process
1147-
const entry = JSON.stringify({ [id]: { sandboxId, sessionKey, settled: false } })
1147+
const entry = JSON.stringify({
1148+
[id]: { sandboxId, sessionKey, settled: false, ownerToken: input.ownerToken },
1149+
})
11481150
const [row] = await tx
11491151
.update(copilotAsyncToolCalls)
11501152
.set({
@@ -1297,7 +1299,7 @@ export async function prepareWorkbenchAccess(
12971299
throw new Error('Workbench does not belong to this tool execution')
12981300
}
12991301
const [claim] = await tx
1300-
.select({ id: copilotAsyncToolCalls.id })
1302+
.select({ processes: copilotAsyncToolCalls.sandboxProcesses })
13011303
.from(copilotAsyncToolCalls)
13021304
.where(
13031305
and(
@@ -1386,7 +1388,19 @@ export async function prepareWorkbenchAccess(
13861388
)
13871389
)
13881390
)
1389-
const processes = previous.flatMap((row) =>
1391+
/**
1392+
* A re-run's earlier attempts may have left commands running on the remote workbench, which
1393+
* outlive the Sim process that started them; they are stopped like any earlier call's.
1394+
*/
1395+
const earlierAttempts = {
1396+
toolCallId: input.toolCallId,
1397+
processes: Object.fromEntries(
1398+
Object.entries(claim.processes).filter(
1399+
([, process]) => process.ownerToken !== input.ownerToken
1400+
)
1401+
),
1402+
}
1403+
const processes = [...previous, earlierAttempts].flatMap((row) =>
13901404
Object.entries(row.processes).flatMap(([id, process]) =>
13911405
process.settled
13921406
? []

‎apps/sim/lib/mothership/request/application/abandoned-runs.integration.ts‎

Lines changed: 121 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@ import { deleteWorkspaceFixture, insertWorkspaceFixture } from '@sim/db/testing/
77
* duplicate-send path does, then answers the tool resume that follows.
88
*/
99
import { authMock, authMockFns } from '@sim/testing/mocks/auth.mock'
10+
import { envFlagsMock, resetEnvFlagsMock, setEnvFlags } from '@sim/testing/mocks/env-flags.mock'
1011
import { afterAll, beforeAll, beforeEach, describe, expect, it, vi } from 'vitest'
1112

1213
const { redisUrl, inheritedEnv, worker } = await vi.hoisted(async () => {
@@ -47,12 +48,13 @@ const { redisUrl, inheritedEnv, worker } = await vi.hoisted(async () => {
4748
REDIS_URL: process.env.REDIS_URL,
4849
SIM_AGENT_API_URL: process.env.SIM_AGENT_API_URL,
4950
}
50-
process.env.REDIS_URL = url
51+
if (url) process.env.REDIS_URL = url
5152
process.env.SIM_AGENT_API_URL = `http://127.0.0.1:${port}`
5253
return { redisUrl: url, inheritedEnv, worker: Object.assign(worker, { server }) }
5354
})
5455

5556
vi.mock('@/lib/auth', () => authMock)
57+
vi.mock('@/lib/core/config/env-flags', () => envFlagsMock)
5658

5759
import { db } from '@sim/db'
5860
import {
@@ -74,20 +76,32 @@ import {
7476
EXHAUSTED_SIM_TOOL_MESSAGE,
7577
SIM_TOOL_EXECUTION_MAX_ATTEMPTS,
7678
} from '@/lib/mothership/async-runs/execution-lease'
77-
import { REVOKED_RUN_ERROR } from '@/lib/mothership/async-runs/orphaned-runs'
79+
import {
80+
ABANDONED_RUN_CURSOR_KEY,
81+
findAbandonedRecoverableRuns,
82+
REVOKED_RUN_ERROR,
83+
} from '@/lib/mothership/async-runs/orphaned-runs'
7884
import {
7985
claimToolExecution,
8086
claimWorkflowToolExecution,
87+
prepareWorkbenchAccess,
88+
recordSimSandboxProcess,
8189
revokeExhaustedSimToolExecutions,
8290
upsertAsyncToolCall,
8391
} from '@/lib/mothership/async-runs/repository'
92+
import { createTrustedCopilotPrincipal } from '@/lib/mothership/auth/application-delegation'
8493
import { chatPubSub } from '@/lib/mothership/chat-status'
8594
import { takeOverAbandonedRuns } from '@/lib/mothership/request/application/abandoned-runs'
95+
import {
96+
RUN_TAKEOVER_AUDIENCE,
97+
takeOverStreamRun,
98+
} from '@/lib/mothership/request/application/recover-stream'
8699
import {
87100
createProviderToolCallIdentity,
88101
scopeProviderToolCallId,
89102
} from '@/lib/mothership/request/go/tool-call-identity'
90103
import { chatStreamLockKey } from '@/lib/mothership/request/session/controller-lease'
104+
import { chatSandboxSessionKey } from '@/lib/mothership/tools/sandbox-session-key'
91105
import { GET as streamGET } from '@/app/api/copilot/chat/stream/route'
92106

93107
const ownerId = generateId()
@@ -119,12 +133,12 @@ function frameFor(streamId: string, chatId: string) {
119133
* chat, and the tool's execution lease lapsed with no result recorded.
120134
*/
121135
async function runAbandonedMidTool(
122-
options: { attempts?: number; leaseLive?: boolean; toolName?: string } = {}
136+
options: { attempts?: number; leaseLive?: boolean; toolName?: string; runId?: string } = {}
123137
) {
124138
const chatId = generateId()
125139
chatIds.push(chatId)
126140
const streamId = generateId()
127-
const runId = generateId()
141+
const runId = options.runId ?? generateId()
128142
const request = { message: 'Run the report', userId, messageId: streamId, chatId, workspaceId }
129143
await db
130144
.insert(copilotChats)
@@ -304,6 +318,8 @@ describe.runIf(Boolean(redisUrl))('taking over a run whose controller died', ()
304318

305319
beforeEach(async () => {
306320
worker.requests.length = 0
321+
resetEnvFlagsMock()
322+
await redis().del(ABANDONED_RUN_CURSOR_KEY)
307323
await db
308324
.insert(permissions)
309325
.values({
@@ -471,4 +487,105 @@ describe.runIf(Boolean(redisUrl))('taking over a run whose controller died', ()
471487
await db.update(copilotRuns).set({ status: 'cancelled' }).where(eq(copilotRuns.id, runId))
472488
}
473489
})
490+
491+
it('reaches an abandoned run behind live ones, examining a bounded window per scan', async () => {
492+
/** The lowest ids, so the scan meets the two live runs before the abandoned one. */
493+
const ids = [1, 2, 3].map((n) => `00000000-0000-4000-8000-00000000000${n}`)
494+
const live = [
495+
await runAbandonedMidTool({ runId: ids[0] }),
496+
await runAbandonedMidTool({ runId: ids[1] }),
497+
]
498+
const abandoned = await runAbandonedMidTool({ runId: ids[2] })
499+
for (const run of live) {
500+
await redis().set(chatStreamLockKey(run.chatId), `${run.streamId}\nlive`, 'EX', 60)
501+
}
502+
try {
503+
const bound = { pageSize: 2, maxExamined: 2 }
504+
expect(await findAbandonedRecoverableRuns(1, bound)).toEqual([])
505+
expect((await findAbandonedRecoverableRuns(1, bound)).map((run) => run.id)).toEqual([
506+
abandoned.runId,
507+
])
508+
} finally {
509+
for (const run of [...live, abandoned]) {
510+
await redis().del(chatStreamLockKey(run.chatId))
511+
await db
512+
.update(copilotRuns)
513+
.set({ status: 'cancelled' })
514+
.where(eq(copilotRuns.id, run.runId))
515+
}
516+
}
517+
})
518+
519+
it("hands a re-run the earlier attempt's workbench commands to stop first", async () => {
520+
const { chatId, runId, toolCallId } = await runAbandonedMidTool({ leaseLive: true })
521+
const sessionKey = chatSandboxSessionKey(chatId)
522+
try {
523+
await recordSimSandboxProcess({
524+
toolCallId,
525+
runId,
526+
userId,
527+
ownerToken: 'dead-controller',
528+
process: { id: 'earlier-command', sandboxId: 'sandbox-1', sessionKey },
529+
})
530+
await db
531+
.update(copilotAsyncToolCalls)
532+
.set({ executionLeaseExpiresAt: sql`clock_timestamp() - interval '1 second'` })
533+
.where(eq(copilotAsyncToolCalls.toolCallId, toolCallId))
534+
const successor = { toolCallId, runId, userId, ownerToken: 'successor' }
535+
expect(await claimToolExecution(successor)).toEqual({ outcome: 'claimed' })
536+
await recordSimSandboxProcess({
537+
...successor,
538+
process: { id: 'own-command', sandboxId: 'sandbox-1', sessionKey },
539+
})
540+
541+
const state = await prepareWorkbenchAccess({ ...successor, sessionKey })
542+
543+
expect(state.processes.map((process) => process.id)).toEqual(['earlier-command'])
544+
} finally {
545+
await db.update(copilotRuns).set({ status: 'cancelled' }).where(eq(copilotRuns.id, runId))
546+
}
547+
})
548+
549+
it('claims no run once its caller is shutting down', async () => {
550+
const { chatId, streamId, runId, toolCallId } = await runAbandonedMidTool()
551+
const before = await storedRun(runId)
552+
const shutdown = new AbortController()
553+
shutdown.abort()
554+
const principal = createTrustedCopilotPrincipal(
555+
{ userId, workspaceId, chatId, delegationId: `test:${runId}` },
556+
{ audience: RUN_TAKEOVER_AUDIENCE, ttlMs: 60_000 }
557+
)
558+
try {
559+
await expect(
560+
takeOverStreamRun.execute({
561+
principal,
562+
input: { streamId },
563+
request: { headers: new Headers(), signal: shutdown.signal },
564+
})
565+
).rejects.toThrow()
566+
567+
expect(await storedRun(runId)).toEqual(before)
568+
expect(requestsFor(streamId, '/api/mothership')).toHaveLength(0)
569+
expect((await storedTool(toolCallId)).attempts).toBe(1)
570+
expect(await redis().get(chatStreamLockKey(chatId))).toBeNull()
571+
} finally {
572+
await db.update(copilotRuns).set({ status: 'cancelled' }).where(eq(copilotRuns.id, runId))
573+
}
574+
})
575+
576+
it('leaves a run recoverable when its takeover fails for a reason other than access', async () => {
577+
/** Hosted recovery refuses a run saved without its billing admission. */
578+
setEnvFlags({ isHosted: true })
579+
const { streamId, runId } = await runAbandonedMidTool()
580+
const before = await storedRun(runId)
581+
try {
582+
await takeOverAbandonedRuns()
583+
584+
expect(await storedRun(runId)).toEqual(before)
585+
expect(requestsFor(streamId, '/api/streams/explicit-abort')).toHaveLength(0)
586+
expect(requestsFor(streamId, '/api/mothership')).toHaveLength(0)
587+
} finally {
588+
await db.update(copilotRuns).set({ status: 'cancelled' }).where(eq(copilotRuns.id, runId))
589+
}
590+
})
474591
})

0 commit comments

Comments
 (0)