Skip to content

Commit 0c6c375

Browse files
authored
fix(mothership): serialize headless run admission with chat soft delete (#9107)
1 parent ebd8efb commit 0c6c375

5 files changed

Lines changed: 153 additions & 39 deletions

File tree

‎apps/sim/lib/knowledge/application/slack-search/assistant.integration.ts‎

Lines changed: 17 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -139,8 +139,11 @@ function job() {
139139
}
140140
}
141141

142-
/** Runs one Slack turn in a fresh private chat and returns the run it recorded. */
143-
async function slackTurn() {
142+
/**
143+
* Runs one Slack turn in a fresh private chat and returns the run it recorded. `deletedAt`
144+
* stands for a delete committed after the turn resolved its chat.
145+
*/
146+
async function slackTurn(deletedAt?: Date) {
144147
const chatId = generateId()
145148
const turnId = generateId()
146149
hoisted.chat = { id: chatId, model: 'default' }
@@ -150,6 +153,7 @@ async function slackTurn() {
150153
userId: hoisted.userId,
151154
organizationId: hoisted.organizationId,
152155
type: 'mothership',
156+
deletedAt,
153157
})
154158
const controller = new AbortController()
155159
hoisted.controller = controller
@@ -237,6 +241,17 @@ describe('Slack Assistant run record', () => {
237241
expect(run.completedAt).not.toBeNull()
238242
})
239243

244+
it('never starts the worker for a turn whose chat was deleted before its run', async () => {
245+
mockRunHeadlessCopilotLifecycle.mockClear()
246+
247+
const { run, outcome } = await slackTurn(new Date())
248+
249+
expect(mockRunHeadlessCopilotLifecycle).not.toHaveBeenCalled()
250+
expect(outcome).toBeInstanceOf(Error)
251+
expect(run.status).toBe('cancelled')
252+
expect(run.toolAdmissionClosedAt).not.toBeNull()
253+
})
254+
240255
it('records a failed turn as an error', async () => {
241256
mockRunHeadlessCopilotLifecycle.mockResolvedValueOnce({
242257
success: false,

‎apps/sim/lib/knowledge/application/slack-search/assistant.ts‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -49,7 +49,7 @@ import {
4949
startAbortPoller,
5050
unregisterActiveStream,
5151
} from '@/lib/mothership/request/session/abort'
52-
import { isExplicitStopReason } from '@/lib/mothership/request/session/abort-reason'
52+
import { AbortReason, isExplicitStopReason } from '@/lib/mothership/request/session/abort-reason'
5353
import type { OrchestratorResult } from '@/lib/mothership/request/types'
5454
import { organizationRoutes } from '@/lib/navigation/paths'
5555
import { SlackSearchAssistantStream } from '@/lib/slack-search/assistant-stream'
@@ -206,6 +206,11 @@ export async function runSlackSearchAssistant(
206206
})
207207
if (!run) throw new Error('Could not persist Assistant execution')
208208
runId = run.id
209+
/** Admission found the turn stopped or its chat deleted: it ends as that Stop would. */
210+
if (run.status === 'cancelled') {
211+
controller.abort(AbortReason.UserStop)
212+
controller.signal.throwIfAborted()
213+
}
209214
const responseStream = new SlackSearchAssistantStream({
210215
token: secret.botToken,
211216
channel: job.message.channelId,

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

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -148,9 +148,10 @@ describe('run admission and early Stop', () => {
148148
[copilotRequestStops.workspaceId, scope.workspaceId],
149149
[copilotRequestStops.streamId, scope.streamId],
150150
]) {
151+
/** The chat row is read first; the Stop lookup follows it. */
151152
expect(
152153
hasMockCondition(
153-
dbChainMockFns.where.mock.calls[0]?.[0],
154+
dbChainMockFns.where.mock.calls[1]?.[0],
154155
(condition) =>
155156
condition.type === 'eq' && condition.left === column && condition.right === value
156157
)

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

Lines changed: 20 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -441,28 +441,31 @@ export async function createRunSegment(input: CreateRunSegmentInput) {
441441
)
442442
}
443443

444-
/** Caller holds the Stop/admission lock; its related chat writes commit with this run. */
444+
/**
445+
* Caller holds the Stop/admission lock; its related chat writes commit with this run. The
446+
* chat row is read FOR SHARE, after that lock as chat admission takes it, so a soft delete
447+
* either commits first and this run starts stopped, or waits and then stops this run.
448+
*/
445449
export async function insertRunSegment(tx: RunAdmissionTransaction, input: CreateRunSegmentInput) {
446-
let workspaceId = input.workspaceId
447-
let organizationId = input.organizationId
448-
if (!workspaceId && !organizationId) {
449-
const [chat] = await tx
450-
.select({
451-
workspaceId: copilotChats.workspaceId,
452-
organizationId: copilotChats.organizationId,
453-
})
454-
.from(copilotChats)
455-
.where(and(eq(copilotChats.id, input.chatId), eq(copilotChats.userId, input.userId)))
456-
.limit(1)
457-
workspaceId = chat?.workspaceId
458-
organizationId = chat?.organizationId
459-
}
450+
const [chat] = await tx
451+
.select({
452+
workspaceId: copilotChats.workspaceId,
453+
organizationId: copilotChats.organizationId,
454+
deletedAt: copilotChats.deletedAt,
455+
})
456+
.from(copilotChats)
457+
.where(and(eq(copilotChats.id, input.chatId), eq(copilotChats.userId, input.userId)))
458+
.for('share')
459+
const inheritsOwner = !input.workspaceId && !input.organizationId
460+
const workspaceId = inheritsOwner ? chat?.workspaceId : input.workspaceId
461+
const organizationId = inheritsOwner ? chat?.organizationId : input.organizationId
460462
const { table, predicate } = stopLocation({ ...input, workspaceId, organizationId })
461463
const [stop] = await tx
462464
.select({ stoppedAt: table.stoppedAt })
463465
.from(table)
464466
.where(predicate)
465467
.limit(1)
468+
const closedAt = stop?.stoppedAt ?? chat?.deletedAt
466469
const [run] = await tx
467470
.insert(copilotRuns)
468471
.values({
@@ -481,8 +484,8 @@ export async function insertRunSegment(tx: RunAdmissionTransaction, input: Creat
481484
provider: input.provider ?? null,
482485
requestContext: input.requestContext ?? {},
483486
desktopDeviceId: input.desktopDeviceId ?? null,
484-
status: stop ? 'cancelled' : (input.status ?? 'active'),
485-
...(stop ? { completedAt: sql`now()`, toolAdmissionClosedAt: stop.stoppedAt } : {}),
487+
status: closedAt ? 'cancelled' : (input.status ?? 'active'),
488+
...(closedAt ? { completedAt: sql`now()`, toolAdmissionClosedAt: closedAt } : {}),
486489
})
487490
.returning()
488491
if (!run) throw new Error('Run admission did not return its persisted identity')

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

Lines changed: 108 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -66,6 +66,7 @@ import { NextRequest } from 'next/server'
6666
import { closeRedisConnection, getRedisClient } from '@/lib/core/config/redis'
6767
import { OrchestrationError } from '@/lib/core/orchestration/types'
6868
import { SIM_TOOL_EXECUTION_VERSION } from '@/lib/mothership/async-runs/lifecycle'
69+
import { createRunSegment } from '@/lib/mothership/async-runs/repository'
6970
import { admitChatTurn } from '@/lib/mothership/chat/application/admit-turn'
7071
import { deleteCopilotChat, softDeleteMothershipChat } from '@/lib/mothership/chat/deletion'
7172
import { chatPubSub } from '@/lib/mothership/chat-status'
@@ -203,6 +204,37 @@ async function seedRun(
203204
return { userId: owner.userId, chatId, streamId, runId, ...scope }
204205
}
205206

207+
/** A chat with no run yet, as a task wake, inbox email, API or Slack turn finds it. */
208+
async function seedChat(userId: string, workspaceId: string) {
209+
const chatId = generateId()
210+
await db.insert(copilotChats).values({ id: chatId, userId, workspaceId, type: 'mothership' })
211+
return chatId
212+
}
213+
214+
/** Admits a headless turn's run the way those surfaces do, without the chat view's admission. */
215+
async function admitHeadlessRun(userId: string, chatId: string, workspaceId: string) {
216+
const streamId = generateId()
217+
testStreams.add(streamId)
218+
const run = await createRunSegment({
219+
id: generateId(),
220+
executionId: generateId(),
221+
chatId,
222+
userId,
223+
workspaceId,
224+
streamId,
225+
requestContext: { source: 'headless_lifecycle' },
226+
})
227+
return { run, seeded: { userId, chatId, streamId, runId: run.id, workspaceId } }
228+
}
229+
230+
async function toolAdmissionClosedAt(run: SeededRun) {
231+
const [row] = await db
232+
.select({ closedAt: copilotRuns.toolAdmissionClosedAt })
233+
.from(copilotRuns)
234+
.where(eq(copilotRuns.id, run.runId))
235+
return row?.closedAt ?? null
236+
}
237+
206238
/** What the worker's heartbeat reads, through the outbound control transport. */
207239
async function runControl(run: SeededRun) {
208240
const result = await executeSimControl({
@@ -294,6 +326,23 @@ async function whileHolding<T>(chatId: string, run: (holder: number) => Promise<
294326
}
295327
}
296328

329+
/** Queues admission and deletion behind a held chat row in the order `first` names. */
330+
async function queueOnChat<A, D>(
331+
chatId: string,
332+
first: 'admission' | 'deletion',
333+
admit: () => Promise<A>,
334+
remove: () => Promise<D>
335+
): Promise<{ admission: Promise<A>; deletion: Promise<D> }> {
336+
return whileHolding(chatId, async (holder) => {
337+
const admission = first === 'admission' ? admit() : undefined
338+
const deletion = admission ? undefined : remove()
339+
const waiter = await lockWaiterBehind(holder)
340+
const queued = { admission: admission ?? admit(), deletion: deletion ?? remove() }
341+
await lockWaiterBehind(holder, waiter)
342+
return queued
343+
})
344+
}
345+
297346
const STOPPED = { status: 200, body: { stopped: true } }
298347
const LIVE = { status: 200, body: { stopped: false } }
299348

@@ -674,12 +723,9 @@ describe.runIf(Boolean(redisUrl))('removing a run’s chat, owner or access stop
674723
async (first) => {
675724
const ownerId = await seedUser()
676725
const workspaceId = await seedWorkspace(ownerId)
677-
const chatId = generateId()
726+
const chatId = await seedChat(ownerId, workspaceId)
678727
const streamId = generateId()
679728
testStreams.add(streamId)
680-
await db
681-
.insert(copilotChats)
682-
.values({ id: chatId, userId: ownerId, workspaceId, type: 'mothership' })
683729
const sendClaim = { normalizedKey: `run-removal:${generateId()}`, claimToken: generateId() }
684730
await db
685731
.insert(idempotencyKey)
@@ -725,20 +771,7 @@ describe.runIf(Boolean(redisUrl))('removing a run’s chat, owner or access stop
725771
.finally(() => releasePendingChatStream(chatId, streamId, lease))
726772
const remove = () => softDeleteMothershipChat({ chatId, userId: ownerId })
727773

728-
const { admission, deletion } = await whileHolding(chatId, async (holder) => {
729-
if (first === 'admission') {
730-
const admission = admit()
731-
const waiter = await lockWaiterBehind(holder)
732-
const deletion = remove()
733-
await lockWaiterBehind(holder, waiter)
734-
return { admission, deletion }
735-
}
736-
const deletion = remove()
737-
const waiter = await lockWaiterBehind(holder)
738-
const admission = admit()
739-
await lockWaiterBehind(holder, waiter)
740-
return { admission, deletion }
741-
})
774+
const { admission, deletion } = await queueOnChat(chatId, first, admit, remove)
742775

743776
const admitted = await admission
744777
expect(await deletion).not.toBeNull()
@@ -751,4 +784,61 @@ describe.runIf(Boolean(redisUrl))('removing a run’s chat, owner or access stop
751784
expect(await runControl(run)).toEqual(STOPPED)
752785
}
753786
)
787+
788+
it('admits a headless run of a live chat as live', async () => {
789+
const ownerId = await seedUser()
790+
const workspaceId = await seedWorkspace(ownerId)
791+
792+
const { run, seeded } = await admitHeadlessRun(
793+
ownerId,
794+
await seedChat(ownerId, workspaceId),
795+
workspaceId
796+
)
797+
798+
expect(run).toMatchObject({ status: 'active', toolAdmissionClosedAt: null, completedAt: null })
799+
expect(await runControl(seeded)).toEqual(LIVE)
800+
})
801+
802+
it('admits a headless run of a chat deleted before it as already stopped', async () => {
803+
const ownerId = await seedUser()
804+
const workspaceId = await seedWorkspace(ownerId)
805+
const chatId = await seedChat(ownerId, workspaceId)
806+
expect(await softDeleteMothershipChat({ chatId, userId: ownerId })).not.toBeNull()
807+
808+
const { run, seeded } = await admitHeadlessRun(ownerId, chatId, workspaceId)
809+
810+
expect(run.status).toBe('cancelled')
811+
expect(run.toolAdmissionClosedAt).not.toBeNull()
812+
expect(run.completedAt).not.toBeNull()
813+
expect(await runControl(seeded)).toEqual(STOPPED)
814+
})
815+
816+
/**
817+
* A headless run is admitted without writing the chat row, so only its read of the chat
818+
* serializes it with a deletion: holding the row queues both in a chosen order.
819+
*/
820+
it.each(['admission', 'deletion'] as const)(
821+
'never leaves a headless run live while its chat is being deleted, when %s reaches the chat first',
822+
async (first) => {
823+
const ownerId = await seedUser()
824+
const workspaceId = await seedWorkspace(ownerId)
825+
const chatId = await seedChat(ownerId, workspaceId)
826+
const admit = () => admitHeadlessRun(ownerId, chatId, workspaceId)
827+
const remove = () => softDeleteMothershipChat({ chatId, userId: ownerId })
828+
829+
const { admission, deletion } = await queueOnChat(chatId, first, admit, remove)
830+
831+
const { run, seeded } = await admission
832+
expect(await deletion).not.toBeNull()
833+
if (first === 'admission') {
834+
await expect.poll(pushedStops).toEqual([seeded.streamId])
835+
expect(await hasStop(seeded)).toBe(true)
836+
} else {
837+
expect(run.status).toBe('cancelled')
838+
expect(pushedStops()).toEqual([])
839+
}
840+
expect(await toolAdmissionClosedAt(seeded)).not.toBeNull()
841+
expect(await runControl(seeded)).toEqual(STOPPED)
842+
}
843+
)
754844
})

0 commit comments

Comments
 (0)