Skip to content

Commit 3f399e7

Browse files
authoredJan 31, 2025··
fix(start): fix stream handling during SSR (#3299)
1 parent 3f1f24f commit 3f399e7

File tree

4 files changed

+124
-164
lines changed

4 files changed

+124
-164
lines changed
 

‎e2e/start/basic/app/routes/stream.tsx

+5-2
Original file line numberDiff line numberDiff line change
@@ -33,8 +33,11 @@ function Home() {
3333
let chunk
3434

3535
while (!(chunk = await reader.read()).done) {
36-
const decoded = decoder.decode(chunk.value, { stream: !chunk.done })
37-
setStreamData((prev) => [...prev, decoded])
36+
let value = chunk.value
37+
if (typeof value !== 'string') {
38+
value = decoder.decode(value, { stream: !chunk.done })
39+
}
40+
setStreamData((prev) => [...prev, value])
3841
}
3942
}
4043

‎packages/start-client/src/ssr-client.tsx

+6-6
Original file line numberDiff line numberDiff line change
@@ -31,19 +31,19 @@ export interface StartSsrGlobal {
3131
initMatch: (match: SsrMatch) => void
3232
resolvePromise: (opts: {
3333
matchId: string
34-
id: string
34+
id: number
3535
promiseState: DeferredPromiseState<any>
3636
}) => void
37-
injectChunk: (opts: { matchId: string; id: string; chunk: string }) => void
38-
closeStream: (opts: { matchId: string; id: string }) => void
37+
injectChunk: (opts: { matchId: string; id: number; chunk: string }) => void
38+
closeStream: (opts: { matchId: string; id: number }) => void
3939
}
4040

4141
export interface SsrMatch {
4242
id: string
43-
__beforeLoadContext?: string
43+
__beforeLoadContext: string
4444
loaderData?: string
4545
error?: string
46-
extracted: Record<string, ClientExtractedEntry>
46+
extracted?: Array<ClientExtractedEntry>
4747
updatedAt: MakeRouteMatch['updatedAt']
4848
status: MakeRouteMatch['status']
4949
}
@@ -160,7 +160,7 @@ export function hydrate(router: AnyRouter) {
160160
}
161161

162162
// Handle extracted
163-
Object.entries((match as any).extracted).forEach(([_, ex]: any) => {
163+
;(match as unknown as SsrMatch).extracted?.forEach((ex) => {
164164
deepMutableSetByPath(match, ['loaderData', ...ex.path], ex.value)
165165
})
166166
} else {

‎packages/start-server/src/ssr-server.ts

+109-152
Original file line numberDiff line numberDiff line change
@@ -27,7 +27,6 @@ export type ServerExtractedEntry =
2727
| ServerExtractedPromise
2828

2929
export interface ServerExtractedBaseEntry extends ClientExtractedBaseEntry {
30-
dataType: '__beforeLoadContext' | 'loaderData'
3130
id: number
3231
matchIndex: number
3332
}
@@ -121,32 +120,21 @@ export function dehydrateRouter(router: AnyRouter) {
121120
)
122121
}
123122

124-
export function extractAsyncDataToMatch(
125-
dataType: '__beforeLoadContext' | 'loaderData',
126-
data: any,
123+
export function extractAsyncLoaderData(
124+
loaderData: any,
127125
ctx: {
128126
match: AnyRouteMatch
129127
router: AnyRouter
130128
},
131129
) {
132-
;(ctx.match as any).extracted = (ctx.match as any).extracted || []
133-
134-
const extracted = (ctx.match as any).extracted
135-
136-
const replacedLoaderData = replaceBy(data, (value, path) => {
137-
const type =
138-
value instanceof ReadableStream
139-
? 'stream'
140-
: value instanceof Promise
141-
? 'promise'
142-
: undefined
130+
const extracted: Array<ServerExtractedEntry> = []
143131

132+
const replaced = replaceBy(loaderData, (value, path) => {
144133
// If it's a stream, we need to tee it so we can read it multiple times
145-
if (type === 'stream') {
134+
if (value instanceof ReadableStream) {
146135
const [copy1, copy2] = value.tee()
147136
const entry: ServerExtractedStream = {
148-
dataType,
149-
type,
137+
type: 'stream',
150138
path,
151139
id: extracted.length,
152140
matchIndex: ctx.match.index,
@@ -155,11 +143,10 @@ export function extractAsyncDataToMatch(
155143

156144
extracted.push(entry)
157145
return copy1
158-
} else if (type === 'promise') {
146+
} else if (value instanceof Promise) {
159147
const deferredPromise = defer(value)
160148
const entry: ServerExtractedPromise = {
161-
dataType,
162-
type,
149+
type: 'promise',
163150
path,
164151
id: extracted.length,
165152
matchIndex: ctx.match.index,
@@ -171,7 +158,7 @@ export function extractAsyncDataToMatch(
171158
return value
172159
})
173160

174-
return replacedLoaderData
161+
return { replaced, extracted }
175162
}
176163

177164
export function onMatchSettled(opts: {
@@ -180,148 +167,118 @@ export function onMatchSettled(opts: {
180167
}) {
181168
const { router, match } = opts
182169

183-
const [serializedBeforeLoadData, serializedLoaderData] = (
184-
['__beforeLoadContext', 'loaderData'] as const
185-
).map((dataType) => {
186-
const data = extractAsyncDataToMatch(dataType, match[dataType], {
187-
router: router,
170+
let extracted: Array<ServerExtractedEntry> | undefined = undefined
171+
let serializedLoaderData: any = undefined
172+
if (match.loaderData !== undefined) {
173+
const result = extractAsyncLoaderData(match.loaderData, {
174+
router,
188175
match,
189176
})
177+
match.loaderData = result.replaced
178+
extracted = result.extracted
179+
serializedLoaderData = extracted.reduce(
180+
(acc: any, entry: ServerExtractedEntry) => {
181+
return deepImmutableSetByPath(acc, ['temp', ...entry.path], undefined)
182+
},
183+
{ temp: result.replaced },
184+
).temp
185+
}
190186

191-
const extracted = (match as any).extracted as
192-
| undefined
193-
| Array<ServerExtractedEntry>
194-
195-
return extracted
196-
? extracted.reduce(
197-
(acc: any, entry: ServerExtractedEntry) => {
198-
if (entry.dataType !== dataType) {
199-
return deepImmutableSetByPath(
200-
acc,
201-
['temp', ...entry.path],
202-
undefined,
203-
)
204-
}
205-
return acc
206-
},
207-
{ temp: data },
208-
).temp
209-
: data
210-
})
187+
const initCode = `__TSR_SSR__.initMatch(${jsesc(
188+
{
189+
id: match.id,
190+
__beforeLoadContext: router.ssr!.serializer.stringify(
191+
match.__beforeLoadContext,
192+
),
193+
loaderData: router.ssr!.serializer.stringify(serializedLoaderData),
194+
error: router.ssr!.serializer.stringify(match.error),
195+
extracted: extracted?.map((entry) => pick(entry, ['type', 'path'])),
196+
updatedAt: match.updatedAt,
197+
status: match.status,
198+
} satisfies SsrMatch,
199+
{
200+
isScriptContext: true,
201+
wrap: true,
202+
json: true,
203+
},
204+
)})`
211205

212-
const extracted = (match as any).extracted as
213-
| undefined
214-
| Array<ServerExtractedEntry>
215-
216-
if (
217-
serializedBeforeLoadData !== undefined ||
218-
serializedLoaderData !== undefined ||
219-
extracted?.length
220-
) {
221-
const initCode = `__TSR_SSR__.initMatch(${jsesc(
222-
{
223-
id: match.id,
224-
__beforeLoadContext: router.ssr!.serializer.stringify(
225-
serializedBeforeLoadData,
226-
),
227-
loaderData: router.ssr!.serializer.stringify(serializedLoaderData),
228-
error: router.ssr!.serializer.stringify(match.error),
229-
extracted: extracted
230-
? Object.fromEntries(
231-
extracted.map((entry) => {
232-
return [entry.id, pick(entry, ['type', 'path'])]
233-
}),
234-
)
235-
: {},
236-
updatedAt: match.updatedAt,
237-
status: match.status,
238-
} satisfies SsrMatch,
239-
{
240-
isScriptContext: true,
241-
wrap: true,
242-
json: true,
243-
},
244-
)})`
206+
router.serverSsr!.injectScript(() => initCode)
245207

246-
router.serverSsr!.injectScript(() => initCode)
208+
if (extracted) {
209+
extracted.forEach((entry) => {
210+
if (entry.type === 'promise') return injectPromise(entry)
211+
return injectStream(entry)
212+
})
213+
}
247214

248-
if (extracted) {
249-
extracted.forEach((entry) => {
250-
if (entry.type === 'promise') return injectPromise(entry)
251-
return injectStream(entry)
252-
})
253-
}
215+
function injectPromise(entry: ServerExtractedPromise) {
216+
router.serverSsr!.injectScript(async () => {
217+
await entry.promise
254218

255-
function injectPromise(entry: ServerExtractedPromise) {
256-
router.serverSsr!.injectScript(async () => {
257-
await entry.promise
258-
259-
return `__TSR_SSR__.resolvePromise(${jsesc(
260-
{
261-
matchId: match.id,
262-
id: entry.id,
263-
promiseState: entry.promise[TSR_DEFERRED_PROMISE],
264-
} satisfies ResolvePromiseState,
265-
{
266-
isScriptContext: true,
267-
wrap: true,
268-
json: true,
269-
},
270-
)})`
271-
})
272-
}
219+
return `__TSR_SSR__.resolvePromise(${jsesc(
220+
{
221+
matchId: match.id,
222+
id: entry.id,
223+
promiseState: entry.promise[TSR_DEFERRED_PROMISE],
224+
} satisfies ResolvePromiseState,
225+
{
226+
isScriptContext: true,
227+
wrap: true,
228+
json: true,
229+
},
230+
)})`
231+
})
232+
}
273233

274-
function injectStream(entry: ServerExtractedStream) {
275-
// Inject a promise that resolves when the stream is done
276-
// We do this to keep the stream open until we're done
277-
router.serverSsr!.injectHtml(async () => {
278-
//
279-
try {
280-
const reader = entry.stream.getReader()
281-
let chunk: ReadableStreamReadResult<any> | null = null
282-
while (!(chunk = await reader.read()).done) {
283-
if (chunk.value) {
284-
const code = `__TSR_SSR__.injectChunk(${jsesc(
285-
{
286-
matchId: match.id,
287-
id: entry.id,
288-
chunk: chunk.value,
289-
},
290-
{
291-
isScriptContext: true,
292-
wrap: true,
293-
json: true,
294-
},
295-
)})`
296-
297-
router.serverSsr!.injectScript(() => code)
298-
}
234+
function injectStream(entry: ServerExtractedStream) {
235+
// Inject a promise that resolves when the stream is done
236+
// We do this to keep the stream open until we're done
237+
router.serverSsr!.injectHtml(async () => {
238+
//
239+
try {
240+
const reader = entry.stream.getReader()
241+
let chunk: ReadableStreamReadResult<any> | null = null
242+
while (!(chunk = await reader.read()).done) {
243+
if (chunk.value) {
244+
const code = `__TSR_SSR__.injectChunk(${jsesc(
245+
{
246+
matchId: match.id,
247+
id: entry.id,
248+
chunk: chunk.value,
249+
},
250+
{
251+
isScriptContext: true,
252+
wrap: true,
253+
json: true,
254+
},
255+
)})`
256+
257+
router.serverSsr!.injectScript(() => code)
299258
}
300-
301-
router.serverSsr!.injectScript(
302-
() =>
303-
`__TSR_SSR__.closeStream(${jsesc(
304-
{
305-
matchId: match.id,
306-
id: entry.id,
307-
},
308-
{
309-
isScriptContext: true,
310-
wrap: true,
311-
json: true,
312-
},
313-
)})`,
314-
)
315-
} catch (err) {
316-
console.error('stream read error', err)
317259
}
318260

319-
return ''
320-
})
321-
}
322-
}
261+
router.serverSsr!.injectScript(
262+
() =>
263+
`__TSR_SSR__.closeStream(${jsesc(
264+
{
265+
matchId: match.id,
266+
id: entry.id,
267+
},
268+
{
269+
isScriptContext: true,
270+
wrap: true,
271+
json: true,
272+
},
273+
)})`,
274+
)
275+
} catch (err) {
276+
console.error('stream read error', err)
277+
}
323278

324-
return null
279+
return ''
280+
})
281+
}
325282
}
326283

327284
function deepImmutableSetByPath<T>(obj: T, path: Array<string>, value: any): T {

‎packages/start-server/src/tsrScript.ts

+4-4
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,7 @@ const __TSR_SSR__: StartSsrGlobal = {
77
initMatch: (match) => {
88
__TSR_SSR__.matches.push(match)
99

10-
Object.entries(match.extracted).forEach(([_id, ex]) => {
10+
match.extracted?.forEach((ex) => {
1111
if (ex.type === 'stream') {
1212
let controller
1313
ex.value = new ReadableStream({
@@ -45,7 +45,7 @@ const __TSR_SSR__: StartSsrGlobal = {
4545
resolvePromise: ({ matchId, id, promiseState }) => {
4646
const match = __TSR_SSR__.matches.find((m) => m.id === matchId)
4747
if (match) {
48-
const ex = match.extracted[id]
48+
const ex = match.extracted?.[id]
4949
if (
5050
ex &&
5151
ex.type === 'promise' &&
@@ -62,7 +62,7 @@ const __TSR_SSR__: StartSsrGlobal = {
6262
const match = __TSR_SSR__.matches.find((m) => m.id === matchId)
6363

6464
if (match) {
65-
const ex = match.extracted[id]
65+
const ex = match.extracted?.[id]
6666
if (ex && ex.type === 'stream' && ex.value?.controller) {
6767
ex.value.controller.enqueue(new TextEncoder().encode(chunk.toString()))
6868
return true
@@ -73,7 +73,7 @@ const __TSR_SSR__: StartSsrGlobal = {
7373
closeStream: ({ matchId, id }) => {
7474
const match = __TSR_SSR__.matches.find((m) => m.id === matchId)
7575
if (match) {
76-
const ex = match.extracted[id]
76+
const ex = match.extracted?.[id]
7777
if (ex && ex.type === 'stream' && ex.value?.controller) {
7878
ex.value.controller.close()
7979
return true

0 commit comments

Comments
 (0)
Please sign in to comment.