fix(chat): deduplicate chat sends server-side instead of probing for them (#6536)

* fix(chat): deduplicate chat sends server-side instead of probing for them

A client cannot tell whether a request it aborted reached the server: the chat
route never reads `request.signal`, so an accepted one still opens the chat,
persists the user message, and bills the turn after the socket drops. #6525
answered that by polling the orphaned stream before retrying — a 2.5s guess
that had to distinguish "no such stream" from "we stopped looking", and still
left a window open.

The codebase already owns the right tool. `IdempotencyService` backs webhook,
polling, and billing dedup, and `billingIdempotency` exists for exactly this
hazard: "a retry would double-record usage — real money". Chat sends now claim
the same way, keyed on the client-generated `userMessageId` and scoped to the
caller so nobody can probe another user's sends. A repeat gets 409 naming the
chat the first attempt opened — deliberately the shape the pending-stream lock
already returns, so the client's existing conflict handler reattaches instead
of starting a turn, with only the chat-adoption line added.

The claim fails open at every step. Deduplication saves a duplicate chat; the
send IS the user's message, so an unreachable bookkeeping store degrades chat
rather than taking it down. It is released when a send fails before recording a
chat, and deliberately kept once recorded.

Retrying now just reuses the id, which deletes the probe outright: the poll and
its two constants, the three-state result, the epoch plumbing that kept a
superseded poll from re-sending, and the chat-adoption branch it needed. The
client hook nets 67 lines smaller.

Idle sends go back to calling `startSendMessage` directly. #6525 routed them
through the durable queue so recovery had a backing entry, which put every
message in the product through the queue store, sessionStorage, and the
dispatch loop for the sake of a rare path — and the recovery never needed it,
since the message, attachments, contexts, and id are all in scope at the abort.
Both callers now share one `handOffWithdrawnSend`.

`startSendMessage` takes its optional tail as an options object; it was at six
positional parameters and the retry id would have been a seventh.

Tests cover both halves: the server dedups, scopes the key per user, records
the chat, and still sends when the claim store is down; the client reuses the
original id on retry and adopts the chat a deduplicated retry names. Each was
confirmed red without its fix.

* fix(chat): keep a withdrawn send in its own chat, and release stranded claims

Audit follow-ups, two of them real defects in the previous commit.

A withdrawn send routed unconditionally through the cross-surface lanes. Those
deliver to whatever chat is mounted next, so sending in one chat and switching
to another re-sent the message into the second one. The dispatcher already drew
the distinction; the idle path now draws it too — a chat-bound key is the stable
chat id, so re-queueing under it both retries durably and keeps the message
where the user put it. Only a chatless key, which dies with its mount, goes to
the lanes.

The claim release sat in `catch`, so the two paths that return a response
without throwing — a rejected branch, and a missing chat — stranded an
in-progress claim for its full 60s TTL, and a retry inside that window got a
spurious "already sent" instead of the real error. Moved to `finally`.

Also: `userMessageId` is now length-bounded, since it becomes part of a Postgres
key and an oversized one would throw inside the claim; `requestId` was still
empty at claim time, so both dedup logs printed a blank prefix; the provider
segment said `mothership` on a handler that also serves the workflow copilot,
and now says what the key identifies; `retryFailures` was dead config, only read
by `executeWithIdempotency`, which this caller never invokes; the doc pointed at
`billingIdempotency`, which has no consumers, and now points at the live Stripe
analogue.

Trimmed: `sendClaimRecorded` folded into clearing `sendClaim`, the unread `kind`
discriminant dropped from a one-arm union, the single-use `claimedChatId`
inlined, and the prose on all three of those cut back to what the code does not
already say.

* fix(chat): make a send's claim permanent only once its turn starts

The claim became permanent as soon as the chat resolved, but three exits still
return without starting a turn — a rejected branch, a missing chat, and a
pending-stream collision. The last one matters: the queued-send-handoff path
deliberately retries under the original `userMessageId` after a collision, and
against a permanent claim that retry deduplicated to a chat whose turn never
ran, reattaching to a stream that does not exist. A send that had merely
collided became unsendable for the claim's full hour.

The claim is now dropped immediately before the stream response is returned, so
`finally` releases it on every other exit. Recording the chat still happens as
early as possible — a concurrent duplicate needs somewhere to go — it just no
longer implies the turn happened.

* refactor(chat): give the send claim a single point of permanence

Recording the chat also dropped the claim when it failed, which left a second
way for a claim to stop being tracked and a compound hole behind it: a failed
record followed by a throw stranded the claim for its in-progress TTL, and a
retry inside that window reattached to a turn that never started.

Only one line now decides permanence — the claim is cleared immediately before
the stream response — so `finally` releases it on every exit that did not start
a turn, including a failed record. The `recorded` flag is gone with it.

Covers the 400 early return with a release assertion: that path returns without
throwing, so it is the one that proves the release has to live in `finally`.
This commit is contained in:
Waleed
2026-08-11 07:58:40 -07:00
committed by GitHub
parent 783e1b542c
commit f8644cc679
11 changed files with 698 additions and 488 deletions
@@ -342,7 +342,7 @@ export function Home({ chatId, userName, userId, tableViewsEnabled }: HomeProps)
if (!detail?.message) return
e.preventDefault()
sendMessage(detail.message, detail.fileAttachments, detail.contexts, {
...(detail.recoverStreamId ? { recoverStreamId: detail.recoverStreamId } : {}),
...(detail.resumeUserMessageId ? { resumeUserMessageId: detail.resumeUserMessageId } : {}),
})
}
window.addEventListener(MOTHERSHIP_SEND_MESSAGE_EVENT, handler)
@@ -373,7 +373,9 @@ export function Home({ chatId, userName, userId, tableViewsEnabled }: HomeProps)
if (!handoff) return
if (handoff.message) {
sendMessage(handoff.message, handoff.fileAttachments, handoff.contexts, {
...(handoff.recoverStreamId ? { recoverStreamId: handoff.recoverStreamId } : {}),
...(handoff.resumeUserMessageId
? { resumeUserMessageId: handoff.resumeUserMessageId }
: {}),
})
return
}
@@ -3,18 +3,20 @@
*
* Regression tests for the remount send loss: a send started on a fresh chat
* surface was silently dropped when the hook's unmount cleanup ran mid-flight
* and aborted the POST. Two things run that cleanup while an auto-send from a
* cross-route handoff is still in flight — StrictMode's dev double-mount, and a
* real client-side navigation away — and because `MothershipHandoffStorage`
* consumes atomically, the second mount finds nothing left to retry.
* and aborted the POST. Two things run that cleanup while an auto-send is still
* in flight — the chat route's `key={chatId}` remount when the user switches
* chats, and StrictMode's dev double-mount — and because
* `MothershipHandoffStorage` consumes atomically, the replacement mount finds
* nothing left to retry.
*
* (A Suspense hide/reveal does NOT cause this: React 19 disappears layout
* effects only, so this passive cleanup never runs for it.)
*
* The fix routes idle sends through the durable queue so every send has a
* recoverable entry, and recovers one the cleanup withdrew — probing the
* orphaned stream first so a request the server had already accepted is
* adopted rather than sent twice.
* Recovery hands the message to the next surface carrying the original
* `userMessageId`. Reusing that id is what makes the retry safe: the server
* deduplicates it against the first attempt rather than opening a second chat
* and billing a second turn, so the client never has to guess whether the
* request it aborted was accepted.
*/
import { act, type ReactNode, StrictMode, useEffect } from 'react'
import { sleep } from '@sim/utils/helpers'
@@ -42,29 +44,20 @@ import { MothershipHandoffStorage } from '@/lib/core/utils/browser-storage'
import { useChat } from '@/app/workspace/[workspaceId]/home/hooks/use-chat'
import { useMothershipQueueStore } from '@/stores/mothership-queue/store'
const DEDUPED_CHAT_ID = 'chat-the-first-attempt-opened'
interface NetworkState {
/** How the chat POST behaves for the next call. */
postBehavior: 'hang' | 'accept'
postCalls: number
/**
* How the orphaned-stream probe answers:
* - `found` — the server accepted the withdrawn request and owns a chat
* - `gone` — 404, it has no such stream (never accepted)
* - `pending` — registered but no owner yet, so the probe keeps polling;
* this is the only mode that leaves a probe in flight to interrupt
* How the chat POST behaves:
* - `hang` — accepted but never answered, the window the cleanup abort lands in
* - `accept` — a normal streaming response
* - `deduped` — the 409 the server returns for an already-claimed send
*/
probeBehavior: 'found' | 'gone' | 'pending'
orphanedStreamChatId: string
streamProbes: number
postBehavior: 'hang' | 'accept' | 'deduped'
postBodies: Array<{ message: string; userMessageId?: string }>
}
const state: NetworkState = {
postBehavior: 'hang',
postCalls: 0,
probeBehavior: 'gone',
orphanedStreamChatId: 'chat-server-already-made',
streamProbes: 0,
}
const state: NetworkState = { postBehavior: 'hang', postBodies: [] }
/** An SSE response whose stream ends immediately without a terminal event. */
function emptySseResponse(): Response {
@@ -79,29 +72,28 @@ function emptySseResponse(): Response {
async function fetchStub(input: RequestInfo | URL, init?: RequestInit): Promise<Response> {
const url = String(input instanceof Request ? input.url : input)
// The orphaned-stream probe: does the server hold a stream for the send the
// cleanup abort withdrew?
/* Stream replay, used by the reconnect a deduplicated send falls into.
`complete` is the terminal status the hook recognises — anything else and
reconnect polls forever. */
if (url.includes('/api/mothership/chat/stream')) {
state.streamProbes++
// 404 is what the server returns for a stream it never registered — i.e.
// the request really was withdrawn before it was accepted.
if (state.probeBehavior === 'gone') {
return new Response(JSON.stringify({ error: 'stream gone' }), { status: 404 })
}
return new Response(
JSON.stringify({
success: true,
events: [],
status: 'streaming',
// `pending` omits the owner, so the probe keeps polling.
...(state.probeBehavior === 'found' ? { chatId: state.orphanedStreamChatId } : {}),
}),
{ status: 200, headers: { 'Content-Type': 'application/json' } }
)
return new Response(JSON.stringify({ success: true, events: [], status: 'complete' }), {
status: 200,
headers: { 'Content-Type': 'application/json' },
})
}
if (url.includes('/api/mothership/chat') && init?.method === 'POST') {
state.postCalls++
state.postBodies.push(JSON.parse(String(init.body)))
if (state.postBehavior === 'deduped') {
return new Response(
JSON.stringify({
error: 'This message was already sent.',
activeStreamId: state.postBodies.at(-1)?.userMessageId,
chatId: DEDUPED_CHAT_ID,
}),
{ status: 409, headers: { 'Content-Type': 'application/json' } }
)
}
if (state.postBehavior === 'accept') return emptySseResponse()
return new Promise<Response>((_, reject) => {
const signal = init?.signal
@@ -154,49 +146,48 @@ function renderUseChat(): {
}
/**
* Mounts the hook under StrictMode with a handoff already in storage, mirroring
* `home.tsx`'s consume-and-auto-send effect. This is the production-shaped
* failure: the dev double-mount runs the passive cleanup between the two
* mounts, aborting the in-flight POST, and `consume` has already cleared the
* entry so the second mount has nothing to replay.
* As `renderUseChat`, but bound to an existing chat rather than chatless. The
* pathname has to match: the hook resets a chat-bound surface back to a fresh
* pending key when it finds itself on the home route.
*/
function renderStrictModeHandoffConsumer(): { unmount: () => void } {
function renderUseChatInChat(chatId: string): {
getResult: () => ReturnType<typeof useChat>
unmount: () => void
} {
navigationMocks.usePathname.mockReturnValue(`/workspace/ws-1/chat/${chatId}`)
;(globalThis as { IS_REACT_ACT_ENVIRONMENT?: boolean }).IS_REACT_ACT_ENVIRONMENT = true
queryClient = new QueryClient({ defaultOptions: { queries: { retry: false } } })
const container = document.createElement('div')
const root = createRoot(container)
mountedRoots.push(root)
let result: ReturnType<typeof useChat> | undefined
function Probe() {
const { sendMessage } = useChat('ws-1', undefined)
useEffect(() => {
const handoff = MothershipHandoffStorage.consume('ws-1')
if (!handoff?.message) return
sendMessage(handoff.message, handoff.fileAttachments, handoff.contexts, {
...(handoff.recoverStreamId ? { recoverStreamId: handoff.recoverStreamId } : {}),
})
}, [sendMessage])
result = useChat('ws-1', chatId)
return null
}
act(() => {
root.render(
<StrictMode>
<QueryClientProvider client={queryClient}>{(<Probe />) as ReactNode}</QueryClientProvider>
</StrictMode>
<QueryClientProvider client={queryClient}>{(<Probe />) as ReactNode}</QueryClientProvider>
)
})
return { unmount: () => act(() => root.unmount()) }
return {
getResult: () => {
if (result === undefined) throw new Error('Hook result is not ready')
return result
},
unmount: () => act(() => root.unmount()),
}
}
/**
* Mounts a surface shaped like `home.tsx`: it drives `useChat` AND registers
* the `mothership-send-message` listener that claims the event with
* `preventDefault`. Unmounting this exercises the ordering question — whether
* the departing surface's own still-attached listener can claim the recovery
* event its own teardown emitted, which would suppress the storage fallback
* and strand the message.
* Mounts a surface shaped like `home.tsx`: it drives `useChat` AND registers the
* `mothership-send-message` listener that claims the event with
* `preventDefault`. Unmounting it exercises whether the departing surface's own
* still-attached listener can claim the recovery event its teardown emitted,
* which would suppress the storage fallback and strand the message.
*/
function renderHomeLikeSurface(): {
getResult: () => ReturnType<typeof useChat>
@@ -215,16 +206,18 @@ function renderHomeLikeSurface(): {
const chat = useChat('ws-1', undefined)
result = chat
const { sendMessage } = chat
// Mirrors home.tsx:339 — declared AFTER useChat, so on unmount React runs
// Mirrors home.tsx — declared AFTER useChat, so on unmount React runs
// useChat's cleanup (which aborts) before this removeEventListener.
useEffect(() => {
const handler = (e: Event) => {
const detail = (e as CustomEvent<{ message?: string; recoverStreamId?: string }>).detail
const detail = (e as CustomEvent<{ message?: string; resumeUserMessageId?: string }>).detail
if (!detail?.message) return
claims++
e.preventDefault()
sendMessage(detail.message, undefined, undefined, {
...(detail.recoverStreamId ? { recoverStreamId: detail.recoverStreamId } : {}),
...(detail.resumeUserMessageId
? { resumeUserMessageId: detail.resumeUserMessageId }
: {}),
})
}
window.addEventListener('mothership-send-message', handler)
@@ -249,6 +242,43 @@ function renderHomeLikeSurface(): {
}
}
/**
* Mounts the hook under StrictMode with a handoff already in storage, mirroring
* `home.tsx`'s consume-and-auto-send effect. This is the production-shaped
* failure: the dev double-mount runs the passive cleanup between the two
* mounts, aborting the in-flight POST, and `consume` has already cleared the
* entry so the second mount has nothing to replay.
*/
function renderStrictModeHandoffConsumer(): void {
;(globalThis as { IS_REACT_ACT_ENVIRONMENT?: boolean }).IS_REACT_ACT_ENVIRONMENT = true
queryClient = new QueryClient({ defaultOptions: { queries: { retry: false } } })
const container = document.createElement('div')
const root = createRoot(container)
mountedRoots.push(root)
function Probe() {
const { sendMessage } = useChat('ws-1', undefined)
useEffect(() => {
const handoff = MothershipHandoffStorage.consume('ws-1')
if (!handoff?.message) return
sendMessage(handoff.message, handoff.fileAttachments, handoff.contexts, {
...(handoff.resumeUserMessageId
? { resumeUserMessageId: handoff.resumeUserMessageId }
: {}),
})
}, [sendMessage])
return null
}
act(() => {
root.render(
<StrictMode>
<QueryClientProvider client={queryClient}>{(<Probe />) as ReactNode}</QueryClientProvider>
</StrictMode>
)
})
}
/** Every queued message across all chat keys, flattened. */
function allQueuedMessages() {
return Object.values(useMothershipQueueStore.getState().queues).flat()
@@ -267,10 +297,9 @@ async function waitFor(predicate: () => boolean, budgetMs = 2000): Promise<void>
describe('useChat remount send recovery', () => {
beforeEach(() => {
vi.stubGlobal('fetch', fetchStub)
navigationMocks.usePathname.mockReturnValue('/workspace/ws-1/home')
state.postBehavior = 'hang'
state.postCalls = 0
state.probeBehavior = 'gone'
state.streamProbes = 0
state.postBodies = []
mockRequestJson.mockResolvedValue({ chats: [] })
useMothershipQueueStore.setState({ queues: {}, editing: {} })
window.sessionStorage.clear()
@@ -286,7 +315,23 @@ describe('useChat remount send recovery', () => {
vi.clearAllMocks()
})
it('delivers an aborted chatless send directly to a live replacement surface', async () => {
it('keeps a cross-route handoff recoverable across a StrictMode double-mount', async () => {
MothershipHandoffStorage.store({ message: 'investigate this failed run' }, 'ws-1')
renderStrictModeHandoffConsumer()
await waitFor(() => state.postBodies.length >= 1)
// Something must still hold the message: the live event was claimed and it
// is in flight again, or it is back in storage for the next mount.
await waitFor(
() =>
window.localStorage.getItem('sim_mothership_handoff') !== null ||
allQueuedMessages().length > 0 ||
state.postBodies.length > 1
)
})
it('delivers a withdrawn chatless send to a live replacement surface', async () => {
const attachment = {
id: 'file-1',
key: 'uploads/file-1',
@@ -294,10 +339,13 @@ describe('useChat remount send recovery', () => {
media_type: 'text/plain',
size: 12,
}
const received: Array<{ message: string; fileAttachments?: unknown[] }> = []
const received: Array<{
message: string
fileAttachments?: unknown[]
resumeUserMessageId?: string
}> = []
const claim = (event: Event) => {
const detail = (event as CustomEvent<{ message: string; fileAttachments?: unknown[] }>).detail
received.push(detail)
received.push((event as CustomEvent<(typeof received)[number]>).detail)
event.preventDefault()
}
window.addEventListener('mothership-send-message', claim)
@@ -307,20 +355,22 @@ describe('useChat remount send recovery', () => {
await act(async () => {
void getResult().sendMessage('hello from the palette', [attachment])
})
await waitFor(() => state.postCalls === 1)
await waitFor(() => state.postBodies.length === 1)
unmount()
await waitFor(() => received.length === 1)
expect(received[0].message).toBe('hello from the palette')
expect(received[0].fileAttachments).toEqual([attachment])
// Carried so the replacement retries as the same send, not a new one.
expect(received[0].resumeUserMessageId).toBe(state.postBodies[0].userMessageId)
expect(window.localStorage.getItem('sim_mothership_handoff')).toBeNull()
} finally {
window.removeEventListener('mothership-send-message', claim)
}
})
it('re-persists an aborted chatless send as a handoff for the next mount', async () => {
it('re-persists a withdrawn chatless send as a handoff for the next mount', async () => {
const attachment = {
id: 'file-2',
key: 'uploads/file-2',
@@ -333,16 +383,11 @@ describe('useChat remount send recovery', () => {
await act(async () => {
void getResult().sendMessage('hello from the palette', [attachment])
})
await waitFor(() => state.postCalls === 1)
await waitFor(() => state.postBodies.length === 1)
// The dispatch claimed the queue head when the optimistic send applied.
// An idle send goes straight out — it never occupies the queue.
expect(allQueuedMessages()).toHaveLength(0)
// The cleanup abort (the same code path a StrictMode remount or a real
// navigation away runs) fires while the POST is still awaiting the server.
// A chatless surface regenerates its queue key per mount, so recovery
// re-persists the send as a one-shot handoff for the next mount's consumer
// instead of restoring the dead instance's queue.
unmount()
await waitFor(() => window.localStorage.getItem('sim_mothership_handoff') !== null)
@@ -350,16 +395,17 @@ describe('useChat remount send recovery', () => {
const handoff = MothershipHandoffStorage.consume('ws-1')
expect(handoff?.message).toBe('hello from the palette')
expect(handoff?.fileAttachments).toEqual([attachment])
expect(handoff?.resumeUserMessageId).toBe(state.postBodies[0].userMessageId)
})
it('does not re-queue a send the server already received', async () => {
it('does not recover a send the server already answered', async () => {
state.postBehavior = 'accept'
const { getResult, unmount } = renderUseChat()
await act(async () => {
void getResult().sendMessage('already accepted')
})
await waitFor(() => state.postCalls === 1)
await waitFor(() => state.postBodies.length === 1)
unmount()
await act(async () => {
@@ -373,15 +419,15 @@ describe('useChat remount send recovery', () => {
/**
* A departing surface's own listener must not claim the recovery event its
* teardown emitted: claiming returns `true`, which suppresses the storage
* fallback, and the enqueue would land under the disposed pending key — the
* message would be stranded exactly where this fix is supposed to save it.
* fallback, and the message would be stranded exactly where this fix is meant
* to save it.
*/
it('does not let a departing surface claim its own recovery event', async () => {
const surface = renderHomeLikeSurface()
await act(async () => {
void surface.getResult().sendMessage('must survive my own teardown')
})
await waitFor(() => state.postCalls === 1)
await waitFor(() => state.postBodies.length === 1)
surface.unmount()
await waitFor(() => window.localStorage.getItem('sim_mothership_handoff') !== null)
@@ -390,166 +436,77 @@ describe('useChat remount send recovery', () => {
expect(MothershipHandoffStorage.consume('ws-1')?.message).toBe('must survive my own teardown')
})
/**
* The end-to-end failure, driven by the thing that actually runs the cleanup
* mid-flight rather than by a hand-rolled unmount. On the unfixed hook the
* handoff is consumed, the POST is aborted, and nothing survives to retry.
*/
it('keeps a cross-route handoff recoverable across a StrictMode double-mount', async () => {
MothershipHandoffStorage.store({ message: 'investigate this failed run' }, 'ws-1')
describe('retrying a withdrawn send', () => {
/**
* The whole point of carrying the id: the server sees one logical send, so
* it deduplicates instead of opening a second chat and billing again.
*/
it('reuses the original message id so the server can deduplicate', async () => {
const { getResult, unmount } = renderUseChat()
await act(async () => {
void getResult().sendMessage('only bill me once')
})
await waitFor(() => state.postBodies.length === 1)
unmount()
await waitFor(() => window.localStorage.getItem('sim_mothership_handoff') !== null)
renderStrictModeHandoffConsumer()
await waitFor(() => state.postCalls >= 1)
const handoff = MothershipHandoffStorage.consume('ws-1')
const replacement = renderUseChat()
await act(async () => {
void replacement.getResult().sendMessage(handoff?.message as string, undefined, undefined, {
resumeUserMessageId: handoff?.resumeUserMessageId as string,
})
})
await waitFor(() => state.postBodies.length === 2)
// Something must still be holding the message: either the live event was
// claimed and it is queued/in flight again, or it is back in storage.
await waitFor(() => {
const stored = window.localStorage.getItem('sim_mothership_handoff')
return stored !== null || allQueuedMessages().length > 0 || state.postCalls > 1
expect(state.postBodies[1].userMessageId).toBe(state.postBodies[0].userMessageId)
})
/**
* When the first attempt did reach the server, the retry comes back 409
* naming the chat it opened. The client adopts that chat rather than
* starting another turn.
*/
it('adopts the chat a deduplicated retry names', async () => {
state.postBehavior = 'deduped'
const { getResult } = renderUseChat()
await act(async () => {
void getResult().sendMessage('this one already landed', undefined, undefined, {
resumeUserMessageId: 'the-first-attempt',
})
})
await waitFor(() => getResult().resolvedChatId === DEDUPED_CHAT_ID)
expect(state.postBodies).toHaveLength(1)
expect(state.postBodies[0].userMessageId).toBe('the-first-attempt')
})
})
/**
* The abort tears down the client socket but the route handler never reads
* `request.signal` — a request the server had already accepted still creates
* the chat, persists the user message, and runs (and bills) the turn. So the
* recovered send has to ask whether that happened before sending again.
* A withdrawn send belongs to the chat it was sent to. The cross-surface
* lanes deliver to whatever chat is mounted next, so routing a chat-bound
* send through them would drop the message into a different conversation —
* exactly what happens if the user switches chats mid-send. Its key is the
* stable chat id, so re-queueing under that key is the durable retry.
*/
describe('recovered send probes the orphaned stream before re-sending', () => {
it('adopts the chat the server already created instead of sending twice', async () => {
// The server accepted the withdrawn request and registered its stream.
state.probeBehavior = 'found'
it('re-queues a withdrawn chat-bound send instead of following the user', async () => {
const { getResult, unmount } = renderUseChatInChat('chat-a')
const { getResult, unmount } = renderUseChat()
await act(async () => {
void getResult().sendMessage('only once please')
})
await waitFor(() => state.postCalls === 1)
unmount()
await waitFor(() => window.localStorage.getItem('sim_mothership_handoff') !== null)
// The next mount consumes the handoff, exactly as home.tsx does.
const handoff = MothershipHandoffStorage.consume('ws-1')
expect(handoff?.recoverStreamId).toBeTruthy()
const replacement = renderUseChat()
await act(async () => {
void replacement.getResult().sendMessage(handoff?.message as string, undefined, undefined, {
recoverStreamId: handoff?.recoverStreamId as string,
})
})
await waitFor(() => state.streamProbes > 0)
await waitFor(() => replacement.getResult().resolvedChatId === 'chat-server-already-made')
expect(state.postCalls).toBe(1)
expect(allQueuedMessages()).toHaveLength(0)
await act(async () => {
void getResult().sendMessage('belongs to chat-a')
})
await waitFor(() => state.postBodies.length === 1)
/**
* Adoption alone does not surface the running turn: hydration reconnects
* only when `chatHistory.activeStreamId` is set, and that query is cached
* for 30s. Without an explicit detail invalidation the adopted chat renders
* with the live response invisible.
*/
it('invalidates the adopted chat detail so hydration can reconnect', async () => {
state.probeBehavior = 'found'
unmount()
await waitFor(() => allQueuedMessages().length === 1)
const { getResult, unmount } = renderUseChat()
await act(async () => {
void getResult().sendMessage('surface the running turn')
})
await waitFor(() => state.postCalls === 1)
unmount()
await waitFor(() => window.localStorage.getItem('sim_mothership_handoff') !== null)
const handoff = MothershipHandoffStorage.consume('ws-1')
const replacement = renderUseChat()
const invalidateSpy = vi.spyOn(queryClient, 'invalidateQueries')
await act(async () => {
void replacement.getResult().sendMessage(handoff?.message as string, undefined, undefined, {
recoverStreamId: handoff?.recoverStreamId as string,
})
})
await waitFor(() =>
invalidateSpy.mock.calls.some(([arg]) => {
const key = (arg as { queryKey?: unknown[] } | undefined)?.queryKey
return Array.isArray(key) && key.includes('chat-server-already-made')
})
)
expect(state.postCalls).toBe(1)
})
/**
* A probe cut short by unmount answers "unknown", not "safe to send".
* Falling through to `startSendMessage` there would open a POST whose
* abort controller the teardown already dropped — an uncancellable request
* that duplicates the send. The entry must stay queued instead.
*/
it('does not send when the probe is cut short by unmount', async () => {
state.probeBehavior = 'gone'
const { getResult, unmount } = renderUseChat()
await act(async () => {
void getResult().sendMessage('do not zombie me')
})
await waitFor(() => state.postCalls === 1)
unmount()
await waitFor(() => window.localStorage.getItem('sim_mothership_handoff') !== null)
const handoff = MothershipHandoffStorage.consume('ws-1')
/* `pending` keeps the probe polling instead of answering on the first
attempt, which is what leaves one in flight to interrupt. A `gone`
probe answers immediately and the re-send would already have happened
before the unmount — that version of this test cannot fail. */
state.probeBehavior = 'pending'
const replacement = renderUseChat()
await act(async () => {
void replacement.getResult().sendMessage(handoff?.message as string, undefined, undefined, {
recoverStreamId: handoff?.recoverStreamId as string,
})
})
await waitFor(() => state.streamProbes > 0)
const postsBeforeUnmount = state.postCalls
replacement.unmount()
// Well past the probe's poll budget: nothing may send after teardown.
await act(async () => {
await sleep(3000)
})
expect(state.postCalls).toBe(postsBeforeUnmount)
/* Not sending is only half of it — the message must still be
recoverable. A chatless surface's `pending::` key is regenerated per
mount, so leaving the entry there would strand it just as surely as
re-sending would have duplicated it. */
expect(MothershipHandoffStorage.consume('ws-1')?.message).toBe('do not zombie me')
expect(allQueuedMessages()).toHaveLength(0)
})
it('re-sends when the server has no stream for it', async () => {
// 404 from the probe: the request really was withdrawn before acceptance.
state.probeBehavior = 'gone'
const { getResult, unmount } = renderUseChat()
await act(async () => {
void getResult().sendMessage('please actually send me')
})
await waitFor(() => state.postCalls === 1)
unmount()
await waitFor(() => window.localStorage.getItem('sim_mothership_handoff') !== null)
const handoff = MothershipHandoffStorage.consume('ws-1')
const replacement = renderUseChat()
await act(async () => {
void replacement.getResult().sendMessage(handoff?.message as string, undefined, undefined, {
recoverStreamId: handoff?.recoverStreamId as string,
})
})
await waitFor(() => state.postCalls === 2)
expect(state.streamProbes).toBeGreaterThan(0)
})
const queues = useMothershipQueueStore.getState().queues
expect(Object.keys(queues)).toEqual(['chat-a'])
expect(queues['chat-a'][0].content).toBe('belongs to chat-a')
// Reused on the retry so the server deduplicates it.
expect(queues['chat-a'][0].resumeUserMessageId).toBe(state.postBodies[0].userMessageId)
// Must NOT have gone to the cross-surface handoff.
expect(MothershipHandoffStorage.consume('ws-1')).toBeNull()
})
})
@@ -138,31 +138,43 @@ import type {
export interface SendMessageOptions {
/**
* Stream id of a send a cleanup abort withdrew, when this call is the
* recovery of it. The dispatcher probes the id before sending and adopts the
* chat the server already created instead, when there is one.
* Message id of a prior attempt this send retries, set when recovering a send
* an unmount cleanup withdrew. Reusing it lets the server deduplicate the two
* attempts instead of opening a second chat.
*/
recoverStreamId?: string
resumeUserMessageId?: string
}
/**
* `true` when the send owns the transcript (rendered, or handed to reconnect),
* `false` when the caller should restore the queue entry, and the object form
* when a cleanup abort withdrew the send while its request was on the wire —
* `streamId` is what recovery probes before re-sending.
* when an unmount cleanup withdrew it — `userMessageId` is what a retry reuses
* so the server deduplicates the two attempts.
*/
type StartSendMessageResult = boolean | { kind: 'recoverable_cleanup_abort'; streamId: string }
type StartSendMessageResult = boolean | { userMessageId: string }
/**
* Outcome of asking the server whether it accepted the request a cleanup abort
* withdrew. `superseded` is deliberately distinct from `not_found`: the former
* means the answer is unknown because this dispatch is no longer current, and
* must not be treated as licence to re-send.
*/
type RecoveredSendProbe =
| { status: 'adopted'; chatId: string }
| { status: 'not_found' }
| { status: 'superseded' }
interface StartSendMessageOptions {
/** Awaited before dispatch. Defaults to the hook's in-flight stop, if any. */
pendingStop?: Promise<void> | null
/** Runs once the optimistic user/assistant pair is in the transcript. */
onOptimisticSendApplied?: () => void
/** Seed for a queued send that superseded a stopped stream. */
queuedSendHandoff?: QueuedSendHandoffSeed
/**
* Message id of a prior attempt this send retries. Reusing it is what makes
* the retry safe: the server deduplicates against that attempt rather than
* opening a second chat and billing a second turn.
*/
resumeUserMessageId?: string
}
/** A send an unmount cleanup withdrew, as handed to the next chat surface. */
interface WithdrawnSend {
content: string
fileAttachments?: FileAttachmentForApi[]
contexts?: ChatContext[]
userMessageId: string
}
export interface UseChatReturn {
messages: ChatMessage[]
@@ -205,16 +217,6 @@ const RECONNECT_MAX_DELAY_MS = 30_000
const STREAM_BATCH_FETCH_TIMEOUT_MS = 10_000
const STREAM_CHAT_ID_RESOLVE_TIMEOUT_MS = 10_000
const CHAT_HISTORY_RECOVERY_TIMEOUT_MS = 10_000
/**
* How long a recovered send waits to find out whether the request its cleanup
* abort withdrew had in fact been accepted by the server. The server registers
* the stream early in the request (before it responds), so a short poll is
* enough; the ceiling is deliberately low because every millisecond here delays
* a send the user is watching for. Timing out re-sends, which is the safe
* direction: a lost message is worse than a rare duplicate.
*/
const RECOVERED_SEND_PROBE_TIMEOUT_MS = 2500
const RECOVERED_SEND_PROBE_INTERVAL_MS = 250
const STOP_REQUEST_TIMEOUT_MS = 15_000
const QUEUED_SEND_HANDOFF_STORAGE_KEY = `${STREAM_STORAGE_KEY}:queued-send-handoff`
const QUEUED_SEND_HANDOFF_CLAIM_STORAGE_KEY = `${STREAM_STORAGE_KEY}:queued-send-handoff-claim`
@@ -3387,7 +3389,7 @@ export function useChat(
message: string,
fileAttachments?: FileAttachmentForApi[],
contexts?: ChatContext[],
recoverStreamId?: string
resumeUserMessageId?: string
): QueuedMothershipMessage => {
const id = generateId()
const handoffChatId = selectedChatIdRef.current ?? chatIdRef.current
@@ -3407,7 +3409,7 @@ export function useChat(
content: message,
fileAttachments,
contexts,
...(recoverStreamId ? { recoverStreamId } : {}),
...(resumeUserMessageId ? { resumeUserMessageId } : {}),
...(supersededStreamId || handoffChatId
? {
queuedSendHandoff: {
@@ -3494,12 +3496,11 @@ export function useChat(
message: string,
fileAttachments?: FileAttachmentForApi[],
contexts?: ChatContext[],
pendingStopOverride?: Promise<void> | null,
onOptimisticSendApplied?: () => void,
queuedSendHandoff?: QueuedSendHandoffSeed
options?: StartSendMessageOptions
): Promise<StartSendMessageResult> => {
if (!message.trim() || !workspaceId) return false
const pendingStop = pendingStopOverride ?? pendingStopPromiseRef.current
const { onOptimisticSendApplied, queuedSendHandoff } = options ?? {}
const pendingStop = options?.pendingStop ?? pendingStopPromiseRef.current
const pendingStopStreamId = pendingStop
? queuedSendHandoff?.supersededStreamId ||
locallyTerminalStreamIdRef.current ||
@@ -3514,7 +3515,10 @@ export function useChat(
setError(null)
setTransportStreaming()
const userMessageId = queuedSendHandoff?.userMessageId ?? generateId()
/* A retry of a withdrawn send reuses its id so the server deduplicates
the two attempts; anything else mints a fresh one. */
const userMessageId =
queuedSendHandoff?.userMessageId ?? options?.resumeUserMessageId ?? generateId()
const assistantId = getLiveAssistantMessageId(userMessageId)
const storedAttachments: PersistedFileAttachment[] | undefined =
@@ -3809,6 +3813,19 @@ export function useChat(
setError('Previous response is still shutting down; queued message was restored.')
return false
}
/* A send deduplicated against an earlier attempt comes back naming
the chat that attempt opened. Adopting it here spares a chatless
surface the stream-to-chat lookup and puts the user in the right
chat before the reconnect below replays it. */
const conflictChatId =
typeof errorData.chatId === 'string' ? errorData.chatId : undefined
if (conflictChatId && !streamTargetChatId) {
adoptResolvedChatId(conflictChatId, {
replaceHomeHistory: true,
invalidateList: true,
})
streamTargetChatId = conflictChatId
}
streamIdRef.current = conflictStreamId
const succeeded = await retryReconnect({
streamId: conflictStreamId,
@@ -3877,19 +3894,18 @@ export function useChat(
(err instanceof Error && err.name === 'AbortError') || sendAbortSignal?.aborted === true
if (sendWasAborted) {
if (sendAbortSignal?.reason === 'unmount:client_cleanup' && !sendReachedServer) {
/* A remount (StrictMode's dev double-mount, or a real navigation
away) ran the unmount cleanup before this send's response headers
arrived. Nothing was rendered from it, so withdraw the optimistic
pair and report the distinct outcome so the dispatcher recovers
the message.
/* A remount ran the unmount cleanup before this send's response
arrived — a chat-route `key` change, or StrictMode's dev
double-mount. Nothing was rendered from it, so withdraw the
optimistic pair and report the message id, which a retry reuses.
`sendReachedServer` only rules out a send whose RESPONSE landed —
the request itself may well have been accepted, and the route
never reads `request.signal`, so it runs to completion either
way. Recovery therefore carries `userMessageId` as the stream id
to probe before re-sending. */
The request itself may well have been accepted: the route never
reads `request.signal`, so it runs to completion regardless of
the abort. Reusing the id is what makes the retry safe — the
server deduplicates it against that turn instead of billing
another one. */
rollbackOptimisticSend()
return { kind: 'recoverable_cleanup_abort', streamId: userMessageId }
return { userMessageId }
}
return consumedByTranscript
}
@@ -3942,6 +3958,32 @@ export function useChat(
setTransportStreaming,
]
)
/**
* Hands a send the unmount cleanup withdrew to whatever chat surface comes
* next: the live replacement's listener when one is mounted, else a one-shot
* stored handoff for the next mount. Both lanes carry `userMessageId`, so
* whoever picks it up retries as the same send rather than a new one.
*/
const handOffWithdrawnSend = useCallback(
(send: WithdrawnSend) => {
if (
sendMothershipMessage(send.content, send.contexts, send.fileAttachments, send.userMessageId)
) {
return
}
MothershipHandoffStorage.store(
{
message: send.content,
...(send.contexts?.length ? { contexts: send.contexts } : {}),
...(send.fileAttachments?.length ? { fileAttachments: send.fileAttachments } : {}),
resumeUserMessageId: send.userMessageId,
},
workspaceId
)
},
[workspaceId]
)
const sendMessage = useCallback(
async (
message: string,
@@ -3975,33 +4017,45 @@ export function useChat(
queueStore.setEditing(activeChatKey, null)
}
const queued = createQueuedMessage(
message,
// An in-flight send drains the queue from `finalize`; a pending stop kicks
// the dispatcher itself, since nothing else will once the stop settles.
if (sendingRef.current || pendingStopPromiseRef.current) {
queueStore.enqueue(
activeChatKey,
createQueuedMessage(message, fileAttachments, contexts, options?.resumeUserMessageId)
)
if (pendingStopPromiseRef.current) {
void enqueueQueueDispatchRef.current({ type: 'send_head' })
}
return
}
const result = await startSendMessage(message, fileAttachments, contexts, options)
if (typeof result !== 'object') return
/* An unmount cleanup withdrew the send. A chat-bound key is the stable
chat id, so re-queueing under the key this was sent to is the durable
retry — and keeps the message in that chat rather than following the
user into whichever one they opened next. Only a chatless surface,
whose key dies with the mount, goes to the cross-surface lanes. */
const withdrawn = {
content: message,
fileAttachments,
contexts,
options?.recoverStreamId
)
if (sendingRef.current) {
queueStore.enqueue(activeChatKey, queued)
userMessageId: result.userMessageId,
}
if (activeChatKey.startsWith(PENDING_CHAT_KEY_PREFIX)) {
handOffWithdrawnSend(withdrawn)
return
}
if (pendingStopPromiseRef.current) {
queueStore.enqueue(activeChatKey, queued)
void enqueueQueueDispatchRef.current({ type: 'send_head' })
return
}
/* Even an idle-path send goes through the durable queue: a direct
startSendMessage has no backing entry, so a cleanup abort mid-flight
silently drops it. The dispatch loop claims the head in the same tick
for a chatless send; a chat-bound one yields once on `cancelQueries`
first, so it can briefly show as queued. */
queueStore.enqueue(activeChatKey, queued)
void enqueueQueueDispatchRef.current({ type: 'send_head' })
useMothershipQueueStore
.getState()
.enqueue(
activeChatKey,
createQueuedMessage(message, fileAttachments, contexts, result.userMessageId)
)
},
[workspaceId, createQueuedMessage]
[workspaceId, createQueuedMessage, startSendMessage, handOffWithdrawnSend]
)
useEffect(() => {
if (typeof window === 'undefined') return
@@ -4170,19 +4224,15 @@ export function useChat(
const claimOwnerId = writeQueuedSendHandoffClaim(handoff.id)
recoveringQueuedSendHandoffRef.current = { id: handoff.id, ownerId: claimOwnerId }
void startSendMessage(
handoff.message,
handoff.fileAttachments,
handoff.contexts,
null,
undefined,
{
void startSendMessage(handoff.message, handoff.fileAttachments, handoff.contexts, {
pendingStop: null,
queuedSendHandoff: {
id: handoff.id,
chatId: handoff.chatId,
supersededStreamId: handoff.supersededStreamId,
userMessageId: handoff.userMessageId,
}
).finally(() => {
},
}).finally(() => {
if (
recoveringQueuedSendHandoffRef.current?.id === handoff.id &&
recoveringQueuedSendHandoffRef.current.ownerId === claimOwnerId
@@ -4509,49 +4559,6 @@ export function useChat(
]
)
/**
* Answers "did the server accept the request that a cleanup abort withdrew?"
* by polling for the stream it would have registered.
*
* The abort tears down the client's socket but the route handler never reads
* `request.signal`, so an accepted request still creates the chat, persists
* the user message, and runs the turn. Re-sending in that case bills a second
* run and leaves the user with two chats, so recovery adopts the existing
* chat instead whenever this resolves one.
*
* @returns The chat the orphaned stream belongs to, or `undefined` when the
* server has no such stream (never accepted, or already gone) — in which case
* the caller re-sends.
*/
const resolveRecoveredSendChatId = useCallback(
async (streamId: string, epoch: number): Promise<RecoveredSendProbe> => {
const deadline = Date.now() + RECOVERED_SEND_PROBE_TIMEOUT_MS
while (true) {
const resolve = resolveDetachedChatForStreamRef.current
if (!resolve) return { status: 'not_found' }
/* "Superseded" is NOT the same answer as "the server has no such
stream", and the caller must not conflate them: adopting rewrites
the URL, and re-sending after an unmount would open a POST whose
abort controller the teardown has already dropped — a zombie request
nothing can cancel, duplicating the very send this recovery exists
to protect. Reported distinctly so the caller leaves the entry
queued for a later mount instead. */
if (epoch !== queueDispatchEpochRef.current) return { status: 'superseded' }
const resolution = await resolve(streamId)
if (epoch !== queueDispatchEpochRef.current) return { status: 'superseded' }
if (resolution.chatId) return { status: 'adopted', chatId: resolution.chatId }
// A terminal status means the stream existed and finished without a
// durable owner; polling cannot improve on that.
if (resolution.terminal) return { status: 'not_found' }
if (Date.now() + RECOVERED_SEND_PROBE_INTERVAL_MS >= deadline) {
return { status: 'not_found' }
}
await sleep(RECOVERED_SEND_PROBE_INTERVAL_MS)
}
},
[]
)
const dispatchQueuedMessage = useCallback(
async (
msg: QueuedMothershipMessage,
@@ -4586,32 +4593,15 @@ export function useChat(
useMothershipQueueStore.getState().remove(dispatchChatKey, msg.id)
}
/**
* Hands a chatless send to whatever surface comes next, because its
* `pending::` queue key is regenerated per mount and anything left under
* this one is unreachable. Prefers the live replacement surface's
* listener and falls back to a one-shot stored handoff for a real
* navigation away. Both lanes carry attachments and the stream id, so the
* next surface probes before it sends.
*/
const handOffChatlessRecovery = (recoverStreamId?: string) => {
if (
!sendMothershipMessage(msg.content, msg.contexts, msg.fileAttachments, recoverStreamId)
) {
MothershipHandoffStorage.store(
{
message: msg.content,
...(msg.contexts?.length ? { contexts: msg.contexts } : {}),
...(msg.fileAttachments?.length ? { fileAttachments: msg.fileAttachments } : {}),
...(recoverStreamId ? { recoverStreamId } : {}),
},
workspaceId
)
}
}
const restoreQueuedMessage = (handoff?: QueuedSendHandoffSeed, recoverStreamId?: string) => {
const recoverableCleanupAbort = recoverStreamId !== undefined
/* What actually went out. `msg` is the snapshot from when the dispatch was
scheduled; the send below uses the re-read live entry, so recovery
tracks that rather than assuming the two still match. */
let dispatched = msg
const restoreQueuedMessage = (
handoff?: QueuedSendHandoffSeed,
withdrawnUserMessageId?: string
) => {
const withdrawnByCleanup = withdrawnUserMessageId !== undefined
if (!handoff) {
clearQueuedSendHandoffState(msg.id)
}
@@ -4619,7 +4609,7 @@ export function useChat(
if (!removedFromQueue) {
return
}
if (options.epoch !== queueDispatchEpochRef.current && !recoverableCleanupAbort) {
if (options.epoch !== queueDispatchEpochRef.current && !withdrawnByCleanup) {
return
}
// If the user explicitly removed this message during dispatch, honor
@@ -4627,22 +4617,22 @@ export function useChat(
if (userRemovedDuringDispatchRef.current.delete(msg.id)) {
return
}
/* A pending (chatless) surface regenerates its chat key per mount, and
the cleanup that aborted this send belongs to a full remount — a
queue restore would orphan the message under the dead instance's
key. Deliver to the replacement surface instead: its send listener
is live by the time this microtask executes. When nothing claims the
event (a real navigation away), a one-shot handoff covers the next
mount; both lanes carry attachments and the stream id to probe.
Chat-bound sends keep the queue restore — their key is the stable
chat id — and carry the stream id on the restored entry. */
if (recoverableCleanupAbort && dispatchChatKey.startsWith(PENDING_CHAT_KEY_PREFIX)) {
handOffChatlessRecovery(recoverStreamId)
/* A chatless surface regenerates its queue key every mount, so a
restore would strand this under the dead instance's key — hand it to
the next surface instead. A chat-bound key is the stable chat id, so
the queue itself is the durable retry. */
if (withdrawnByCleanup && dispatchChatKey.startsWith(PENDING_CHAT_KEY_PREFIX)) {
handOffWithdrawnSend({
content: dispatched.content,
fileAttachments: dispatched.fileAttachments,
contexts: dispatched.contexts,
userMessageId: withdrawnUserMessageId,
})
return
}
useMothershipQueueStore.getState().insertAt(dispatchChatKey, originalIndex, {
...msg,
...(recoverStreamId ? { recoverStreamId } : {}),
...dispatched,
...(withdrawnUserMessageId ? { resumeUserMessageId: withdrawnUserMessageId } : {}),
})
}
@@ -4660,66 +4650,27 @@ export function useChat(
// Re-read live: the user may have applied an in-place edit (`replaceAt`)
// between dispatch scheduling and this send.
const liveMsg = queueAtSend[currentIndex]
dispatched = liveMsg
activeQueuedSendHandoff = options.queuedSendHandoff ?? liveMsg.queuedSendHandoff
/* This entry is the recovery of a send a cleanup abort withdrew while
its request was already on the wire. The server never sees that
abort, so if it had accepted the request it created the chat and
persisted the message regardless — re-sending would duplicate both
the chat and the billed run. Probe the orphaned stream first and
adopt its chat instead when it exists. */
if (liveMsg.recoverStreamId) {
const probe = await resolveRecoveredSendChatId(liveMsg.recoverStreamId, options.epoch)
/* Unknown, not "safe to send". A chat-bound key is the stable chat
id, so leaving the entry queued IS the retry — the next mount's
drain probes again. A chatless `pending::` key is regenerated per
mount, so the same move would strand the message under a dead key;
hand it to the recovery lanes instead, still carrying the stream
id. Skipped when the entry is no longer under this key (adoption
migrated it to a live chat), where it is already recoverable. */
if (probe.status === 'superseded') {
if (dispatchChatKey.startsWith(PENDING_CHAT_KEY_PREFIX)) {
const queueStore = useMothershipQueueStore.getState()
const stillUnderDeadKey = (queueStore.queues[dispatchChatKey] ?? []).some(
(queued) => queued.id === msg.id
)
if (stillUnderDeadKey) {
queueStore.remove(dispatchChatKey, msg.id)
handOffChatlessRecovery(liveMsg.recoverStreamId)
}
}
return
}
if (probe.status === 'adopted') {
removeQueuedMessage()
adoptResolvedChatId(probe.chatId, {
replaceHomeHistory: true,
invalidateList: true,
})
/* Adoption alone does not surface the running turn. Hydration only
reconnects when `chatHistory.activeStreamId` is set, and that
query is cached for `MOTHERSHIP_CHAT_HISTORY_STALE_TIME` — on a
chat-bound recover the client usually holds a copy predating this
stream, so without an explicit detail invalidation the adopted
chat renders with the live response invisible. */
invalidateChatQueries({ includeDetail: true, targetChatId: probe.chatId })
return
}
}
const sendResult = await startSendMessage(
liveMsg.content,
liveMsg.fileAttachments,
liveMsg.contexts,
options.pendingStop,
removeQueuedMessage,
activeQueuedSendHandoff
{
pendingStop: options.pendingStop,
onOptimisticSendApplied: removeQueuedMessage,
queuedSendHandoff: activeQueuedSendHandoff,
...(liveMsg.resumeUserMessageId
? { resumeUserMessageId: liveMsg.resumeUserMessageId }
: {}),
}
)
if (sendResult !== true) {
restoreQueuedMessage(
activeQueuedSendHandoff,
typeof sendResult === 'object' ? sendResult.streamId : undefined
typeof sendResult === 'object' ? sendResult.userMessageId : undefined
)
}
} catch {
@@ -4730,13 +4681,7 @@ export function useChat(
userRemovedDuringDispatchRef.current.delete(msg.id)
}
},
[
startSendMessage,
workspaceId,
resolveRecoveredSendChatId,
adoptResolvedChatId,
invalidateChatQueries,
]
[startSendMessage, handOffWithdrawnSend]
)
const runQueueDispatchLoop = useCallback(async () => {
@@ -490,7 +490,7 @@ export const Panel = memo(function Panel() {
e.preventDefault()
setActiveTab('copilot')
copilotSendMessage(detail.message, detail.fileAttachments, detail.contexts, {
...(detail.recoverStreamId ? { recoverStreamId: detail.recoverStreamId } : {}),
...(detail.resumeUserMessageId ? { resumeUserMessageId: detail.resumeUserMessageId } : {}),
})
}
window.addEventListener(MOTHERSHIP_SEND_MESSAGE_EVENT, handler)
+172
View File
@@ -36,6 +36,9 @@ const {
appendCopilotChatMessages,
persistChatResources,
mockPublishStatusChanged,
atomicallyClaimChatSend,
storeChatSendResult,
releaseChatSendClaim,
} = vi.hoisted(() => ({
generateWorkspaceSnapshot: vi.fn(),
processContextsServer: vi.fn(),
@@ -51,6 +54,9 @@ const {
appendCopilotChatMessages: vi.fn(),
persistChatResources: vi.fn(),
mockPublishStatusChanged: vi.fn(),
atomicallyClaimChatSend: vi.fn(),
storeChatSendResult: vi.fn(),
releaseChatSendClaim: vi.fn(),
}))
const getSession = authMockFns.mockGetSession
@@ -103,6 +109,14 @@ vi.mock('@/lib/copilot/chat/lifecycle', () => ({
resolveOrCreateChat,
}))
vi.mock('@/lib/core/idempotency', () => ({
chatSendIdempotency: {
atomicallyClaim: atomicallyClaimChatSend,
storeResult: storeChatSendResult,
release: releaseChatSendClaim,
},
}))
vi.mock('@/lib/copilot/chat/terminal-state', () => ({
finalizeAssistantTurn,
}))
@@ -132,6 +146,14 @@ describe('handleUnifiedChatPost', () => {
beforeEach(() => {
vi.clearAllMocks()
resetDbChainMock()
atomicallyClaimChatSend.mockResolvedValue({
claimed: true,
normalizedKey: 'chat-send:user-message:msg-1:userId=user-1',
storageMethod: 'database',
claimToken: 'claim-1',
})
storeChatSendResult.mockResolvedValue(true)
releaseChatSendClaim.mockResolvedValue(undefined)
getSession.mockResolvedValue({ user: { id: 'user-1' } })
resolveWorkflowIdForUser.mockResolvedValue({
status: 'resolved',
@@ -721,8 +743,158 @@ describe('handleUnifiedChatPost', () => {
)
expect(response.status).toBe(400)
// Returns without throwing, so only a `finally` can free the claim.
expect(releaseChatSendClaim).toHaveBeenCalled()
await expect(response.json()).resolves.toMatchObject({
error: 'workspaceId is required when workflowId is not provided',
})
})
describe('deduplicating a repeated send', () => {
/**
* The client cannot tell whether a request it aborted reached the server —
* the route never reads `request.signal`, so an accepted one runs to
* completion regardless. Recovering such a send therefore retries it under
* the original `userMessageId`, and this is what makes that safe.
*/
it('answers an already-claimed send with the chat the first attempt opened', async () => {
atomicallyClaimChatSend.mockResolvedValue({
claimed: false,
normalizedKey: 'chat-send:user-message:msg-1:userId=user-1',
storageMethod: 'database',
existingResult: { success: true, status: 'completed', result: { chatId: 'chat-first' } },
})
const response = await handleUnifiedChatPost(
new NextRequest('http://localhost/api/mothership/chat', {
method: 'POST',
body: JSON.stringify({
message: 'Hello',
workspaceId: 'ws-1',
userMessageId: 'msg-1',
createNewChat: true,
}),
})
)
expect(response.status).toBe(409)
await expect(response.json()).resolves.toMatchObject({
activeStreamId: 'msg-1',
chatId: 'chat-first',
})
// The whole point: no second chat, no second billed turn.
expect(resolveOrCreateChat).not.toHaveBeenCalled()
expect(createSSEStream).not.toHaveBeenCalled()
})
it('scopes the claim to the caller so one user cannot probe another', async () => {
await handleUnifiedChatPost(
new NextRequest('http://localhost/api/mothership/chat', {
method: 'POST',
body: JSON.stringify({
message: 'Hello',
workspaceId: 'ws-1',
userMessageId: 'msg-1',
createNewChat: true,
}),
})
)
expect(atomicallyClaimChatSend).toHaveBeenCalledWith('user-message', 'msg-1', {
userId: 'user-1',
})
})
it('records the chat against the send so a retry resolves to it', async () => {
await handleUnifiedChatPost(
new NextRequest('http://localhost/api/mothership/chat', {
method: 'POST',
body: JSON.stringify({
message: 'Hello',
workspaceId: 'ws-1',
userMessageId: 'msg-1',
createNewChat: true,
}),
})
)
expect(storeChatSendResult).toHaveBeenCalledWith(
'chat-send:user-message:msg-1:userId=user-1',
expect.objectContaining({ result: { chatId: 'chat-1' } }),
'database',
'claim-1'
)
})
/**
* Deduplication saves a duplicate chat; the send IS the user's message.
* An unreachable bookkeeping store must degrade chat, never take it down.
*/
it('sends normally when the claim store is unavailable', async () => {
atomicallyClaimChatSend.mockRejectedValue(new Error('idempotency store down'))
const response = await handleUnifiedChatPost(
new NextRequest('http://localhost/api/mothership/chat', {
method: 'POST',
body: JSON.stringify({
message: 'Hello',
workspaceId: 'ws-1',
userMessageId: 'msg-1',
createNewChat: true,
}),
})
)
expect(response.status).toBe(200)
expect(createSSEStream).toHaveBeenCalled()
expect(storeChatSendResult).not.toHaveBeenCalled()
})
/**
* The queued-send-handoff path deliberately retries under the original
* `userMessageId` after a stream collision. If the collided attempt left a
* permanent claim, that retry would deduplicate against a chat whose turn
* never started and reattach to a stream that does not exist.
*/
it('releases the claim when a stream collision stops the turn from starting', async () => {
acquirePendingChatStream.mockResolvedValue(false)
getPendingChatStreamId.mockResolvedValue('other-stream')
const response = await handleUnifiedChatPost(
new NextRequest('http://localhost/api/mothership/chat', {
method: 'POST',
body: JSON.stringify({
message: 'Hello',
workspaceId: 'ws-1',
userMessageId: 'msg-1',
createNewChat: true,
}),
})
)
expect(response.status).toBe(409)
expect(releaseChatSendClaim).toHaveBeenCalledWith(
'chat-send:user-message:msg-1:userId=user-1',
'database',
'claim-1'
)
})
it('keeps the claim once a turn is actually streaming', async () => {
const response = await handleUnifiedChatPost(
new NextRequest('http://localhost/api/mothership/chat', {
method: 'POST',
body: JSON.stringify({
message: 'Hello',
workspaceId: 'ws-1',
userMessageId: 'msg-1',
createNewChat: true,
}),
})
)
expect(response.status).toBe(200)
expect(releaseChatSendClaim).not.toHaveBeenCalled()
})
})
})
+111 -1
View File
@@ -58,6 +58,8 @@ import {
sanitizeChatResources,
} from '@/lib/copilot/resources/types'
import { prepareExecutionContext } from '@/lib/copilot/tools/handlers/context'
import type { AtomicClaimResult } from '@/lib/core/idempotency'
import { chatSendIdempotency } from '@/lib/core/idempotency'
import { captureServerEvent } from '@/lib/posthog/server'
import { resolveWorkflowIdForUser } from '@/lib/workflows/utils'
import {
@@ -258,7 +260,10 @@ const ChatContextSchema = z
const ChatMessageSchema = z.object({
message: z.string().min(1, 'Message is required'),
userMessageId: z.string().optional(),
/* Bounded because it becomes part of a Postgres key in `chatSendIdempotency`;
a client-supplied id longer than the btree entry limit would throw there.
A generated id is 36 chars. */
userMessageId: z.string().max(128).optional(),
chatId: z.string().optional(),
workflowId: z.string().optional(),
workspaceId: z.string().optional(),
@@ -943,10 +948,68 @@ async function resolveBranch(params: {
}
}
/** Names what the key identifies: `chat-send:user-message:<id>:userId=<id>`. */
const CHAT_SEND_IDEMPOTENCY_PROVIDER = 'user-message'
/**
* Claims this send so a retry of it can be recognised.
*
* Fails open: a missed deduplication costs a duplicate chat and turn, but
* refusing the send loses the user's message. Returns `undefined` when the
* store is unreachable, which sends normally with no claim to finalize.
*
* The key is scoped to the caller — `userMessageId` is client-supplied, so an
* unscoped one would let a user probe another's sends for their chat id.
*/
async function claimChatSend(
userMessageId: string,
userId: string
): Promise<AtomicClaimResult | undefined> {
try {
return await chatSendIdempotency.atomicallyClaim(
CHAT_SEND_IDEMPOTENCY_PROVIDER,
userMessageId,
{ userId }
)
} catch (error) {
logger.warn('Could not claim chat send; proceeding without deduplication', {
userMessageId,
error: getErrorMessage(error, 'Unknown error'),
})
return undefined
}
}
/**
* Answers a send whose `userMessageId` was already claimed.
*
* Deliberately the same 409 shape the pending-stream lock returns, because the
* client's conflict handler already knows how to reattach to `activeStreamId`
* instead of starting a turn — a duplicate send and a send that collided with
* an in-flight one want exactly the same thing. `chatId` rides along when the
* first attempt got far enough to resolve one, letting a chatless client adopt
* it without a stream-to-chat lookup.
*/
function duplicateChatSendResponse(claim: AtomicClaimResult, userMessageId: string): NextResponse {
const claimed = claim.existingResult?.result?.chatId
const chatId = typeof claimed === 'string' && claimed ? claimed : undefined
logger.info('Deduplicated a repeated chat send', { userMessageId, chatId })
return NextResponse.json(
{
error: 'This message was already sent.',
activeStreamId: userMessageId,
...(chatId ? { chatId } : {}),
},
{ status: 409 }
)
}
export async function handleUnifiedChatPost(req: NextRequest) {
let actualChatId: string | undefined
let userMessageId = ''
let chatStreamLockAcquired = false
/** Cleared once the chat is recorded against it, which makes it permanent. */
let sendClaim: AtomicClaimResult | undefined
// Started once we've parsed the body (need userMessageId to stamp as
// streamId). Every subsequent span (persistUserMessage,
// createRunSegment, the whole SSE stream, etc.) nests under this
@@ -981,6 +1044,11 @@ export async function handleUnifiedChatPost(req: NextRequest) {
const normalizedContexts = normalizeContexts(body.contexts) ?? []
userMessageId = body.userMessageId || generateId()
sendClaim = await claimChatSend(userMessageId, authenticatedUserId)
if (sendClaim?.claimed === false) {
return duplicateChatSendResponse(sendClaim, userMessageId)
}
otelRoot = startCopilotOtelRoot({
streamId: userMessageId,
executionId,
@@ -1073,6 +1141,27 @@ export async function handleUnifiedChatPost(req: NextRequest) {
}
}
/* Record the chat as soon as it is known — the earliest a retry can be
answered with somewhere to go. This does not make the claim permanent:
several exits below still return without starting a turn, and a retry
of those must be free to start one. Failing to record only costs a
retry the chat-id shortcut, so it must not fail the send. */
if (sendClaim?.claimToken && actualChatId) {
await chatSendIdempotency
.storeResult(
sendClaim.normalizedKey,
{ success: true, status: 'completed', result: { chatId: actualChatId } },
sendClaim.storageMethod,
sendClaim.claimToken
)
.catch((error) => {
logger.warn(`[${requestId}] Could not record the chat for this send`, {
userMessageId,
error: getErrorMessage(error, 'Unknown error'),
})
})
}
if (chatIsNew && actualChatId && body.resourceAttachments?.length) {
// Canonicalizes here, not just inside `persistChatResources`: several
// browser tabs collapse onto the one Browser panel before they are
@@ -1366,6 +1455,11 @@ export async function handleUnifiedChatPost(req: NextRequest) {
const rootTraceparent = `00-${rootCtx.traceId}-${rootCtx.spanId}-${
(rootCtx.traceFlags & 0x1) === 0x1 ? '01' : '00'
}`
/* A turn is running. Only now is the claim permanent, so the `finally`
below leaves it in place and a retry of this send resolves to this
chat instead of opening another. Every earlier exit returns without a
turn, and releases. */
sendClaim = undefined
return new Response(stream, {
headers: {
...SSE_RESPONSE_HEADERS,
@@ -1406,5 +1500,21 @@ export async function handleUnifiedChatPost(req: NextRequest) {
},
{ status: 500 }
)
} finally {
/* A claim still held here never started a turn — the send threw, or
returned early on a rejected branch, a missing chat, or a chat that
already has a stream running. Release it so a retry may start one rather
than deduplicating against a turn that never happened. Must be
`finally`: those early returns skip `catch`. */
if (sendClaim?.claimToken) {
await chatSendIdempotency
.release(sendClaim.normalizedKey, sendClaim.storageMethod, sendClaim.claimToken)
.catch((releaseError) => {
logger.warn('Could not release the claim for an unfinished send', {
userMessageId,
error: getErrorMessage(releaseError, 'Unknown error'),
})
})
}
}
}
+25
View File
@@ -733,3 +733,28 @@ export const billingIdempotency = new IdempotencyService({
ttlSeconds: 60 * 60, // 1 hour
forceStorage: 'database',
})
/**
* Dedupes a chat send by its client-generated `userMessageId`, so re-sending
* one is safe.
*
* The client cannot tell whether a request it aborted reached the server: the
* chat route never reads `request.signal`, so an accepted request creates the
* chat, persists the user message, and runs the (billed) turn even after the
* browser drops the socket. Without this, a client that recovers an aborted
* send has to choose between losing the message and duplicating the run.
*
* Storage is forced to Postgres for the same reason Stripe webhook claims are
* (see `stripeWebhookIdempotency`): a missed deduplication is a second LLM turn
* billed to the workspace, so the key must not be evictable under Redis memory
* pressure. The added 1-5ms is invisible next to the LLM call that follows.
*
* `inProgressTtlSeconds` is short so a crashed pod cannot block a genuine retry
* for the full hour, while completed sends stay deduplicated for it.
*/
export const chatSendIdempotency = new IdempotencyService({
namespace: 'chat-send',
ttlSeconds: 60 * 60, // 1 hour
inProgressTtlSeconds: 60,
forceStorage: 'database',
})
+8 -8
View File
@@ -311,12 +311,12 @@ export interface MothershipHandoff {
/** Already-uploaded attachment references riding along with the message. */
fileAttachments?: FileAttachmentForApi[]
/**
* Set only when a cleanup abort withdrew an in-flight send and this handoff
* is the recovery of it: the aborted send's stream id. The consuming chat
* probes it before sending, so a request the server had already accepted is
* adopted rather than sent a second time.
* Set only when this handoff is the recovery of a send an unmount cleanup
* withdrew: that attempt's message id. The consuming chat reuses it so the
* server deduplicates against the first attempt instead of opening a second
* chat and billing a second turn.
*/
recoverStreamId?: string
resumeUserMessageId?: string
}
interface StoredHandoff extends MothershipHandoff {
@@ -364,7 +364,7 @@ export class MothershipHandoffStorage {
? contexts
: [...MothershipHandoffStorage.pendingContexts(workspaceId), ...contexts],
...(handoff.fileAttachments?.length ? { fileAttachments: handoff.fileAttachments } : {}),
...(handoff.recoverStreamId ? { recoverStreamId: handoff.recoverStreamId } : {}),
...(handoff.resumeUserMessageId ? { resumeUserMessageId: handoff.resumeUserMessageId } : {}),
workspaceId,
timestamp: Date.now(),
})
@@ -427,8 +427,8 @@ export class MothershipHandoffStorage {
...(Array.isArray(data.fileAttachments) && data.fileAttachments.length > 0
? { fileAttachments: data.fileAttachments }
: {}),
...(typeof data.recoverStreamId === 'string' && data.recoverStreamId
? { recoverStreamId: data.recoverStreamId }
...(typeof data.resumeUserMessageId === 'string' && data.resumeUserMessageId
? { resumeUserMessageId: data.resumeUserMessageId }
: {}),
}
}
+7 -7
View File
@@ -28,12 +28,12 @@ export interface MothershipSendMessageDetail {
/** Already-uploaded attachments riding along with the message. */
fileAttachments?: FileAttachmentForApi[]
/**
* Set only when a cleanup abort withdrew an in-flight send and this is the
* recovery of it: the aborted send's stream id. The receiving chat probes it
* before sending, so a request the server had already accepted is adopted
* rather than sent a second time.
* Set only when this message is the recovery of a send an unmount cleanup
* withdrew: that attempt's message id. The receiving chat reuses it so the
* server deduplicates against the first attempt instead of opening a second
* chat and billing a second turn.
*/
recoverStreamId?: string
resumeUserMessageId?: string
}
/**
@@ -49,7 +49,7 @@ export function sendMothershipMessage(
message: string,
contexts?: ChatContext[],
fileAttachments?: FileAttachmentForApi[],
recoverStreamId?: string
resumeUserMessageId?: string
): boolean {
const trimmed = message.trim()
if (!trimmed) {
@@ -60,7 +60,7 @@ export function sendMothershipMessage(
message: trimmed,
contexts,
fileAttachments,
...(recoverStreamId ? { recoverStreamId } : {}),
...(resumeUserMessageId ? { resumeUserMessageId } : {}),
})
logger.info('Dispatched mothership message event', { messageLength: trimmed.length, consumed })
return consumed
+4 -4
View File
@@ -99,12 +99,12 @@ export const useMothershipQueueStore = create<MothershipQueueState>()(
const next = [...current]
// Strip `queuedSendHandoff` — references the stream active at
// original enqueue time; the dispatcher mints a fresh one at send.
// Strip `recoverStreamId` too: it dedupes against a server-side copy
// of the PRE-edit text, which the edited message is no longer a
// duplicate of, so probing it would wrongly suppress this send.
// Strip `resumeUserMessageId` too: it deduplicates against a
// server-side copy of the PRE-edit text, which the edited message is
// not a duplicate of, so reusing it would suppress this send.
const {
queuedSendHandoff: _stale,
recoverStreamId: _staleRecover,
resumeUserMessageId: _staleResume,
...rest
} = next[index]
next[index] = {
+6 -7
View File
@@ -11,14 +11,13 @@ export interface QueuedSendHandoffSeed {
export type QueuedMothershipMessage = QueuedMessage & {
queuedSendHandoff?: QueuedSendHandoffSeed
/**
* Stream id (the aborted send's `userMessageId`) of a send a cleanup abort
* withdrew while its request was already on the wire. The dispatcher probes
* it before re-sending: the server does not observe `request.signal`, so a
* request it had already accepted still creates the chat and persists the
* message, and re-sending blind would duplicate both. Persisted, so a
* nav-back restore probes too.
* Message id of a prior attempt at this send that an unmount cleanup
* withdrew. Reused when the entry is dispatched so the server deduplicates
* against that attempt it never sees the client's abort, so a request it
* had already accepted still opened the chat and billed the turn. Persisted,
* so a retry after a reload deduplicates too.
*/
recoverStreamId?: string
resumeUserMessageId?: string
}
// Mutable fields an in-place edit overwrites; id and index are preserved by `replaceAt`.