Skip to content

Commit ac20e3f

Browse files
committed
fix(mothership): keep a withdrawn first message ahead of its follow-ups on remount
Holding a dying chatless queue for the next mount let that mount send a follow-up before the first message's cross-surface handoff arrived, so the two went out in reverse order. When follow-ups are queued, the unmount now puts the withdrawn first message at their head, under its own id, and skips its separate handoff; both reach the same surface in the order written. A first message after a Stop is not withdrawn, so only the follow-up goes out.
1 parent 999b5eb commit ac20e3f

2 files changed

Lines changed: 133 additions & 28 deletions

File tree

‎apps/sim/app/workspace/[workspaceId]/home/hooks/use-chat.dom.test.tsx‎

Lines changed: 82 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -2053,31 +2053,92 @@ describe('useChat remount send recovery', () => {
20532053
})
20542054

20552055
/**
2056-
* A follow-up sent on the new-chat surface while the first message's Stop is
2057-
* pending waits in that surface's queue until the first message's chat is
2058-
* known. If the Stop fails and the surface remounts first, the queue key dies
2059-
* with the mount: the next new-chat surface must show the follow-up (and may
2060-
* send it), not leave it under the dead key, neither sent nor visible.
2056+
* The first POST on the new-chat surface never answers; later POSTs open a
2057+
* turn in the chat the first message created. The abort endpoint and the
2058+
* stream lookup fail, as they would for a Stop that cannot reach the server.
20612059
*/
2062-
it('keeps a follow-up queued behind a failed chatless Stop when the surface remounts', async () => {
2060+
function stubFirstPostPendingThenAdmitted() {
20632061
vi.stubGlobal('fetch', async (input: RequestInfo | URL, init?: RequestInit) => {
20642062
const url = String(input)
20652063
if (url === '/api/mothership/chat' && init?.method === 'POST') {
20662064
state.postBodies.push(JSON.parse(String(init.body)))
2067-
return new Promise<Response>((_, reject) => {
2068-
init.signal?.addEventListener('abort', () => reject(init.signal?.reason), { once: true })
2069-
})
2065+
if (state.postBodies.length === 1) {
2066+
return new Promise<Response>((_, reject) => {
2067+
init.signal?.addEventListener('abort', () => reject(init.signal?.reason), {
2068+
once: true,
2069+
})
2070+
})
2071+
}
2072+
return new Response(
2073+
new ReadableStream<Uint8Array>({
2074+
start(controller) {
2075+
controller.close()
2076+
},
2077+
}),
2078+
{
2079+
status: 200,
2080+
headers: {
2081+
'Content-Type': 'text/event-stream',
2082+
'x-mothership-chat-id': DEDUPED_CHAT_ID,
2083+
},
2084+
}
2085+
)
20702086
}
20712087
if (url.includes('/api/copilot/chat/abort')) {
20722088
state.abortBodies.push(JSON.parse(String(init?.body)))
20732089
return Response.json({ error: 'Internal error' }, { status: 500 })
20742090
}
2075-
if (url.includes('/api/mothership/chat/stream')) {
2091+
if (url.includes('/api/mothership/chat/stream') && state.postBodies.length === 1) {
20762092
return Response.json({ error: 'Internal error' }, { status: 500 })
20772093
}
20782094
return fetchStub(input, init)
20792095
})
2080-
const first = renderUseChat()
2096+
}
2097+
2098+
/**
2099+
* A follow-up typed on the new-chat surface while the first message waits for
2100+
* the server is queued behind it. If the surface remounts then, the first
2101+
* message is withdrawn and both must reach the next mount in the order they
2102+
* were written: the first message (under its own id), then the follow-up.
2103+
*/
2104+
it('sends a withdrawn first message before its follow-up when the new-chat surface remounts', async () => {
2105+
stubFirstPostPendingThenAdmitted()
2106+
const first = renderHomeLikeSurface()
2107+
await act(async () => {
2108+
void first.getResult().sendMessage('inspect the workspace')
2109+
})
2110+
await waitFor(() => state.postBodies.length === 1)
2111+
await act(async () => {
2112+
void first.getResult().sendMessage('follow-up while admission pending')
2113+
})
2114+
await waitFor(() => allQueuedMessages().length === 1)
2115+
first.unmount()
2116+
2117+
const second = renderHomeLikeSurface()
2118+
await waitFor(() => state.postBodies.length >= 3, 4_000)
2119+
await act(async () => {
2120+
await sleep(300)
2121+
})
2122+
2123+
const afterRemount = state.postBodies.slice(1)
2124+
expect(afterRemount.map((body) => body.message)).toEqual([
2125+
'inspect the workspace',
2126+
'follow-up while admission pending',
2127+
])
2128+
expect(afterRemount[0].userMessageId).toBe(state.postBodies[0].userMessageId)
2129+
expect(afterRemount[1].chatId).toBe(DEDUPED_CHAT_ID)
2130+
expect(second.claimedByOwnListener()).toBe(0)
2131+
expect(allQueuedMessages()).toHaveLength(0)
2132+
})
2133+
2134+
/**
2135+
* After a Stop of the first message, only the follow-up was the user's
2136+
* intent: the Stop's POST is left to the server, nothing withdraws it, and the
2137+
* next mount sends just the follow-up, once.
2138+
*/
2139+
it('sends only the follow-up after a failed Stop when the new-chat surface remounts', async () => {
2140+
stubFirstPostPendingThenAdmitted()
2141+
const first = renderHomeLikeSurface()
20812142
await act(async () => {
20822143
void first.getResult().sendMessage('inspect the workspace')
20832144
})
@@ -2092,15 +2153,16 @@ describe('useChat remount send recovery', () => {
20922153
})
20932154
first.unmount()
20942155

2095-
const second = renderUseChat()
2096-
await waitFor(
2097-
() =>
2098-
second
2099-
.getResult()
2100-
.messageQueue.some((message) => message.content === 'Sent while the Stop was failing') ||
2101-
state.postBodies.some((body) => body.message === 'Sent while the Stop was failing'),
2102-
4_000
2103-
)
2156+
renderHomeLikeSurface()
2157+
await waitFor(() => state.postBodies.length >= 2, 4_000)
2158+
await act(async () => {
2159+
await sleep(500)
2160+
})
2161+
2162+
expect(state.postBodies.slice(1).map((body) => body.message)).toEqual([
2163+
'Sent while the Stop was failing',
2164+
])
2165+
expect(allQueuedMessages()).toHaveLength(0)
21042166
})
21052167

21062168
it('stopping a chat preserves an unrelated manual workflow execution', async () => {

‎apps/sim/app/workspace/[workspaceId]/home/hooks/use-chat.ts‎

Lines changed: 51 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -267,6 +267,8 @@ interface PendingChatAdmission {
267267
chatKey: string
268268
controller: AbortController
269269
settled: Promise<string | undefined>
270+
/** What an unmount would withdraw, if it ran before the server answered. */
271+
send: WithdrawnSend
270272
}
271273

272274
/** A send an unmount cleanup withdrew, as handed to the next chat surface. */
@@ -772,6 +774,8 @@ export function useChat(
772774
const heldSendSurface = `${scopeKey}:${options?.workflowId ?? 'home'}`
773775
const heldSendSurfaceRef = useRef(heldSendSurface)
774776
heldSendSurfaceRef.current = heldSendSurface
777+
/** Withdrawn first messages the unmount queued itself, so their handoff is skipped. */
778+
const withdrawnHeldAtUnmountRef = useRef<Set<string> | null>(null)
775779
const onToolResultRef = useRef(options?.onToolResult)
776780
onToolResultRef.current = options?.onToolResult
777781
const onTitleUpdateRef = useRef(options?.onTitleUpdate)
@@ -3784,6 +3788,17 @@ export function useChat(
37843788
settled: new Promise((resolve) => {
37853789
resolveAdmission = resolve
37863790
}),
3791+
send: {
3792+
content: message,
3793+
userMessageId,
3794+
...(fileAttachments ? { fileAttachments } : {}),
3795+
...(contexts ? { contexts } : {}),
3796+
...(options?.requestMode ? { requestMode: options.requestMode } : {}),
3797+
...(options?.assistantSearch ? { assistantSearch: options.assistantSearch } : {}),
3798+
...(options?.assistantSearchLevel !== undefined
3799+
? { assistantSearchLevel: options.assistantSearchLevel }
3800+
: {}),
3801+
},
37873802
}
37883803
pendingChatAdmissionRef.current = admission
37893804
}
@@ -4243,6 +4258,8 @@ export function useChat(
42434258
*/
42444259
const handOffWithdrawnSend = useCallback(
42454260
(send: WithdrawnSend) => {
4261+
/** The unmount already queued it ahead of its follow-ups; see the unmount cleanup. */
4262+
if (withdrawnHeldAtUnmountRef.current?.delete(send.userMessageId)) return
42464263
if (
42474264
sendMothershipMessage(
42484265
send.content,
@@ -5305,6 +5322,40 @@ export function useChat(
53055322

53065323
useEffect(() => {
53075324
return () => {
5325+
/* A chatless mount's queue key dies with it, so messages still queued there
5326+
go to the next mount of this surface, as held sends do. A first message
5327+
this unmount withdraws (its POST not yet answered, and not stopped) goes
5328+
at their head: the follow-ups were written after it, and the next mount
5329+
would otherwise send them before its handoff arrives. Alone, it keeps
5330+
the usual cross-surface handoff. */
5331+
const deadKey = chatKeyRef.current
5332+
if (deadKey.startsWith(PENDING_CHAT_KEY_PREFIX)) {
5333+
const queueStore = useMothershipQueueStore.getState()
5334+
const withdrawing = pendingChatAdmissionRef.current
5335+
if (
5336+
withdrawing &&
5337+
withdrawing.chatKey === deadKey &&
5338+
abortControllerRef.current === withdrawing.controller &&
5339+
(queueStore.queues[deadKey]?.length ?? 0) > 0
5340+
) {
5341+
const { send } = withdrawing
5342+
queueStore.insertAt(deadKey, 0, {
5343+
id: generateId(),
5344+
content: send.content,
5345+
resumeUserMessageId: send.userMessageId,
5346+
...(send.fileAttachments ? { fileAttachments: send.fileAttachments } : {}),
5347+
...(send.contexts ? { contexts: send.contexts } : {}),
5348+
...(send.requestMode ? { requestMode: send.requestMode } : {}),
5349+
...(send.assistantSearch ? { assistantSearch: send.assistantSearch } : {}),
5350+
...(send.assistantSearchLevel !== undefined
5351+
? { assistantSearchLevel: send.assistantSearchLevel }
5352+
: {}),
5353+
})
5354+
withdrawnHeldAtUnmountRef.current ??= new Set()
5355+
withdrawnHeldAtUnmountRef.current.add(send.userMessageId)
5356+
}
5357+
queueStore.holdForSurface(deadKey, heldSendSurfaceRef.current)
5358+
}
53085359
cancelActiveStreamRecovery()
53095360
clearQueueDispatchState()
53105361
streamGenRef.current++
@@ -5319,14 +5370,6 @@ export function useChat(
53195370
sendingRef.current = false
53205371
// Release the editing slot — the composer it binds to is unmounting.
53215372
useMothershipQueueStore.getState().setEditing(chatKeyRef.current, null)
5322-
/* A chatless mount's queue key dies with it. Messages still queued there
5323-
(behind a first message whose chat is not known yet, or a Stop that
5324-
failed) go to the next mount of this surface, as held sends do. */
5325-
if (chatKeyRef.current.startsWith(PENDING_CHAT_KEY_PREFIX)) {
5326-
useMothershipQueueStore
5327-
.getState()
5328-
.holdForSurface(chatKeyRef.current, heldSendSurfaceRef.current)
5329-
}
53305373
}
53315374
}, [
53325375
cancelActiveStreamRecovery,

0 commit comments

Comments
 (0)