diff --git a/apps/content/blog/nextjs-ssr-orpc-tanstack-query.mdx b/apps/content/blog/nextjs-ssr-orpc-tanstack-query.mdx index ea8f6bbd0..9f02da2d0 100644 --- a/apps/content/blog/nextjs-ssr-orpc-tanstack-query.mdx +++ b/apps/content/blog/nextjs-ssr-orpc-tanstack-query.mdx @@ -415,10 +415,10 @@ Wire it into the page, inside its own suspense boundary: ``` -Think about what SSR means for a query that never completes. Two problems: +A query that never completes causes two SSR problems: -1. **On the server**, waiting for the stream to end would block rendering forever. -2. **In the browser**, the hydrated reading is a static snapshot, and `staleTime` marks it fresh, so TanStack Query sees no reason to refetch. Nothing reopens the stream, and the widget silently freezes at the SSR value. +1. **On the server**, waiting for the stream to end blocks rendering forever. +2. **In the browser**, `staleTime` marks the hydrated reading fresh, so nothing reopens the stream. The widget freezes at the SSR value. `StreamQuerySSRUtilsPlugin` fixes both. Register it once as a [plugin](/docs/integrations/tanstack-query#plugins) in `lib/orpc.ts`: @@ -430,7 +430,7 @@ export const orpc = createTanstackQueryUtils(client, { }) ``` -On the server, the plugin settles live queries on their first event and closes the stream, so the HTML ships with a real temperature in it. In the browser, it sets `refetchOnMount: 'always'` on streamed and live queries, so the client reopens the stream after hydration. It also gives streamed queries `initialData: []`, so they render an empty list instead of suspending or reporting `pending` before their stream connects. Per-call options still override anything the plugin sets. The [Streamed and Live Queries](/docs/integrations/tanstack-query#streamed-and-live-queries) docs cover the details. +On the server, the plugin stops a live query at its first event (a streamed query once it connects) and closes the stream. The HTML ships with a real temperature. The plugin also sets `refetchOnMount: 'always'`, so the browser reopens the stream on mount while the server value stays on screen. See [Streamed and Live Queries](/docs/integrations/tanstack-query#streamed-and-live-queries) for details. Reload the page one last time and check your work: diff --git a/apps/content/docs/integrations/pinia-colada.mdx b/apps/content/docs/integrations/pinia-colada.mdx index 687df6304..e666ff6fe 100644 --- a/apps/content/docs/integrations/pinia-colada.mdx +++ b/apps/content/docs/integrations/pinia-colada.mdx @@ -61,46 +61,63 @@ Options accept plain values only. For reactive inputs, pass a callback to `useQu ## Streamed Query Options Utility -Use `.streamedOptions` to build streamed query options for an [AsyncIteratorObject](/docs/async-iterator-object). The resulting data is an array of chunks, and each new chunk is appended as it arrives. It works with `useQuery` and any other API that accepts query options. +Use `.streamedOptions` for a procedure that returns an [AsyncIteratorObject](/docs/async-iterator-object). The data is an array of chunks: each new chunk is appended, and `maxChunks` keeps only the latest. ```ts const query = useQuery( orpc.streamed.streamedOptions({ - input: { id: 123 }, // Specify input if needed - context: { cache: true }, // Provide client context if needed - fnOptions: { - // Configure streamed query behavior - refetchMode: 'reset', - maxChunks: 3, - }, - // additional options... + input: { id: 123 }, + fnOptions: { refetchMode: 'append', maxChunks: 3 }, }), ) ``` -:::info -`refetchMode` determines how data is handled when the query is fetched again: +On refetch, existing data stays until the new stream connects. Then `refetchMode` decides what happens to it: -- `'reset'` _(default)_: Clears existing data and returns the query to a pending state. -- `'append'`: Adds new streamed chunks to the existing data. -- `'replace'`: Buffers streamed data and replaces the cache after the stream completes. +| `refetchMode` | When the stream connects | +| --------------------- | ------------------------------------ | +| `'reset'` _(default)_ | Data becomes `[]` | +| `'append'` | Data is kept, trimmed to `maxChunks` | -::: +:::warning +`refetchMode` and `maxChunks` are part of the query key (when not the default): + +- Queries with different options get separate cache entries and separate streams. To share one, use the same options, and show fewer chunks when reading, e.g. `computed(() => data.value?.slice(-3))`. +- To read or write the cache by key, pass the same options: `streamedKey({ input, fnOptions: { refetchMode, maxChunks } })`. + ::: ## Live Query Options Utility -Use `.liveOptions` to build live query options for an [AsyncIteratorObject](/docs/async-iterator-object). The data always reflects the latest chunk, replacing the previous value whenever a new one arrives. It works with `useQuery` and any other API that accepts query options. +Use `.liveOptions` for a procedure that returns an [AsyncIteratorObject](/docs/async-iterator-object). The data is always the latest chunk. ```ts const query = useQuery( orpc.live.liveOptions({ - input: { id: 123 }, // Specify input if needed - context: { cache: true }, // Provide client context if needed - // additional options... + input: { id: 123 }, }), ) ``` +### Settle on First Data + +By default, a streamed or live fetch resolves only when its stream ends. A stream also closes once the last component using it unmounts. Set `settleOnFirstData` to resolve at the first data instead: when a streamed query connects, or on a live query's first event. The stream then closes and the data is kept. + +| `settleOnFirstData` | Settles at first data | +| ------------------- | --------------------------------------------------- | +| `true` | Always | +| `'unobserved'` | Only if no component or active scope uses the query | +| `false` _(default)_ | Never | + +```ts +const options = orpc.live.liveOptions({ + input: { id: 123 }, + settleOnFirstData: 'unobserved', + staleTime: 0, +}) +``` + +Use `staleTime: 0` so the stream reopens when a component mounts the query. `settleOnFirstData` is not part of the query key, and is ignored with a custom `query`. + ## Infinite Query Options Utility Use `.infiniteOptions` to build infinite query options. It works with `useInfiniteQuery` and any other API that accepts infinite query options. @@ -387,6 +404,39 @@ app.use(PiniaColada, { +## Streamed and Live Queries in SSR + +`StreamQuerySSRUtilsPlugin` renders [streamed](#streamed-query-options-utility) and [live](#live-query-options-utility) queries on the server without waiting for the stream to end, then reopens the streams in the browser. + +```ts +import { StreamQuerySSRUtilsPlugin } from '@orpc/pinia-colada' + +export const orpc = createPiniaColadaUtils(client, { + plugins: [new StreamQuerySSRUtilsPlugin()], +}) +``` + +It sets these defaults on `.streamedOptions` and `.liveOptions`. Per-call and [default options](#default-options) override them. + +| Option | Server | Browser | +| -------------------------------------------- | ------ | -------------- | +| [`settleOnFirstData`](#settle-on-first-data) | `true` | `'unobserved'` | +| `staleTime` | - | `0` | + +- The server renders the first data. +- In the browser, mounting the query reopens the stream while the hydrated data stays visible. +- A fetch no component uses yet, such as `queryCache.fetch` in a navigation guard, resolves at the first data. Pass `settleOnFirstData: false` to wait for the stream to end. + +To make this work: + +- On the server, call `useQuery` in a component setup, so the fetch runs in `onServerPrefetch`. +- Keep `refetchOnMount` at its default. `'always'` restarts a stream that is already open. +- An environment without `window` (or Deno) counts as the server. Pass `isServer` to override. + +:::info +With `staleTime: 0`, a finite stream that has ended refetches when the window regains focus. Set `refetchOnWindowFocus: false` to avoid this. +::: + ## Client Context When a client is invoked through the Pinia Colada integration, an **operation context** is automatically added to the [client context](/docs/client/client-side#client-context). You can use this context to configure request behavior, such as selecting the HTTP method for [RPC Link](/docs/rpc/link#request-method). diff --git a/apps/content/docs/integrations/swr.mdx b/apps/content/docs/integrations/swr.mdx index e75db8ab9..5831bf91a 100644 --- a/apps/content/docs/integrations/swr.mdx +++ b/apps/content/docs/integrations/swr.mdx @@ -84,10 +84,19 @@ import useSWRSubscription from 'swr/subscription' const { data, error } = useSWRSubscription( orpc.streamed.key({ input: { id: 3 } }), - orpc.streamed.subscriber({ context: { cache: true }, maxChunks: 10 }), // Provide client context if needed + orpc.streamed.subscriber({ maxChunks: 10 }), ) ``` +When the subscription restarts, existing data stays until the new stream connects. Then `refetchMode` decides: + +| `refetchMode` | At connect | +| --------------------- | ------------------------------------ | +| `'reset'` _(default)_ | Data becomes `[]` | +| `'append'` | Data is kept, trimmed to `maxChunks` | + +Each chunk is then appended. `maxChunks` keeps only the latest chunks. + Use `.liveSubscriber` to subscribe to the latest events without chunking: ```ts diff --git a/apps/content/docs/integrations/tanstack-query.mdx b/apps/content/docs/integrations/tanstack-query.mdx index 0518a1123..913c482b8 100644 --- a/apps/content/docs/integrations/tanstack-query.mdx +++ b/apps/content/docs/integrations/tanstack-query.mdx @@ -74,52 +74,72 @@ const query = useQuery( ## Streamed Query Options -Use `.streamedOptions` to build streamed query options for an [AsyncIteratorObject](/docs/async-iterator-object). The resulting data is an array of events, and each new event is appended as it arrives. - -It works with `useQuery`, `useSuspenseQuery`, and `prefetchQuery`, and any other API that accepts query options. +Use `.streamedOptions` to build query options for an [AsyncIteratorObject](/docs/async-iterator-object). The data is an array of events, and each new event is appended to it. It works with `useQuery`, `useSuspenseQuery`, `prefetchQuery`, and any other API that accepts query options. ```ts const query = useQuery( orpc.streamed.streamedOptions({ - input: { id: 123 }, // Specify input if needed - context: { cache: true }, // Provide client context if needed + input: { id: 123 }, queryFnOptions: { - // Configure streamed query behavior - refetchMode: 'reset', - maxChunks: 3, + refetchMode: 'append', + maxChunks: 3, // Keep only the last 3 events }, - retry: true, // Infinite retry for more reliable streaming - // additional options... + retry: true, // Retry forever to keep the stream alive }), ) ``` :::info -`refetchMode` determines how data is handled when the query is fetched again: +On refetch, existing data stays until the new stream connects. Then `refetchMode` decides: -- `'reset'` _(default)_: Clears existing data and returns the query to a pending state. -- `'append'`: Adds new streamed chunks to the existing data. -- `'replace'`: Buffers streamed data and replaces the cache after the stream completes. +- `'reset'` _(default)_: data becomes `[]`. +- `'append'`: data is kept, trimmed to `maxChunks`. +For a loading state, call `queryClient.resetQueries({ queryKey })` before refetching. ::: -## Live Query Options +:::warning +`refetchMode` and `maxChunks` are part of the query key (when not the default): + +- Queries with different options get separate cache entries and separate streams. To share one, use the same options, and show fewer chunks with `select`, e.g. `select: data => data.slice(-3)`. +- To read or write the cache by key, pass the same options: `streamedKey({ input, queryFnOptions: { refetchMode, maxChunks } })`. + ::: -Use `.liveOptions` to build live query options for an [AsyncIteratorObject](/docs/async-iterator-object). The data always reflects the latest event, replacing the previous value whenever a new one arrives. +## Live Query Options -It works with `useQuery`, `useSuspenseQuery`, and `prefetchQuery`, and any other API that accepts query options. +Use `.liveOptions` to build query options for an [AsyncIteratorObject](/docs/async-iterator-object). The data is always the latest event. An `undefined` event rejects the fetch, since TanStack Query cannot store `undefined`. ```ts const query = useQuery( orpc.live.liveOptions({ - input: { id: 123 }, // Specify input if needed - context: { cache: true }, // Provide client context if needed - retry: true, // Infinite retry for more reliable streaming - // additional options... + input: { id: 123 }, + retry: true, // Retry forever to keep the stream alive }), ) ``` +### Settle on First Data + +By default, a streamed or live query follows its stream until it ends. So awaiting an endless stream (`fetchQuery`, `ensureQueryData`, route loaders) never resolves. + +Pass `settleOnFirstData` to settle the fetch at the first data instead: when a streamed query connects, or on a live query's first event. The stream closes and the data stays, marked stale. The stream reopens once the query is observed, for example on mount. + +- `true`: always settle. +- `'unobserved'`: settle only if nothing observes the query, such as a mounted `useQuery`. +- `false` _(default)_: never settle. + +```ts +const query = useSuspenseQuery( + orpc.live.liveOptions({ + input: { id: 123 }, + settleOnFirstData: 'unobserved', + }), +) +``` + +- Not part of the query key, and ignored with a custom `queryFn`. +- A stream cut short (cancelled, or its last observer unmounts) is marked stale, so it reopens when observed again. + ## Infinite Query Options Use `.infiniteOptions` to build infinite query options. It works with `useInfiniteQuery`, `useSuspenseInfiniteQuery`, and `prefetchInfiniteQuery`, and any other API that accepts infinite query options. @@ -570,7 +590,7 @@ const queryClient = new QueryClient({ ### Streamed and Live Queries -`StreamQuerySSRUtilsPlugin` lets [streamed](#streamed-query-options) and [live](#live-query-options) queries render on the server without blocking forever: streamed queries are cancelled once they connect and live queries on their first event, keeping the data received so far. In the browser, they refetch on mount to reopen the stream after hydration. +Add `StreamQuerySSRUtilsPlugin` so [streamed](#streamed-query-options) and [live](#live-query-options) queries render on the server without blocking forever. ```ts import { StreamQuerySSRUtilsPlugin } from '@orpc/tanstack-query' @@ -580,4 +600,19 @@ export const orpc = createTanstackQueryUtils(client, { }) ``` -An environment without `window` counts as the server; pass `isServer` to override. +It sets the [`settleOnFirstData`](#settle-on-first-data) option of `.streamedOptions` and `.liveOptions`: + +- Server: `true`. The query is dehydrated with its first data. +- Browser: `'unobserved'`. A mounted `useQuery` follows the stream, as does `suspense()` in Vue, Svelte, and Solid. + +In the browser it also sets `refetchOnMount: 'always'`, so after hydration, mounting reopens the stream (even with a `staleTime`) while the hydrated data stays on screen. A finished finite stream also reruns on each mount; pass `refetchOnMount: true` to respect `staleTime` instead. On the server it sets `refetchOnMount: false`, so components sharing a settled query do not reopen it. + +Per-call and [default options](#default-options) override the plugin. The server is detected with TanStack Query's `environmentManager.isServer()`; pass `isServer` to override. + +:::warning + +- In the browser, `await queryClient.fetchQuery(orpc.x.streamedOptions())` resolves once the stream connects, unless a mounted query already follows that stream, in which case it waits for that stream. Pass `settleOnFirstData: false` to wait for the stream to end. +- Avoid `staleTime: 'static'`, which stops the stream from reopening on mount. +- When [streaming](https://tanstack.com/query/latest/docs/framework/react/guides/advanced-ssr#streaming-with-server-components) pending queries, read them with `useSuspenseQuery`. A `useQuery` that mounts before the hydrated data arrives keeps that data until it goes stale. +- Do not mix suspense and non-suspense consumers of one streamed key across suspense boundaries. A mounted `useQuery` keeps the suspending fetch from settling, so the stream stays open. + ::: diff --git a/apps/content/docs/migrations/from-v1.mdx b/apps/content/docs/migrations/from-v1.mdx index 87fc6e465..9861d09c3 100644 --- a/apps/content/docs/migrations/from-v1.mdx +++ b/apps/content/docs/migrations/from-v1.mdx @@ -1153,10 +1153,25 @@ orpc.live.experimental_liveOptions({ input: {} }) `experimental_defaults` became `scoped`, and the hydration serializer changed from `StandardRPCJsonSerializer` to `RPCJsonSerializer` (see [RPC JSON Serializer](/docs/rpc/serializer#rpc-json-serializer)). +Streamed and live queries changed (see [Streamed Query Options](/docs/integrations/tanstack-query#streamed-query-options)): + +- `refetchMode: 'reset'` keeps the old data until the new stream connects, then sets it to `[]`. (v1 cleared it immediately, returning the query to `pending`.) +- `.liveOptions` rejects an `undefined` event, which TanStack Query cannot store. +- `refetchMode: 'replace'` was removed. To keep the old data until the stream ends, collect the stream in a custom `queryFn`: + +```ts +orpc.streamed.streamedOptions({ + input: {}, + queryFn: async ({ signal }) => Array.fromAsync(await client.streamed({}, { signal })), +}) +``` + ### SWR and Pinia Colada `@orpc/experimental-react-swr` is now `@orpc/swr` (see [SWR Integration](/docs/integrations/swr)), and `@orpc/vue-colada` is now `@orpc/pinia-colada` with `createORPCVueColadaUtils` renamed to `createPiniaColadaUtils` (see [Pinia Colada Integration](/docs/integrations/pinia-colada)). Both switched from `path` to `prefix`, same as TanStack Query. +In SWR, `.subscriber` now writes the data when the stream connects: `[]` with `refetchMode: 'reset'` (instead of `undefined`), and the existing data trimmed to `maxChunks` with `'append'`. + ### NestJS Import `implement`, `ORPCError`, and `onError` from `@orpc/server` instead of `@orpc/nest`, and augment `DefaultInitialContext` instead of `ORPCGlobalContext`. See [NestJS Integration](/docs/integrations/nest). diff --git a/packages/pinia-colada/package.json b/packages/pinia-colada/package.json index cc244710f..276df31ca 100644 --- a/packages/pinia-colada/package.json +++ b/packages/pinia-colada/package.json @@ -59,7 +59,7 @@ "zod": "^4.6.7" }, "peerDependencies": { - "@pinia/colada": ">=1.2.0", + "@pinia/colada": ">=1.4.9", "vue": ">=3.5.17" } } diff --git a/packages/pinia-colada/src/index.test.ts b/packages/pinia-colada/src/index.test.ts index 19c9414af..63c901e5b 100644 --- a/packages/pinia-colada/src/index.test.ts +++ b/packages/pinia-colada/src/index.test.ts @@ -1,7 +1,8 @@ -it('exports createPiniaColadaUtils, PINIA_COLADA_OPERATION_CONTEXT_SYMBOL', async () => { +it('exports createPiniaColadaUtils, PINIA_COLADA_OPERATION_CONTEXT_SYMBOL, StreamQuerySSRUtilsPlugin', async () => { await expect(import('./index.ts')).resolves.toMatchObject({ createPiniaColadaUtils: expect.any(Function), createRouterUtils: expect.any(Function), PINIA_COLADA_OPERATION_CONTEXT_SYMBOL: expect.any(Symbol), + StreamQuerySSRUtilsPlugin: expect.any(Function), }) }) diff --git a/packages/pinia-colada/src/index.ts b/packages/pinia-colada/src/index.ts index a9b69a738..ba5d99001 100644 --- a/packages/pinia-colada/src/index.ts +++ b/packages/pinia-colada/src/index.ts @@ -31,6 +31,8 @@ export { } from './router-utils.ts' export { createRouterUtils as createPiniaColadaUtils } from './router-utils.ts' export { SharedUtils } from './shared-utils.ts' +export * from './ssr.ts' +export type { SettleOnFirstDataOptions } from './stream-cache.ts' export type { SerializableStreamedQueryOptions } from './stream-query.ts' export * from './types.ts' export { diff --git a/packages/pinia-colada/src/key.test.ts b/packages/pinia-colada/src/key.test.ts index e6173a5de..7fd152995 100644 --- a/packages/pinia-colada/src/key.test.ts +++ b/packages/pinia-colada/src/key.test.ts @@ -43,6 +43,24 @@ it('generateOperationKey', () => { ).toEqual(['__prefix__', ['path'], { type: 'query', input: { a: 1 } }]) }) +it('generateOperationKey keeps only non-default streamed fnOptions', () => { + expect( + generateOperationKey(['planet', 'stream'], { + type: 'streamed', + fnOptions: { refetchMode: 'append', maxChunks: 3 }, + }), + ).toEqual([ + ['planet', 'stream'], + { type: 'streamed', fnOptions: { refetchMode: 'append', maxChunks: 3 } }, + ]) + expect( + generateOperationKey(['planet', 'stream'], { + type: 'streamed', + fnOptions: { refetchMode: 'reset', maxChunks: Number.POSITIVE_INFINITY }, + }), + ).toEqual([['planet', 'stream'], { type: 'streamed' }]) +}) + describe('generateOperationKey input meta', () => { it('keys meta by JSON-encoded path', () => { expect( diff --git a/packages/pinia-colada/src/key.ts b/packages/pinia-colada/src/key.ts index 31fd7ecf0..689469465 100644 --- a/packages/pinia-colada/src/key.ts +++ b/packages/pinia-colada/src/key.ts @@ -1,6 +1,6 @@ import { RPCJsonSerializer } from '@orpc/client' import type { PartialDeep, Public } from '@orpc/shared' -import { stringifyJSON } from '@orpc/shared' +import { streamedArrayKeyOptions, stringifyJSON } from '@orpc/shared' import type { EntryKey } from '@pinia/colada' import type { SerializableStreamedQueryOptions } from './stream-query.ts' @@ -62,6 +62,8 @@ export function generateOperationKey( path: string[], options: OperationKeyOptions = {}, ): OperationKey { + const fnOptions = streamedArrayKeyOptions(options.fnOptions) + return [ ...(options.prefix !== undefined ? [options.prefix] : []), path, @@ -70,7 +72,7 @@ export function generateOperationKey( ? serializeInput(options.serializer ?? defaultSerializer, options.input) : {}), ...(options.type !== undefined ? { type: options.type } : {}), - ...(options.fnOptions !== undefined ? { fnOptions: options.fnOptions } : {}), + ...(fnOptions !== undefined ? { fnOptions } : {}), }, ] as OperationKey } diff --git a/packages/pinia-colada/src/live-query.test.tsx b/packages/pinia-colada/src/live-query.test.tsx index a0071817f..79afdbf5b 100644 --- a/packages/pinia-colada/src/live-query.test.tsx +++ b/packages/pinia-colada/src/live-query.test.tsx @@ -1,71 +1,9 @@ -import type { UseQueryOptions } from '@pinia/colada' -import { PiniaColada, useQuery, useQueryCache } from '@pinia/colada' -import { mount } from '@vue/test-utils' -import { createPinia } from 'pinia' -import { defineComponent } from 'vue' +import { promiseWithResolvers } from '@orpc/shared' +import { effectScope, watch } from 'vue' +import { createChunkController, createQueryCache, mountQuery } from '../tests/__shared__/utils.ts' import { liveQuery } from './live-query.ts' -/** - * Mounts a component that runs a real Pinia Colada `useQuery`, - * so tests can observe exactly what users see. - */ -function mountQuery(options: UseQueryOptions) { - return mount( - defineComponent({ - setup() { - const queryCache = useQueryCache() - const query = useQuery(options) - - return { query, queryCache } - }, - render: () => null, - }), - { - global: { - plugins: [createPinia(), PiniaColada], - }, - }, - ) -} - -/** - * An AsyncIterable whose chunks are pushed manually, so tests can assert - * intermediate states deterministically while the stream is still open. - */ -function createChunkController() { - const buffer: ({ value: T } | { done: true })[] = [] - let notify: (() => void) | undefined - - return { - async *stream() { - while (true) { - while (buffer.length === 0) { - await new Promise((resolve) => { - notify = resolve - }) - } - - const item = buffer.shift()! - - if ('done' in item) { - return - } - - yield item.value - } - }, - push(value: T) { - buffer.push({ value }) - notify?.() - }, - close() { - buffer.push({ done: true }) - notify?.() - }, - } -} - describe('liveQuery', () => { it('publishes each chunk to the cache, replacing the previous value', async () => { const controller = createChunkController() @@ -75,50 +13,56 @@ describe('liveQuery', () => { }) controller.push(1) - await vi.waitFor(() => expect(mounted.vm.query.data.value).toEqual(1)) + await vi.waitFor(() => expect(mounted.vm.query.data.value).toEqual(1), { interval: 1 }) expect(mounted.vm.query.asyncStatus.value).toEqual('loading') controller.push(2) - await vi.waitFor(() => expect(mounted.vm.query.data.value).toEqual(2)) + await vi.waitFor(() => expect(mounted.vm.query.data.value).toEqual(2), { interval: 1 }) expect(mounted.vm.query.asyncStatus.value).toEqual('loading') controller.close() - await vi.waitFor(() => expect(mounted.vm.query.asyncStatus.value).toEqual('idle')) + await vi.waitFor(() => expect(mounted.vm.query.asyncStatus.value).toEqual('idle'), { + interval: 1, + }) expect(mounted.vm.query.data.value).toEqual(2) expect(mounted.vm.query.status.value).toEqual('success') expect(mounted.vm.queryCache.getQueryData(['live'])).toEqual(2) }) - it('errors when the stream yields no chunks', async () => { + it('keeps data written to the entry after the last event', async () => { const controller = createChunkController() const mounted = mountQuery({ key: ['live'], query: liveQuery(() => controller.stream()), }) - controller.close() + controller.push(1) + await vi.waitFor(() => expect(mounted.vm.query.data.value).toEqual(1), { interval: 1 }) - await vi.waitFor(() => expect(mounted.vm.query.error.value).toBeInstanceOf(TypeError)) - expect((mounted.vm.query.error.value as any).message).toContain('did not yield any data') - expect(mounted.vm.query.status.value).toEqual('error') - }) + mounted.vm.queryCache.setQueryData(['live'], 2) - it('closes the stream without reading it when aborted before it resolves', async () => { - const started = vi.fn() - const controller = new AbortController() + controller.close() + await vi.waitFor(() => expect(mounted.vm.query.asyncStatus.value).toEqual('idle'), { + interval: 1, + }) + expect(mounted.vm.query.data.value).toEqual(2) + expect(mounted.vm.query.status.value).toEqual('success') + }) - const queryFn = liveQuery(async () => { - controller.abort() - return (async function* () { - started() - yield 1 - })() + it('errors when the stream yields no chunks', async () => { + const controller = createChunkController() + const mounted = mountQuery({ + key: ['live'], + query: liveQuery(() => controller.stream()), }) - await expect(queryFn({ signal: controller.signal } as any)).rejects.toBe( - controller.signal.reason, - ) - expect(started).not.toHaveBeenCalled() + controller.close() + + await vi.waitFor(() => expect(mounted.vm.query.error.value).toBeInstanceOf(TypeError), { + interval: 1, + }) + expect((mounted.vm.query.error.value as any).message).toContain('did not yield any data') + expect(mounted.vm.query.status.value).toEqual('error') }) it('stops streaming when the query is cancelled', async () => { @@ -129,13 +73,140 @@ describe('liveQuery', () => { }) controller.push(1) - await vi.waitFor(() => expect(mounted.vm.query.data.value).toEqual(1)) + await vi.waitFor(() => expect(mounted.vm.query.data.value).toEqual(1), { interval: 1 }) mounted.vm.queryCache.cancelQueries({ key: ['live'] }, new Error('__cancelled__')) controller.push(2) // resumes the stream loop, which now observes the aborted signal - await vi.waitFor(() => expect(mounted.vm.query.asyncStatus.value).toEqual('idle')) + await vi.waitFor(() => expect(mounted.vm.query.asyncStatus.value).toEqual('idle'), { + interval: 1, + }) expect(mounted.vm.query.data.value).toEqual(1) expect(mounted.vm.query.error.value).toBeNull() }) + + it('keeps existing data until the first event on refetch', async () => { + const first = createChunkController() + const second = createChunkController() + + const queryFn = vi + .fn() + .mockImplementationOnce(() => first.stream()) + .mockImplementationOnce(() => second.stream()) + + const mounted = mountQuery({ + key: ['live'], + query: liveQuery(queryFn), + }) + + first.push(1) + first.close() + await vi.waitFor(() => expect(mounted.vm.query.asyncStatus.value).toEqual('idle'), { + interval: 1, + }) + + const statuses: string[] = [] + watch( + () => mounted.vm.query.status.value, + (status) => statuses.push(status), + { flush: 'sync' }, + ) + + const refetching = mounted.vm.query.refetch() + + await vi.waitFor(() => expect(queryFn).toHaveBeenCalledTimes(2), { interval: 1 }) + expect(mounted.vm.query.data.value).toEqual(1) + expect(mounted.vm.query.asyncStatus.value).toEqual('loading') + + second.push(2) + await vi.waitFor(() => expect(mounted.vm.query.data.value).toEqual(2), { interval: 1 }) + + second.close() + await refetching + expect(statuses).not.toContain('pending') + }) + + it.each([true, 'unobserved'] as const)( + 'settles a fetch with no component at the first event with settleOnFirstData: %s', + async (settleOnFirstData) => { + const cleanup = vi.fn() + const queryCache = createQueryCache() + + const entry = queryCache.ensure({ + key: ['live'], + query: liveQuery( + async () => + (async function* () { + try { + yield 1 + await new Promise(() => {}) + } finally { + cleanup() + } + })(), + { settleOnFirstData }, + ), + }) + + await expect(queryCache.fetch(entry)).resolves.toEqual({ + status: 'success', + data: 1, + error: null, + }) + expect(cleanup).toHaveBeenCalledTimes(1) + expect(entry.asyncStatus.value).toEqual('idle') + }, + ) + + it('follows the stream of a mounted useQuery with settleOnFirstData: unobserved', async () => { + const controller = createChunkController() + const mounted = mountQuery({ + key: ['live'], + query: liveQuery(() => controller.stream(), { settleOnFirstData: 'unobserved' }), + }) + + controller.push(1) + await vi.waitFor(() => expect(mounted.vm.query.data.value).toEqual(1), { interval: 1 }) + + controller.push(2) + await vi.waitFor(() => expect(mounted.vm.query.data.value).toEqual(2), { interval: 1 }) + expect(mounted.vm.query.asyncStatus.value).toEqual('loading') + + controller.close() + await vi.waitFor(() => expect(mounted.vm.query.asyncStatus.value).toEqual('idle'), { + interval: 1, + }) + }) + + it('keeps the settled data when the stream reopened for a mounted component yields nothing', async () => { + const closing = promiseWithResolvers() + const queryFn = vi.fn(async () => + queryFn.mock.calls.length > 1 + ? (async function* () {})() + : { + [Symbol.asyncIterator]: () => ({ + next: async () => ({ done: false as const, value: 1 }), + return: async () => { + await closing.promise + return { done: true as const, value: undefined } + }, + }), + }, + ) + const queryCache = createQueryCache() + const entry = queryCache.ensure({ + key: ['live'], + query: liveQuery(queryFn, { settleOnFirstData: 'unobserved' }), + }) + + const fetching = queryCache.fetch(entry) + await vi.waitFor(() => expect(entry.state.value.data).toBe(1), { interval: 1 }) + const scope = effectScope() + queryCache.track(entry, scope) + closing.resolve() + + await expect(fetching).resolves.toEqual({ status: 'success', data: 1, error: null }) + expect(queryFn).toHaveBeenCalledTimes(2) + queryCache.untrack(entry, scope) + }) }) diff --git a/packages/pinia-colada/src/live-query.ts b/packages/pinia-colada/src/live-query.ts index e581bec97..b4ba7d7b0 100644 --- a/packages/pinia-colada/src/live-query.ts +++ b/packages/pinia-colada/src/live-query.ts @@ -1,6 +1,8 @@ import type { Promisable } from '@orpc/shared' -import { stringifyJSON } from '@orpc/shared' +import { pipeStreamToCache, stringifyJSON } from '@orpc/shared' +import type { SettleOnFirstDataOptions } from './stream-cache.ts' +import { followEntryStream } from './stream-cache.ts' import type { UseQueryFnContext } from './types.ts' /** @@ -9,34 +11,22 @@ import type { UseQueryFnContext } from './types.ts' */ export function liveQuery( queryFn: (context: UseQueryFnContext) => Promisable>, + { settleOnFirstData }: SettleOnFirstDataOptions = {}, ): (context: UseQueryFnContext) => Promise { return async (context) => { - const { entry, signal } = context - - const stream = await queryFn(context) - - if (signal.aborted) { - await stream[Symbol.asyncIterator]().return?.() - throw signal.reason - } - - let last: { chunk: T } | undefined - - for await (const chunk of stream) { - if (signal.aborted) { - throw signal.reason - } - - last = { chunk } - entry.state.value = { status: 'success', data: chunk, error: null } as any - } - - if (!last) { + const result = await followEntryStream( + context, + settleOnFirstData, + async (callContext, writes) => + pipeStreamToCache(await queryFn(callContext), { ...writes, onChunk: (_, chunk) => chunk }), + ) + + if (!result) { throw new TypeError( - `Live query for ${stringifyJSON(entry.key)} did not yield any data. Ensure the query function returns an AsyncIterable with at least one chunk.`, + `Live query for ${stringifyJSON(context.entry.key)} did not yield any data. Ensure the query function returns an AsyncIterable with at least one chunk.`, ) } - return last.chunk + return result.data } } diff --git a/packages/pinia-colada/src/procedure-utils.test-d.ts b/packages/pinia-colada/src/procedure-utils.test-d.ts index ffb351c32..1ef8cbd3e 100644 --- a/packages/pinia-colada/src/procedure-utils.test-d.ts +++ b/packages/pinia-colada/src/procedure-utils.test-d.ts @@ -9,6 +9,7 @@ import * as z from 'zod' import type { ProcedureUtils, ProcedureUtilsOptions } from './procedure-utils.ts' import type { InfiniteOptionsIn, + LiveOptionsIn, MutationKeyOptions, MutationOptionsIn, QueryKeyOptions, @@ -113,6 +114,23 @@ describe('ProcedureUtils', () => { streamUtils.streamedOptions({ fnOptions: { refetchMode: 'invalid' } }) // @ts-expect-error invalid input streamUtils.liveOptions({ input: 'invalid' }) + // @ts-expect-error 'replace' is removed + streamUtils.streamedOptions({ fnOptions: { refetchMode: 'replace' } }) + }) + + it('accepts settleOnFirstData on .streamedOptions & .liveOptions only', () => { + streamUtils.streamedOptions({ settleOnFirstData: true }) + streamUtils.streamedOptions({ settleOnFirstData: 'unobserved' }) + streamUtils.liveOptions({ settleOnFirstData: false }) + streamUtils.liveOptions({ settleOnFirstData: 'unobserved' }) + // @ts-expect-error invalid settleOnFirstData + streamUtils.streamedOptions({ settleOnFirstData: 'invalid' }) + // @ts-expect-error invalid settleOnFirstData + streamUtils.liveOptions({ settleOnFirstData: 'invalid' }) + // @ts-expect-error settleOnFirstData is not part of fnOptions + streamUtils.streamedOptions({ fnOptions: { settleOnFirstData: true } }) + // @ts-expect-error settleOnFirstData is not a query option + streamUtils.queryOptions({ settleOnFirstData: true }) }) }) @@ -608,7 +626,7 @@ describe('ProcedureUtilsOptions', () => { } }) - it('liveOptions should accept Partial | modifier function', () => { + it('liveOptions should accept Partial | modifier function', () => { const _defaults: TestOptions = { liveOptions: { input: { search: 'test' }, @@ -626,7 +644,7 @@ describe('ProcedureUtilsOptions', () => { const _defaults2: TestOptions = { liveOptions: (options) => { expectTypeOf(options).toEqualTypeOf< - QueryOptionsIn + LiveOptionsIn >() return { diff --git a/packages/pinia-colada/src/procedure-utils.test.ts b/packages/pinia-colada/src/procedure-utils.test.ts index 5e8032bdf..17ec2db02 100644 --- a/packages/pinia-colada/src/procedure-utils.test.ts +++ b/packages/pinia-colada/src/procedure-utils.test.ts @@ -172,7 +172,7 @@ describe('procedureUtils', () => { expect(streamedQuerySpy).toHaveBeenCalledWith(expect.any(Function), { maxChunks: 3 }) expect(client).toHaveBeenCalledTimes(1) expect(client).toHaveBeenCalledWith(1, { - signal, + signal: expect.any(AbortSignal), context: { [OPERATION_CONTEXT_SYMBOL]: { key: options.key, @@ -182,6 +182,33 @@ describe('procedureUtils', () => { }) }) + it('passes settleOnFirstData to serializableStreamedQuery, not to the key or pinia colada', async () => { + const options = utils.streamedOptions({ + input: 1, + fnOptions: { refetchMode: 'append' }, + settleOnFirstData: 'unobserved', + }) as any + + expect(generateOperationKeySpy).toHaveBeenCalledWith(['ping'], { + prefix: undefined, + type: 'streamed', + input: 1, + fnOptions: { refetchMode: 'append' }, + }) + expect(options).not.toHaveProperty('settleOnFirstData') + + client.mockResolvedValueOnce(stream('a')) + await options.query({ + signal, + entry: { ...entry, state: { value: { ...entry.state.value } } }, + }) + + expect(streamedQuerySpy).toHaveBeenCalledWith(expect.any(Function), { + refetchMode: 'append', + settleOnFirstData: 'unobserved', + }) + }) + it('throws when output is not an AsyncIteratorObject', async () => { const options = utils.streamedOptions({ input: 1 } as any) as any @@ -251,7 +278,7 @@ describe('procedureUtils', () => { ) expect(client).toHaveBeenCalledWith('__override__', { - signal, + signal: expect.any(AbortSignal), context: { batch: true, [OPERATION_CONTEXT_SYMBOL]: { @@ -369,7 +396,7 @@ describe('procedureUtils', () => { expect(liveQuerySpy).toHaveBeenCalledTimes(1) expect(client).toHaveBeenCalledTimes(1) expect(client).toHaveBeenCalledWith(1, { - signal, + signal: expect.any(AbortSignal), context: { [OPERATION_CONTEXT_SYMBOL]: { key: options.key, @@ -379,6 +406,24 @@ describe('procedureUtils', () => { }) }) + it('passes settleOnFirstData to liveQuery, not to the key or pinia colada', async () => { + const options = utils.liveOptions({ input: 1, settleOnFirstData: true }) as any + + expect(generateOperationKeySpy).toHaveBeenCalledWith(['ping'], { + prefix: undefined, + type: 'live', + input: 1, + }) + expect(options).not.toHaveProperty('settleOnFirstData') + + client.mockResolvedValueOnce(stream('a', 'b')) + await expect( + options.query({ signal, entry: { ...entry, state: { value: { ...entry.state.value } } } }), + ).resolves.toEqual('a') + + expect(liveQuerySpy).toHaveBeenCalledWith(expect.any(Function), { settleOnFirstData: true }) + }) + it('throws when output is not an AsyncIteratorObject', async () => { const options = utils.liveOptions({ input: 1 } as any) as any @@ -448,7 +493,7 @@ describe('procedureUtils', () => { ) expect(client).toHaveBeenCalledWith('__override__', { - signal, + signal: expect.any(AbortSignal), context: { batch: true, [OPERATION_CONTEXT_SYMBOL]: { diff --git a/packages/pinia-colada/src/procedure-utils.ts b/packages/pinia-colada/src/procedure-utils.ts index b149e3e66..e085e3a3f 100644 --- a/packages/pinia-colada/src/procedure-utils.ts +++ b/packages/pinia-colada/src/procedure-utils.ts @@ -23,6 +23,7 @@ import type { InfiniteKeyOptions, InfiniteOptionsIn, InfiniteOptionsOut, + LiveOptionsIn, MutationKeyOptions, MutationOptionsIn, MutationOptionsOut, @@ -206,7 +207,7 @@ export interface ProcedureUtilsOptions< */ liveOptions?: | ProcedureUtilsModifier< - QueryOptionsIn< + LiveOptionsIn< TClientContext, TInput, InferLiveQueryOutput, @@ -523,6 +524,7 @@ export class ProcedureUtils< /** * Configure queries for [AsyncIteratorObject](https://orpc.dev/docs/async-iterator-object). * The resulting data is an array of chunks, and each new chunk is appended as it arrives. + * It never returns to pending once it has data. * Works with `useQuery` and any other API that accepts query options. */ streamedOptions | undefined = undefined>( @@ -550,6 +552,7 @@ export class ProcedureUtils< key: _keyIn, query: queryIn, fnOptions, + settleOnFirstData, ...restOptions } = optionsIn as Record @@ -581,15 +584,18 @@ export class ProcedureUtils< > } - return serializableStreamedQuery(async (queryContext) => { - const output = await this.call(input, { signal: queryContext.signal, context }) + return serializableStreamedQuery( + async (queryContext) => { + const output = await this.call(input, { signal: queryContext.signal, context }) - if (!isAsyncIteratorObject(output)) { - throw new Error('streamedQuery requires an AsyncIteratorObject output') - } + if (!isAsyncIteratorObject(output)) { + throw new Error('streamedQuery requires an AsyncIteratorObject output') + } - return output - }, fnOptions)(fnContext) as PromiseWithError, TError> + return output + }, + { ...fnOptions, settleOnFirstData }, + )(fnContext) as PromiseWithError, TError> }, ) }, @@ -625,7 +631,7 @@ export class ProcedureUtils< */ liveOptions | undefined = undefined>( ...rest: MaybeOptionalOptions< - QueryOptionsIn, TError, UInitialData> + LiveOptionsIn, TError, UInitialData> > ): NoInfer, TError, UInitialData>> { let optionsIn = resolveMaybeOptionalOptions(rest) @@ -641,6 +647,7 @@ export class ProcedureUtils< context, key: _keyIn, query: queryIn, + settleOnFirstData, ...restOptions } = optionsIn as Record @@ -669,15 +676,18 @@ export class ProcedureUtils< return queryIn(fnContext) as PromiseWithError, TError> } - return liveQuery(async (queryContext) => { - const output = await this.call(input, { signal: queryContext.signal, context }) + return liveQuery( + async (queryContext) => { + const output = await this.call(input, { signal: queryContext.signal, context }) - if (!isAsyncIteratorObject(output)) { - throw new Error('liveQuery requires an AsyncIteratorObject output') - } + if (!isAsyncIteratorObject(output)) { + throw new Error('liveQuery requires an AsyncIteratorObject output') + } - return output - })(fnContext) + return output + }, + { settleOnFirstData }, + )(fnContext) }, ) }, diff --git a/packages/pinia-colada/src/ssr.test.ts b/packages/pinia-colada/src/ssr.test.ts new file mode 100644 index 000000000..256e87837 --- /dev/null +++ b/packages/pinia-colada/src/ssr.test.ts @@ -0,0 +1,76 @@ +import type { Client } from '@orpc/client' + +import { createQueryCache } from '../tests/__shared__/utils.ts' +import { createRouterUtils } from './router-utils.ts' +import { StreamQuerySSRUtilsPlugin } from './ssr.ts' + +describe('streamQuerySSRUtilsPlugin', () => { + afterEach(() => { + vi.unstubAllGlobals() + }) + + it.each([ + [true, { settleOnFirstData: true }], + [false, { settleOnFirstData: 'unobserved', staleTime: 0 }], + ] as const)( + 'injects options with isServer: %s, given or detected by window', + (isServer, defaults) => { + const expected = { streamedOptions: defaults, liveOptions: defaults } + + expect(new StreamQuerySSRUtilsPlugin({ isServer }).initProcedureOptions([], {})).toEqual( + expected, + ) + + if (!isServer) { + vi.stubGlobal('window', {}) + } + expect(new StreamQuerySSRUtilsPlugin().initProcedureOptions([], {})).toEqual(expected) + }, + ) + + it('detects Deno as the server', () => { + vi.stubGlobal('window', {}) + vi.stubGlobal('Deno', {}) + + expect(new StreamQuerySSRUtilsPlugin().initProcedureOptions([], {})).toEqual({ + streamedOptions: { settleOnFirstData: true }, + liveOptions: { settleOnFirstData: true }, + }) + }) + + it('lets utils and call options override its defaults', async () => { + const cleanup = vi.fn() + const client = { + stream: vi.fn(async function* () { + try { + yield 1 + yield 2 + } finally { + cleanup() + } + }), + } as unknown as { stream: Client, Error> } + + const plugin = new StreamQuerySSRUtilsPlugin({ isServer: true }) + + expect( + plugin.initProcedureOptions([], { + streamedOptions: { settleOnFirstData: false, staleTime: 1000 }, + }), + ).toEqual({ + streamedOptions: { settleOnFirstData: false, staleTime: 1000 }, + liveOptions: { settleOnFirstData: true }, + }) + + const utils = createRouterUtils(client, { plugins: [plugin] }) + const queryCache = createQueryCache() + + const options = utils.stream.streamedOptions({ settleOnFirstData: false, staleTime: 1000 }) + expect(options.staleTime).toBe(1000) + + await expect(queryCache.fetch(queryCache.ensure(options))).resolves.toMatchObject({ + data: [1, 2], + }) + expect(cleanup).toHaveBeenCalledTimes(1) + }) +}) diff --git a/packages/pinia-colada/src/ssr.ts b/packages/pinia-colada/src/ssr.ts new file mode 100644 index 000000000..09f716cc0 --- /dev/null +++ b/packages/pinia-colada/src/ssr.ts @@ -0,0 +1,50 @@ +import type { AnyNestedClient, InferClientContext, InferClientError } from '@orpc/client' + +import type { RouterUtilsPlugin } from './plugin.ts' +import type { ProcedureUtilsOptions } from './procedure-utils.ts' +import { mergeProcedureUtilsOptions } from './procedure-utils.ts' + +export interface StreamQuerySSRUtilsPluginOptions { + /** + * Whether utils run on the server, where streamed and live queries always settle at their first data. + * Defaults to `true` when there is no global `window`, or in Deno. + */ + isServer?: boolean | undefined +} + +/** + * Router utils plugin for SSR of streamed and live queries: sets `settleOnFirstData` (`true` on the server, `'unobserved'` in the browser) + * and, in the browser, `staleTime: 0`, so settled data is stale and its stream reopens when a component mounts. + * + * @see {@link https://orpc.dev/docs/integrations/pinia-colada#streamed-and-live-queries-in-ssr | Pinia Colada Integration - Streamed and Live Queries in SSR} + */ +export class StreamQuerySSRUtilsPlugin< + T extends AnyNestedClient = AnyNestedClient, +> implements RouterUtilsPlugin { + readonly name = '~stream-query-ssr' + + private readonly defaults: + | { settleOnFirstData: true } + | { settleOnFirstData: 'unobserved'; staleTime: 0 } + + constructor(options: StreamQuerySSRUtilsPluginOptions = {}) { + // the server always settles, even when observed, and keeps the staleTime so components sharing a query do not reopen it + this.defaults = + (options.isServer ?? (!('window' in globalThis) || 'Deno' in globalThis)) + ? { settleOnFirstData: true } + : { settleOnFirstData: 'unobserved', staleTime: 0 } + } + + initProcedureOptions( + _path: string[], + options: ProcedureUtilsOptions, any, any, InferClientError>, + ): ProcedureUtilsOptions, any, any, InferClientError> { + return mergeProcedureUtilsOptions( + { + streamedOptions: this.defaults, + liveOptions: this.defaults, + }, + options, + ) + } +} diff --git a/packages/pinia-colada/src/stream-cache.ts b/packages/pinia-colada/src/stream-cache.ts new file mode 100644 index 000000000..b16953c51 --- /dev/null +++ b/packages/pinia-colada/src/stream-cache.ts @@ -0,0 +1,78 @@ +import type { PipeStreamToCacheOptions } from '@orpc/shared' +import { runWithOwnSignal } from '@orpc/shared' + +import type { UseQueryFnContext } from './types.ts' + +export interface SettleOnFirstDataOptions { + /** + * Settles the fetch at its first data (streamed: when the stream connects; live: on the first event) + * instead of following the stream to its end. The stream is closed and the data kept. + * The stream reopens on mount only once the entry is stale (e.g. with `staleTime: 0`, as `StreamQuerySSRUtilsPlugin` sets). + * - `true`: always + * - `'unobserved'`: only if no observer watches the query at that moment + * - `false`: never + * @default false + */ + settleOnFirstData?: boolean | 'unobserved' | undefined +} + +export type StreamWrites = Required< + Pick, 'signal' | 'update'> +> & + Pick, 'settle'> + +/** + * Runs a fetch's stream against its entry, giving each call its own signal, aborted once its stream is no longer read. + * - A fetch aborted before it ends (cancelled, superseded, or its last component left) marks the entry stale, + * so the stream reopens when a component uses it again. + * - If a component mounted while a settled stream was closing, a fresh stream is followed instead. + * @returns the data to resolve the fetch with, or `undefined` if nothing was written + */ +export async function followEntryStream( + context: UseQueryFnContext, + settleOnFirstData: SettleOnFirstDataOptions['settleOnFirstData'], + run: ( + callContext: UseQueryFnContext, + writes: StreamWrites, + ) => Promise<{ data: TData } | undefined>, +): Promise<{ data: TData } | undefined> { + const { entry, signal } = context + + signal.addEventListener( + 'abort', + () => { + entry.when = 0 + }, + { once: true }, + ) + + const update: StreamWrites['update'] = (updater) => { + const { status, data: current } = entry.state.value + const data = updater(current as TData | undefined) + + // an unchanged write (e.g. 'append' at connect) is not new data + if (data !== current || status !== 'success') { + entry.state.value = { status: 'success', data, error: null } as any + } + } + + let settled = false + const settle = settleOnFirstData + ? () => (settled = settleOnFirstData === true || !entry.active) + : undefined + + const follow = (settle: StreamWrites['settle']) => + runWithOwnSignal(signal, (callSignal) => + run({ ...context, signal: callSignal }, { signal, update, settle }), + ) + + let result = await follow(settle) + + if (settled && settleOnFirstData === 'unobserved' && entry.active && !signal.aborted) { + // a fresh stream that yields nothing keeps the settled data + result = (await follow(undefined)) ?? result + } + + // Colada writes the returned data again, so return the entry data to keep later writes + return result && { data: entry.state.value.data as TData } +} diff --git a/packages/pinia-colada/src/stream-query.test.tsx b/packages/pinia-colada/src/stream-query.test.tsx index 7de38d684..7f2116ff6 100644 --- a/packages/pinia-colada/src/stream-query.test.tsx +++ b/packages/pinia-colada/src/stream-query.test.tsx @@ -1,72 +1,9 @@ import { promiseWithResolvers } from '@orpc/shared' -import type { UseQueryOptions } from '@pinia/colada' -import { PiniaColada, useQuery, useQueryCache } from '@pinia/colada' -import { mount } from '@vue/test-utils' -import { createPinia } from 'pinia' -import { defineComponent } from 'vue' +import { effectScope, watch } from 'vue' +import { createChunkController, createQueryCache, mountQuery } from '../tests/__shared__/utils.ts' import { serializableStreamedQuery } from './stream-query.ts' -/** - * Mounts a component that runs a real Pinia Colada `useQuery`, - * so tests can observe exactly what users see. - */ -function mountQuery(options: UseQueryOptions) { - return mount( - defineComponent({ - setup() { - const queryCache = useQueryCache() - const query = useQuery(options) - - return { query, queryCache } - }, - render: () => null, - }), - { - global: { - plugins: [createPinia(), PiniaColada], - }, - }, - ) -} - -/** - * An AsyncIterable whose chunks are pushed manually, so tests can assert - * intermediate states deterministically while the stream is still open. - */ -function createChunkController() { - const buffer: ({ value: T } | { done: true })[] = [] - let notify: (() => void) | undefined - - return { - async *stream() { - while (true) { - while (buffer.length === 0) { - await new Promise((resolve) => { - notify = resolve - }) - } - - const item = buffer.shift()! - - if ('done' in item) { - return - } - - yield item.value - } - }, - push(value: T) { - buffer.push({ value }) - notify?.() - }, - close() { - buffer.push({ done: true }) - notify?.() - }, - } -} - describe('serializableStreamedQuery', () => { it('publishes each chunk to the cache as it arrives on first fetch', async () => { const controller = createChunkController() @@ -75,37 +12,35 @@ describe('serializableStreamedQuery', () => { query: serializableStreamedQuery(() => controller.stream()), }) - await vi.waitFor(() => expect(mounted.vm.query.data.value).toEqual([])) + await vi.waitFor(() => expect(mounted.vm.query.data.value).toEqual([]), { interval: 1 }) expect(mounted.vm.query.asyncStatus.value).toEqual('loading') controller.push(1) - await vi.waitFor(() => expect(mounted.vm.query.data.value).toEqual([1])) + await vi.waitFor(() => expect(mounted.vm.query.data.value).toEqual([1]), { interval: 1 }) expect(mounted.vm.query.asyncStatus.value).toEqual('loading') controller.push(2) - await vi.waitFor(() => expect(mounted.vm.query.data.value).toEqual([1, 2])) + await vi.waitFor(() => expect(mounted.vm.query.data.value).toEqual([1, 2]), { interval: 1 }) controller.close() - await vi.waitFor(() => expect(mounted.vm.query.asyncStatus.value).toEqual('idle')) + await vi.waitFor(() => expect(mounted.vm.query.asyncStatus.value).toEqual('idle'), { + interval: 1, + }) expect(mounted.vm.query.data.value).toEqual([1, 2]) expect(mounted.vm.query.status.value).toEqual('success') expect(mounted.vm.queryCache.getQueryData(['stream'])).toEqual([1, 2]) }) - it('resets previous data to pending by default on refetch', async () => { + it('keeps previous data until connect, then resets to [] by default on refetch', async () => { const first = createChunkController() const second = createChunkController() - - let releaseSecond!: () => void - const secondReady = new Promise((resolve) => { - releaseSecond = resolve - }) + const secondReady = promiseWithResolvers() const queryFn = vi .fn() .mockImplementationOnce(() => first.stream()) .mockImplementationOnce(async () => { - await secondReady + await secondReady.promise return second.stream() }) @@ -116,25 +51,37 @@ describe('serializableStreamedQuery', () => { first.push('old') first.close() - await vi.waitFor(() => expect(mounted.vm.query.data.value).toEqual(['old'])) + await vi.waitFor(() => expect(mounted.vm.query.asyncStatus.value).toEqual('idle'), { + interval: 1, + }) + expect(mounted.vm.query.data.value).toEqual(['old']) + + const statuses: string[] = [] + watch( + () => mounted.vm.query.status.value, + (status) => statuses.push(status), + { flush: 'sync' }, + ) const refetching = mounted.vm.query.refetch() - // previous data is cleared and the query goes back to pending - await vi.waitFor(() => expect(mounted.vm.query.status.value).toEqual('pending')) - expect(mounted.vm.query.data.value).toBeUndefined() + await vi.waitFor(() => expect(queryFn).toHaveBeenCalledTimes(2), { interval: 1 }) + expect(mounted.vm.query.data.value).toEqual(['old']) + expect(mounted.vm.query.status.value).toEqual('success') + expect(mounted.vm.query.asyncStatus.value).toEqual('loading') - releaseSecond() - await vi.waitFor(() => expect(mounted.vm.query.data.value).toEqual([])) + secondReady.resolve() + await vi.waitFor(() => expect(mounted.vm.query.data.value).toEqual([]), { interval: 1 }) second.push('a') - await vi.waitFor(() => expect(mounted.vm.query.data.value).toEqual(['a'])) + await vi.waitFor(() => expect(mounted.vm.query.data.value).toEqual(['a']), { interval: 1 }) expect(mounted.vm.query.asyncStatus.value).toEqual('loading') second.close() await refetching expect(mounted.vm.query.data.value).toEqual(['a']) expect(mounted.vm.query.asyncStatus.value).toEqual('idle') + expect(statuses).not.toContain('pending') }) it('appends to previous data with refetchMode=append', async () => { @@ -153,13 +100,14 @@ describe('serializableStreamedQuery', () => { first.push('old') first.close() - await vi.waitFor(() => expect(mounted.vm.query.data.value).toEqual(['old'])) + await vi.waitFor(() => expect(mounted.vm.query.data.value).toEqual(['old']), { interval: 1 }) const refetching = mounted.vm.query.refetch() second.push('a') - // new chunks are appended to the previous data while the stream is open - await vi.waitFor(() => expect(mounted.vm.query.data.value).toEqual(['old', 'a'])) + await vi.waitFor(() => expect(mounted.vm.query.data.value).toEqual(['old', 'a']), { + interval: 1, + }) expect(mounted.vm.query.asyncStatus.value).toEqual('loading') second.close() @@ -167,51 +115,78 @@ describe('serializableStreamedQuery', () => { expect(mounted.vm.query.data.value).toEqual(['old', 'a']) }) - it('buffers and replaces with refetchMode=replace', async () => { + it('trims previous data to maxChunks at connect with refetchMode=append', async () => { const first = createChunkController() const second = createChunkController() + const secondReady = promiseWithResolvers() const queryFn = vi .fn() .mockImplementationOnce(() => first.stream()) - .mockImplementationOnce(() => second.stream()) + .mockImplementationOnce(async () => { + await secondReady.promise + return second.stream() + }) const mounted = mountQuery({ key: ['stream'], - query: serializableStreamedQuery(queryFn, { refetchMode: 'replace' }), + query: serializableStreamedQuery(queryFn, { refetchMode: 'append', maxChunks: 2 }), }) - first.push('old') + first.push('old1') + first.push('old2') + first.push('old3') first.close() - await vi.waitFor(() => expect(mounted.vm.query.data.value).toEqual(['old'])) + await vi.waitFor(() => expect(mounted.vm.query.asyncStatus.value).toEqual('idle'), { + interval: 1, + }) + expect(mounted.vm.query.data.value).toEqual(['old2', 'old3']) + mounted.vm.queryCache.setQueryData(['stream'], ['old1', 'old2', 'old3']) const refetching = mounted.vm.query.refetch() + await vi.waitFor(() => expect(queryFn).toHaveBeenCalledTimes(2), { interval: 1 }) + expect(mounted.vm.query.data.value).toEqual(['old1', 'old2', 'old3']) + + secondReady.resolve() + await vi.waitFor(() => expect(mounted.vm.query.data.value).toEqual(['old2', 'old3']), { + interval: 1, + }) + second.push('a') - second.push('b') - // previous data stays visible while the new stream is buffered - await vi.waitFor(() => expect(mounted.vm.query.asyncStatus.value).toEqual('loading')) - expect(mounted.vm.query.data.value).toEqual(['old']) + await vi.waitFor(() => expect(mounted.vm.query.data.value).toEqual(['old3', 'a']), { + interval: 1, + }) second.close() await refetching - expect(mounted.vm.query.data.value).toEqual(['a', 'b']) }) - it('updates the cache during stream on first fetch even with refetchMode=replace', async () => { + it('keeps data written to the entry during the stream and after the last chunk', async () => { const controller = createChunkController() const mounted = mountQuery({ key: ['stream'], - query: serializableStreamedQuery(() => controller.stream(), { refetchMode: 'replace' }), + query: serializableStreamedQuery(() => controller.stream()), }) controller.push('a') - await vi.waitFor(() => expect(mounted.vm.query.data.value).toEqual(['a'])) - expect(mounted.vm.query.asyncStatus.value).toEqual('loading') + await vi.waitFor(() => expect(mounted.vm.query.data.value).toEqual(['a']), { interval: 1 }) + + mounted.vm.queryCache.setQueryData(['stream'], ['during']) + + controller.push('b') + await vi.waitFor(() => expect(mounted.vm.query.data.value).toEqual(['during', 'b']), { + interval: 1, + }) + + mounted.vm.queryCache.setQueryData(['stream'], ['during', 'b', 'after']) controller.close() - await vi.waitFor(() => expect(mounted.vm.query.asyncStatus.value).toEqual('idle')) - expect(mounted.vm.query.data.value).toEqual(['a']) + await vi.waitFor(() => expect(mounted.vm.query.asyncStatus.value).toEqual('idle'), { + interval: 1, + }) + expect(mounted.vm.query.data.value).toEqual(['during', 'b', 'after']) + expect(mounted.vm.query.status.value).toEqual('success') }) it('limits chunks with maxChunks', async () => { @@ -231,15 +206,19 @@ describe('serializableStreamedQuery', () => { first.push('old1') first.push('old2') first.close() - await vi.waitFor(() => expect(mounted.vm.query.data.value).toEqual(['old1', 'old2'])) + await vi.waitFor(() => expect(mounted.vm.query.data.value).toEqual(['old1', 'old2']), { + interval: 1, + }) const refetching = mounted.vm.query.refetch() second.push('a') - await vi.waitFor(() => expect(mounted.vm.query.data.value).toEqual(['old2', 'a'])) + await vi.waitFor(() => expect(mounted.vm.query.data.value).toEqual(['old2', 'a']), { + interval: 1, + }) second.push('b') - await vi.waitFor(() => expect(mounted.vm.query.data.value).toEqual(['a', 'b'])) + await vi.waitFor(() => expect(mounted.vm.query.data.value).toEqual(['a', 'b']), { interval: 1 }) second.close() await refetching @@ -269,14 +248,14 @@ describe('serializableStreamedQuery', () => { query: serializableStreamedQuery(queryFn), }) - await vi.waitFor(() => expect(queryFn).toHaveBeenCalledTimes(1)) + await vi.waitFor(() => expect(queryFn).toHaveBeenCalledTimes(1), { interval: 1 }) const refetching = mounted.vm.query.refetch() second.push('new') - await vi.waitFor(() => expect(mounted.vm.query.data.value).toEqual(['new'])) + await vi.waitFor(() => expect(mounted.vm.query.data.value).toEqual(['new']), { interval: 1 }) firstReady.resolve() - await vi.waitFor(() => expect(firstClosed).toHaveBeenCalled()) + await vi.waitFor(() => expect(firstClosed).toHaveBeenCalled(), { interval: 1 }) expect(firstStarted).not.toHaveBeenCalled() expect(mounted.vm.query.data.value).toEqual(['new']) @@ -292,13 +271,274 @@ describe('serializableStreamedQuery', () => { }) controller.push('a') - await vi.waitFor(() => expect(mounted.vm.query.data.value).toEqual(['a'])) + await vi.waitFor(() => expect(mounted.vm.query.data.value).toEqual(['a']), { interval: 1 }) mounted.vm.queryCache.cancelQueries({ key: ['stream'] }, new Error('__cancelled__')) controller.push('b') // resumes the stream loop, which now observes the aborted signal - await vi.waitFor(() => expect(mounted.vm.query.asyncStatus.value).toEqual('idle')) + await vi.waitFor(() => expect(mounted.vm.query.asyncStatus.value).toEqual('idle'), { + interval: 1, + }) expect(mounted.vm.query.data.value).toEqual(['a']) expect(mounted.vm.query.error.value).toBeNull() }) + + it('closes the stream when the last component unmounts before it connects', async () => { + const ready = promiseWithResolvers() + const started = vi.fn() + const stream = (async function* () { + started() + yield 'new' + })() + const closed = vi.spyOn(stream, 'return') + let signal: AbortSignal | undefined + + const queryFn = vi + .fn() + .mockImplementationOnce(async () => + (async function* () { + yield 'old' + })(), + ) + .mockImplementationOnce(async (context) => { + signal = context.signal + await ready.promise + return stream + }) + + const mounted = mountQuery({ + key: ['stream'], + query: serializableStreamedQuery(queryFn), + }) + const queryCache = mounted.vm.queryCache + + await vi.waitFor(() => expect(mounted.vm.query.asyncStatus.value).toEqual('idle'), { + interval: 1, + }) + void mounted.vm.query.refetch() + await vi.waitFor(() => expect(queryFn).toHaveBeenCalledTimes(2), { interval: 1 }) + + mounted.unmount() + ready.resolve() + + await vi.waitFor(() => expect(closed).toHaveBeenCalled(), { interval: 1 }) + expect(signal?.aborted).toBe(true) + expect(started).not.toHaveBeenCalled() + expect(queryCache.getQueryData(['stream'])).toEqual(['old']) + }) + + it('closes the stream once the last component unmounts mid-stream', async () => { + const controller = createChunkController() + const cleanup = vi.fn() + let signal: AbortSignal | undefined + + const mounted = mountQuery({ + key: ['stream'], + query: serializableStreamedQuery(async (context) => { + signal = context.signal + return (async function* () { + try { + yield* controller.stream() + } finally { + cleanup() + } + })() + }), + }) + const queryCache = mounted.vm.queryCache + + controller.push(1) + await vi.waitFor(() => expect(mounted.vm.query.data.value).toEqual([1]), { interval: 1 }) + + mounted.unmount() + controller.push(2) + + await vi.waitFor(() => expect(cleanup).toHaveBeenCalledTimes(1), { interval: 1 }) + expect(signal?.aborted).toBe(true) + expect(queryCache.getQueryData(['stream'])).toEqual([1]) + }) + + it('marks an entry stale when its refetch is aborted because the last component unmounted', async () => { + const controller = createChunkController() + const signals: AbortSignal[] = [] + + const mounted = mountQuery({ + key: ['stream'], + staleTime: 60_000, + query: serializableStreamedQuery(async ({ signal }) => { + signals.push(signal) + return signals.length === 1 + ? (async function* () { + yield 1 + })() + : controller.stream() + }), + }) + const queryCache = mounted.vm.queryCache + const entry = () => queryCache.getEntries({ key: ['stream'], exact: true })[0]! + + await vi.waitFor( + () => expect(entry().state.value).toMatchObject({ status: 'success', data: [1] }), + { interval: 1 }, + ) + await vi.waitFor(() => expect(entry().pending).toBeNull(), { interval: 1 }) + expect(entry().stale).toBe(false) + + void mounted.vm.query.refetch() + controller.push(2) + await vi.waitFor(() => expect(entry().state.value.data).toEqual([2]), { interval: 1 }) + + mounted.unmount() + await new Promise((resolve) => setTimeout(resolve, 0)) + controller.push(3) + + await vi.waitFor(() => expect(signals[1]!.aborted).toBe(true), { interval: 1 }) + expect(entry().stale).toBe(true) + expect(entry().state.value.data).toEqual([2]) + }) + + it('follows the stream of a fetch no component ever used', async () => { + const controller = createChunkController() + const queryCache = createQueryCache() + queryCache.setQueryData(['stream'], [0]) + + const entry = queryCache.ensure({ + key: ['stream'], + query: serializableStreamedQuery(() => controller.stream()), + }) + const fetching = queryCache.fetch(entry) + + controller.push(1) + controller.push(2) + controller.close() + + await expect(fetching).resolves.toEqual({ status: 'success', data: [1, 2], error: null }) + }) + + it.each([true, 'unobserved'] as const)( + 'settles a fetch with no component at connect (settleOnFirstData: %s)', + async (settleOnFirstData) => { + const started = vi.fn() + const queryCache = createQueryCache() + queryCache.setQueryData(['stream'], ['old']) + + const entry = queryCache.ensure({ + key: ['stream'], + query: serializableStreamedQuery( + async () => + (async function* () { + started() + yield 'a' + })(), + { settleOnFirstData }, + ), + }) + + await expect(queryCache.fetch(entry)).resolves.toEqual({ + status: 'success', + data: [], + error: null, + }) + expect(started).not.toHaveBeenCalled() + expect(entry.asyncStatus.value).toEqual('idle') + }, + ) + + it('follows the stream of a mounted useQuery with settleOnFirstData: unobserved', async () => { + const controller = createChunkController() + const mounted = mountQuery({ + key: ['stream'], + query: serializableStreamedQuery(() => controller.stream(), { + settleOnFirstData: 'unobserved', + }), + }) + + await vi.waitFor(() => expect(mounted.vm.query.data.value).toEqual([]), { interval: 1 }) + + controller.push(1) + controller.push(2) + await vi.waitFor(() => expect(mounted.vm.query.data.value).toEqual([1, 2]), { interval: 1 }) + expect(mounted.vm.query.asyncStatus.value).toEqual('loading') + + controller.close() + await vi.waitFor(() => expect(mounted.vm.query.asyncStatus.value).toEqual('idle'), { + interval: 1, + }) + expect(mounted.vm.query.data.value).toEqual([1, 2]) + }) + + it('settles a mounted useQuery with settleOnFirstData: true', async () => { + const controller = createChunkController() + const mounted = mountQuery({ + key: ['stream'], + query: serializableStreamedQuery(() => controller.stream(), { settleOnFirstData: true }), + }) + + await vi.waitFor(() => expect(mounted.vm.query.asyncStatus.value).toEqual('idle'), { + interval: 1, + }) + expect(mounted.vm.query.data.value).toEqual([]) + expect(mounted.vm.query.status.value).toEqual('success') + }) + + it('follows a fresh stream when a component mounts while a settled stream closes (unobserved)', async () => { + const closing = promiseWithResolvers() + const controller = createChunkController() + const queryFn = vi.fn(async () => + queryFn.mock.calls.length > 1 + ? controller.stream() + : { + [Symbol.asyncIterator]: () => ({ + next: async () => ({ done: false as const, value: 0 }), + return: async () => { + await closing.promise + return { done: true as const, value: undefined } + }, + }), + }, + ) + const queryCache = createQueryCache() + const entry = queryCache.ensure({ + key: ['stream'], + query: serializableStreamedQuery(queryFn, { settleOnFirstData: 'unobserved' }), + }) + + const fetching = queryCache.fetch(entry) + await vi.waitFor(() => expect(entry.state.value.data).toEqual([]), { interval: 1 }) + const scope = effectScope() + queryCache.track(entry, scope) + closing.resolve() + + await vi.waitFor(() => expect(queryFn).toHaveBeenCalledTimes(2), { interval: 1 }) + controller.push(1) + controller.close() + await expect(fetching).resolves.toEqual({ status: 'success', data: [1], error: null }) + queryCache.untrack(entry, scope) + }) + + it('does not write unchanged data at connect (append)', async () => { + const queryCache = createQueryCache() + queryCache.setQueryData(['stream'], ['old']) + const entry = queryCache.ensure({ + key: ['stream'], + query: serializableStreamedQuery( + async () => + (async function* () { + yield 'a' + })(), + { refetchMode: 'append', settleOnFirstData: true }, + ), + }) + const writes = vi.fn() + const stop = watch(entry.state, writes, { flush: 'sync' }) + + await expect(queryCache.fetch(entry)).resolves.toEqual({ + status: 'success', + data: ['old'], + error: null, + }) + // only Colada's own write once the fetch resolves + expect(writes).toHaveBeenCalledTimes(1) + stop() + }) }) diff --git a/packages/pinia-colada/src/stream-query.ts b/packages/pinia-colada/src/stream-query.ts index f754c452a..11d922967 100644 --- a/packages/pinia-colada/src/stream-query.ts +++ b/packages/pinia-colada/src/stream-query.ts @@ -1,27 +1,11 @@ -import type { Promisable } from '@orpc/shared' +import type { Promisable, StreamedArrayOptions } from '@orpc/shared' +import { limitArraySize, pipeStreamToCache, streamedArrayWrites } from '@orpc/shared' +import type { SettleOnFirstDataOptions } from './stream-cache.ts' +import { followEntryStream } from './stream-cache.ts' import type { UseQueryFnContext } from './types.ts' -export interface SerializableStreamedQueryOptions { - /** - * Determines how data is handled when the query is fetched again - * - * - `'reset'`: Clears existing data and returns the query to a pending state. - * - `'append'`: Adds new streamed chunks to the existing data. - * - `'replace'`: Buffers streamed data and replaces the cache after the stream completes. - * - * @default 'reset' - */ - refetchMode?: 'append' | 'reset' | 'replace' | undefined - - /** - * Limits the number of data chunks stored in the query result. - * Older chunks are removed when the limit is reached. - * - * @default Number.POSITIVE_INFINITY (unlimited) - */ - maxChunks?: number | undefined -} +export type SerializableStreamedQueryOptions = StreamedArrayOptions /** * A Pinia Colada query function that accumulates streamed chunks into an array, @@ -34,56 +18,21 @@ export interface SerializableStreamedQueryOptions { export function serializableStreamedQuery( queryFn: (context: UseQueryFnContext) => Promisable>, { - refetchMode = 'reset', - maxChunks = Number.POSITIVE_INFINITY, - }: SerializableStreamedQueryOptions = {}, + settleOnFirstData, + ...options + }: SerializableStreamedQueryOptions & SettleOnFirstDataOptions = {}, ): (context: UseQueryFnContext) => Promise { return async (context) => { - const { entry, signal } = context - - const previousData = entry.state.value.data as T[] | undefined - const hasPreviousData = previousData !== undefined - - if (hasPreviousData && refetchMode === 'reset') { - entry.state.value = { status: 'pending', data: undefined, error: null } - } - - const stream = await queryFn(context) - - if (signal.aborted) { - await stream[Symbol.asyncIterator]().return?.() - throw signal.reason - } - - const shouldUpdateCacheDuringStream = !hasPreviousData || refetchMode !== 'replace' - - let result: T[] = - hasPreviousData && refetchMode === 'append' ? limitArraySize(previousData, maxChunks) : [] - - if (shouldUpdateCacheDuringStream) { - entry.state.value = { status: 'success', data: result, error: null } as any - } - - for await (const chunk of stream) { - if (signal.aborted) { - throw signal.reason - } - - result = limitArraySize([...result, chunk], maxChunks) - - if (shouldUpdateCacheDuringStream) { - entry.state.value = { status: 'success', data: result, error: null } as any - } - } - - return result + const result = await followEntryStream( + context, + settleOnFirstData, + async (callContext, writes) => + pipeStreamToCache(await queryFn(callContext), { + ...writes, + ...streamedArrayWrites(options), + }), + ) + + return limitArraySize(result!.data, options.maxChunks ?? Number.POSITIVE_INFINITY) } } - -function limitArraySize(items: T[], maxSize: number): T[] { - if (items.length <= maxSize) { - return items - } - - return items.slice(items.length - maxSize) -} diff --git a/packages/pinia-colada/src/types.ts b/packages/pinia-colada/src/types.ts index 294eac037..0fa2a684a 100644 --- a/packages/pinia-colada/src/types.ts +++ b/packages/pinia-colada/src/types.ts @@ -12,6 +12,7 @@ import type { UseQueryOptions, } from '@pinia/colada' +import type { SettleOnFirstDataOptions } from './stream-cache.ts' import type { SerializableStreamedQueryOptions } from './stream-query.ts' export type InferStreamedQueryOutput = TOutput extends AsyncIterable ? U[] : never @@ -87,7 +88,7 @@ export type StreamedOptionsIn< TInitialData extends TOutput | undefined, > = QueryOptionsIn & { fnOptions?: SerializableStreamedQueryOptions | undefined -} +} & SettleOnFirstDataOptions export type StreamedOptionsOut< TOutput, @@ -95,6 +96,14 @@ export type StreamedOptionsOut< TInitialData extends TOutput | undefined, > = QueryOptionsOut +export type LiveOptionsIn< + TClientContext extends ClientContext, + TInput, + TOutput, + TError, + TInitialData extends TOutput | undefined, +> = QueryOptionsIn & SettleOnFirstDataOptions + export type InfiniteKeyOptions = | Pick, 'input' | 'initialPageParam'> | { key: EntryKey } diff --git a/packages/pinia-colada/tests/__shared__/utils.ts b/packages/pinia-colada/tests/__shared__/utils.ts new file mode 100644 index 000000000..8411f8cb3 --- /dev/null +++ b/packages/pinia-colada/tests/__shared__/utils.ts @@ -0,0 +1,72 @@ +import type { UseQueryOptions } from '@pinia/colada' +import { PiniaColada, useQuery, useQueryCache } from '@pinia/colada' +import { mount } from '@vue/test-utils' +import { createPinia } from 'pinia' +import { createApp, defineComponent } from 'vue' + +/** + * Mounts a component that runs a real Pinia Colada `useQuery`, + * so tests can observe exactly what users see. + */ +export function mountQuery(options: UseQueryOptions) { + return mount( + defineComponent({ + setup() { + const queryCache = useQueryCache() + const query = useQuery(options) + + return { query, queryCache } + }, + render: () => null, + }), + { + global: { + plugins: [createPinia(), PiniaColada], + }, + }, + ) +} + +export function createQueryCache() { + const app = createApp({ render: () => null }) + .use(createPinia()) + .use(PiniaColada) + return app.runWithContext(() => useQueryCache()) +} + +/** + * An AsyncIterable whose chunks are pushed manually, so tests can assert + * intermediate states deterministically while the stream is still open. + */ +export function createChunkController() { + const buffer: ({ value: T } | { done: true })[] = [] + let notify: (() => void) | undefined + + return { + async *stream() { + while (true) { + while (buffer.length === 0) { + await new Promise((resolve) => { + notify = resolve + }) + } + + const item = buffer.shift()! + + if ('done' in item) { + return + } + + yield item.value + } + }, + push(value: T) { + buffer.push({ value }) + notify?.() + }, + close() { + buffer.push({ done: true }) + notify?.() + }, + } +} diff --git a/packages/pinia-colada/tests/ssr.test.tsx b/packages/pinia-colada/tests/ssr.test.tsx new file mode 100644 index 000000000..0003f078c --- /dev/null +++ b/packages/pinia-colada/tests/ssr.test.tsx @@ -0,0 +1,137 @@ +import type { Client } from '@orpc/client' +import { + hydrateQueryCache, + PiniaColada, + serializeQueryCache, + useQuery, + useQueryCache, +} from '@pinia/colada' +import { mount } from '@vue/test-utils' +import { createPinia } from 'pinia' +import { createSSRApp, defineComponent, h, nextTick, ref, watch } from 'vue' +import { renderToString } from 'vue/server-renderer' + +import type { RouterUtils } from '../src/index.ts' +import { createRouterUtils, StreamQuerySSRUtilsPlugin } from '../src/index.ts' +import { createChunkController } from './__shared__/utils.ts' + +type StreamClient = { stream: Client, Error> } + +describe.each([ + [ + 'streamed', + (utils: RouterUtils) => utils.stream.streamedOptions() as any, + [], + [1, 2], + [[1], [1, 2]], + ], + [ + 'live', + (utils: RouterUtils) => utils.stream.liveOptions() as any, + 1, + [2, 3], + [2, 3], + ], +] as const)('ssr of %s queries', (_, options, serverData, events, browserData) => { + const started = vi.fn() + const server = createChunkController() + const browser = createChunkController() + + async function renderOnServer() { + const client = { + stream: vi.fn(() => { + started() + return server.stream() + }), + } as unknown as StreamClient + const utils = createRouterUtils(client, { + plugins: [new StreamQuerySSRUtilsPlugin({ isServer: true })], + }) + const pinia = createPinia() + + const app = createSSRApp( + defineComponent({ + setup() { + const query = useQuery(options(utils)) + return () => h('div', JSON.stringify(query.data.value)) + }, + }), + ) + .use(pinia) + .use(PiniaColada) + + server.push(1) + const html = await renderToString(app) + + return { html, state: JSON.parse(JSON.stringify(serializeQueryCache(useQueryCache(pinia)))) } + } + + it('reopens the stream on mount without going back to pending, and later components join it', async () => { + const { html, state } = await renderOnServer() + expect(html).toBe(`
${JSON.stringify(serverData)}
`) + expect(started).toHaveBeenCalledTimes(1) + + const client = { stream: vi.fn(() => browser.stream()) } as unknown as StreamClient + const utils = createRouterUtils(client, { + plugins: [new StreamQuerySSRUtilsPlugin({ isServer: false })], + }) + + const statuses: string[] = [] + const queries: any[] = [] + const Consumer = defineComponent({ + setup() { + const query = useQuery(options(utils)) + queries.push(query) + watch( + () => query.status.value, + (status) => statuses.push(status), + { flush: 'sync', immediate: true }, + ) + return () => null + }, + }) + + const showSecond = ref(false) + const mounted = mount( + defineComponent({ + setup: () => () => [h(Consumer), showSecond.value ? h(Consumer) : null], + }), + { + global: { + plugins: [ + createPinia(), + PiniaColada, + { + install: (app) => + hydrateQueryCache( + app.runWithContext(() => useQueryCache()), + state, + ), + }, + ], + }, + }, + ) + + expect(queries[0].data.value).toEqual(serverData) + await vi.waitFor(() => expect(client.stream).toHaveBeenCalledTimes(1), { interval: 1 }) + expect(queries[0].asyncStatus.value).toBe('loading') + + browser.push(events[0]) + await vi.waitFor(() => expect(queries[0].data.value).toEqual(browserData[0]), { interval: 1 }) + + showSecond.value = true + await nextTick() + expect(queries[1].data.value).toEqual(browserData[0]) + + browser.push(events[1]) + await vi.waitFor(() => expect(queries[1].data.value).toEqual(browserData[1]), { interval: 1 }) + expect(queries[0].data.value).toEqual(browserData[1]) + expect(queries[0].asyncStatus.value).toBe('loading') + + expect(client.stream).toHaveBeenCalledTimes(1) + expect(statuses).not.toContain('pending') + + mounted.unmount() + }) +}) diff --git a/packages/shared/src/array.test.ts b/packages/shared/src/array.test.ts index 3cc467525..a18b8c811 100644 --- a/packages/shared/src/array.test.ts +++ b/packages/shared/src/array.test.ts @@ -1,4 +1,4 @@ -import { splitInHalf, toArray } from './array.ts' +import { limitArraySize, splitInHalf, toArray } from './array.ts' it('toArray', () => { expect(toArray(undefined)).toEqual([]) @@ -21,3 +21,12 @@ it('splitInHalf', () => { expect(splitInHalf([1])).toEqual([[1], []]) expect(splitInHalf([])).toEqual([[], []]) }) + +it('limitArraySize', () => { + const items = [1, 2, 3] + expect(limitArraySize(items, 3)).toBe(items) + expect(limitArraySize(items, Number.POSITIVE_INFINITY)).toBe(items) + expect(limitArraySize(items, 2)).toEqual([2, 3]) + expect(limitArraySize(items, 0)).toEqual([]) + expect(items).toEqual([1, 2, 3]) +}) diff --git a/packages/shared/src/array.ts b/packages/shared/src/array.ts index ba37808f0..d6d3bd056 100644 --- a/packages/shared/src/array.ts +++ b/packages/shared/src/array.ts @@ -10,3 +10,13 @@ export function splitInHalf(arr: readonly T[]): [T[], T[]] { const half = Math.ceil(arr.length / 2) return [arr.slice(0, half), arr.slice(half)] } + +/** + * Keeps only the last `maxSize` items, returning `items` itself when it already fits. + */ +export function limitArraySize(items: T[], maxSize: number): T[] { + if (items.length <= maxSize) { + return items + } + return items.slice(items.length - maxSize) +} diff --git a/packages/shared/src/index.test.ts b/packages/shared/src/index.test.ts index 4e4364feb..806ba6e85 100644 --- a/packages/shared/src/index.test.ts +++ b/packages/shared/src/index.test.ts @@ -1,5 +1,10 @@ it('exports somethings', async () => { await expect(import('./index.ts')).resolves.toMatchObject({ allAbortSignal: expect.any(Function), + limitArraySize: expect.any(Function), + pipeStreamToCache: expect.any(Function), + runWithOwnSignal: expect.any(Function), + streamedArrayKeyOptions: expect.any(Function), + streamedArrayWrites: expect.any(Function), }) }) diff --git a/packages/shared/src/index.ts b/packages/shared/src/index.ts index 56ba4dc19..399c04dc5 100644 --- a/packages/shared/src/index.ts +++ b/packages/shared/src/index.ts @@ -17,6 +17,7 @@ export * from './proxy.ts' export * from './queue.ts' export * from './signal.ts' export * from './stream.ts' +export * from './stream-to-cache.ts' export * from './tracing.ts' export * from './types.ts' export * from './value.ts' diff --git a/packages/shared/src/stream-to-cache.test.ts b/packages/shared/src/stream-to-cache.test.ts new file mode 100644 index 000000000..9204600e0 --- /dev/null +++ b/packages/shared/src/stream-to-cache.test.ts @@ -0,0 +1,490 @@ +import { sleep } from '@standard-server/shared' + +import { promiseWithResolvers } from './promise.ts' +import { + pipeStreamToCache, + runWithOwnSignal, + streamedArrayKeyOptions, + streamedArrayWrites, +} from './stream-to-cache.ts' + +function createStore(initial?: number[]) { + const store = { + value: initial, + writes: [] as (number[] | undefined)[], + update: (updater: (current: number[] | undefined) => number[]): number[] => { + store.value = updater(store.value) + store.writes.push(store.value) + return store.value + }, + } + return store +} + +const append = (current: number[] | undefined, chunk: number): number[] => [ + ...(current ?? []), + chunk, +] + +describe('pipeStreamToCache', () => { + it('writes once at connect, before the first next(), then once per chunk', async () => { + const log: string[] = [] + const store = createStore([0]) + + async function* gen() { + log.push('start') + yield 1 + yield 2 + } + + const onConnect = vi.fn((current: number[] | undefined) => { + log.push('connect') + return current ?? [] + }) + + await expect( + pipeStreamToCache(gen(), { + signal: new AbortController().signal, + update: store.update, + onConnect, + onChunk: append, + }), + ).resolves.toEqual({ data: [0, 1, 2] }) + + expect(onConnect).toHaveBeenCalledTimes(1) + expect(onConnect).toHaveBeenCalledWith([0]) + expect(log).toEqual(['connect', 'start']) + expect(store.writes).toEqual([[0], [0, 1], [0, 1, 2]]) + }) + + it('writes nothing before the first chunk without onConnect', async () => { + const store = createStore([0]) + const gate = promiseWithResolvers() + + async function* gen() { + await gate.promise + yield 1 + } + + const promise = pipeStreamToCache(gen(), { + signal: new AbortController().signal, + update: store.update, + onChunk: append, + }) + + await sleep(10) + expect(store.writes).toEqual([]) + + gate.resolve() + await expect(promise).resolves.toEqual({ data: [0, 1] }) + expect(store.writes).toEqual([[0, 1]]) + }) + + it('resolves with undefined when nothing was written', async () => { + const store = createStore([0]) + + async function* gen(): AsyncGenerator {} + + await expect( + pipeStreamToCache(gen(), { + signal: new AbortController().signal, + update: store.update, + onChunk: append, + }), + ).resolves.toBeUndefined() + + expect(store.value).toEqual([0]) + }) + + it('reads the cache live, so external writes between chunks are kept', async () => { + const store = createStore() + + async function* gen() { + yield 1 + store.value = [...store.value!, 99] + yield 2 + } + + await expect( + pipeStreamToCache(gen(), { + signal: new AbortController().signal, + update: store.update, + onConnect: () => [], + onChunk: append, + }), + ).resolves.toEqual({ data: [1, 99, 2] }) + }) + + it('closes the stream without starting it when aborted before the call', async () => { + const store = createStore([0]) + const controller = new AbortController() + const reason = new Error('aborted') + controller.abort(reason) + + let started = false + async function* gen() { + started = true + yield 1 + } + const stream = gen() + const returnSpy = vi.spyOn(stream, 'return') + const onConnect = vi.fn(() => []) + + await expect( + pipeStreamToCache(stream, { + signal: controller.signal, + update: store.update, + onConnect, + onChunk: append, + }), + ).rejects.toBe(reason) + + expect(returnSpy).toHaveBeenCalledTimes(1) + expect(started).toBe(false) + expect(onConnect).not.toHaveBeenCalled() + expect(store.writes).toEqual([]) + expect(store.value).toEqual([0]) + }) + + it('closes the stream and rethrows when onConnect throws', async () => { + const error = new Error('connect error') + async function* gen() { + yield 1 + } + const stream = gen() + const returnSpy = vi.spyOn(stream, 'return') + + await expect( + pipeStreamToCache(stream, { + signal: new AbortController().signal, + update: createStore().update, + onConnect: () => { + throw error + }, + onChunk: append, + }), + ).rejects.toBe(error) + + expect(returnSpy).toHaveBeenCalledTimes(1) + }) + + it('stops writing and closes the stream when aborted while awaiting next()', async () => { + const store = createStore() + const controller = new AbortController() + const reason = new Error('aborted') + const gate = promiseWithResolvers() + const cleanup = vi.fn() + + async function* gen() { + try { + yield 1 + await gate.promise + yield 2 + yield 3 + } finally { + cleanup() + } + } + + const promise = pipeStreamToCache(gen(), { + signal: controller.signal, + update: store.update, + onChunk: append, + }) + + await vi.waitFor(() => expect(store.writes).toEqual([[1]])) + controller.abort(reason) + gate.resolve() + + await expect(promise).rejects.toBe(reason) + expect(store.writes).toEqual([[1]]) + expect(cleanup).toHaveBeenCalledTimes(1) + }) + + it('stops at once when update aborts the signal', async () => { + const store = createStore() + const controller = new AbortController() + const reason = new Error('aborted') + const afterFirst = vi.fn() + const cleanup = vi.fn() + + async function* gen() { + try { + yield 1 + afterFirst() + yield 2 + } finally { + cleanup() + } + } + + await expect( + pipeStreamToCache(gen(), { + signal: controller.signal, + update: (updater) => { + const data = store.update(updater) + controller.abort(reason) + return data + }, + onChunk: append, + }), + ).rejects.toBe(reason) + + expect(afterFirst).not.toHaveBeenCalled() + expect(cleanup).toHaveBeenCalledTimes(1) + expect(store.writes).toEqual([[1]]) + }) + + it('settles at the first chunk without waiting for the next one', async () => { + const store = createStore() + const cleanup = vi.fn() + + async function* gen() { + try { + yield 1 + await new Promise(() => {}) + yield 2 + } finally { + cleanup() + } + } + + await expect( + pipeStreamToCache(gen(), { + signal: new AbortController().signal, + update: store.update, + onChunk: append, + settle: () => true, + }), + ).resolves.toEqual({ data: [1] }) + + expect(cleanup).toHaveBeenCalledTimes(1) + expect(store.writes).toEqual([[1]]) + }) + + it('settles at connect without starting the stream', async () => { + const store = createStore([0]) + let started = false + + async function* gen() { + started = true + yield 1 + } + + await expect( + pipeStreamToCache(gen(), { + signal: new AbortController().signal, + update: store.update, + onConnect: () => [], + onChunk: append, + settle: () => true, + }), + ).resolves.toEqual({ data: [] }) + + expect(started).toBe(false) + expect(store.writes).toEqual([[]]) + }) + + it('rejects and keeps the first write when settle aborts the signal', async () => { + const store = createStore() + const controller = new AbortController() + const reason = new Error('aborted') + const cleanup = vi.fn() + + async function* gen() { + try { + yield 1 + yield 2 + } finally { + cleanup() + } + } + + await expect( + pipeStreamToCache(gen(), { + signal: controller.signal, + update: store.update, + onChunk: append, + settle: () => { + controller.abort(reason) + return false + }, + }), + ).rejects.toBe(reason) + + expect(cleanup).toHaveBeenCalledTimes(1) + expect(store.value).toEqual([1]) + }) + + it('calls settle once and follows the stream to its end when it returns false', async () => { + const store = createStore() + const settle = vi.fn(() => false) + + async function* gen() { + yield 1 + yield 2 + } + + await expect( + pipeStreamToCache(gen(), { + signal: new AbortController().signal, + update: store.update, + onConnect: () => [], + onChunk: append, + settle, + }), + ).resolves.toEqual({ data: [1, 2] }) + + expect(settle).toHaveBeenCalledTimes(1) + expect(store.writes).toEqual([[], [1], [1, 2]]) + }) + + it('rejects on a stream error before the first write and keeps existing data', async () => { + const store = createStore([0]) + const error = new Error('stream error') + + async function* gen(): AsyncGenerator { + throw error + } + + await expect( + pipeStreamToCache(gen(), { + signal: new AbortController().signal, + update: store.update, + onChunk: append, + }), + ).rejects.toBe(error) + + expect(store.writes).toEqual([]) + expect(store.value).toEqual([0]) + }) + + it('rejects on a stream error after the first write and keeps the written data', async () => { + const store = createStore() + const error = new Error('stream error') + + async function* gen() { + yield 1 + throw error + } + + await expect( + pipeStreamToCache(gen(), { + signal: new AbortController().signal, + update: store.update, + onConnect: () => [], + onChunk: append, + }), + ).rejects.toBe(error) + + expect(store.value).toEqual([1]) + }) + + it('ignores a stream that fails to close', async () => { + const store = createStore() + const stream = { + [Symbol.asyncIterator]: () => ({ + next: async () => ({ done: false, value: 1 }), + return: async () => { + throw new Error('close error') + }, + }), + } + + await expect( + pipeStreamToCache(stream, { + signal: new AbortController().signal, + update: store.update, + onChunk: append, + settle: () => true, + }), + ).resolves.toEqual({ data: [1] }) + }) + + it('rejects with the abort reason over a stream error', async () => { + const controller = new AbortController() + const reason = new Error('aborted') + const stream = { + [Symbol.asyncIterator]: () => ({ + next: async () => { + controller.abort(reason) + throw new Error('stream error') + }, + }), + } + + await expect( + pipeStreamToCache(stream, { + signal: controller.signal, + update: createStore().update, + onChunk: append, + }), + ).rejects.toBe(reason) + }) +}) + +describe('runWithOwnSignal', () => { + it('aborts its own signal once fn ends', async () => { + const parent = new AbortController() + let signal: AbortSignal | undefined + + await expect( + runWithOwnSignal(parent.signal, async (s) => { + signal = s + expect(s.aborted).toBe(false) + return 1 + }), + ).resolves.toBe(1) + + expect(signal!.aborted).toBe(true) + expect(parent.signal.aborted).toBe(false) + }) + + it('follows the parent signal', async () => { + const parent = new AbortController() + const reason = new Error('aborted') + + await runWithOwnSignal(parent.signal, async (signal) => { + parent.abort(reason) + expect(signal.reason).toBe(reason) + }) + + const aborted = AbortSignal.abort(reason) + await runWithOwnSignal(aborted, async (signal) => expect(signal.reason).toBe(reason)) + }) +}) + +describe('streamedArrayWrites', () => { + it.each([ + ['reset', undefined, [], [1]], + ['reset', [7, 8, 9], [], [1]], + ['append', undefined, [], [1]], + ['append', [7, 8, 9], [8, 9], [9, 1]], + ] as const)('%s over %j', (refetchMode, existing, connected, afterChunk) => { + const { onConnect, onChunk } = streamedArrayWrites({ refetchMode, maxChunks: 2 }) + + expect(onConnect(existing && [...existing])).toEqual(connected) + expect(onChunk(onConnect(existing && [...existing]), 1)).toEqual(afterChunk) + }) + + it('defaults to reset with unlimited chunks', () => { + const { onConnect, onChunk } = streamedArrayWrites() + + expect(onConnect([1, 2])).toEqual([]) + expect(onChunk([1, 2, 3], 4)).toEqual([1, 2, 3, 4]) + }) +}) + +it('streamedArrayKeyOptions keeps only non-default options', () => { + expect(streamedArrayKeyOptions({ refetchMode: 'append', maxChunks: 3 })).toEqual({ + refetchMode: 'append', + maxChunks: 3, + }) + expect(streamedArrayKeyOptions({ refetchMode: 'append' })).toEqual({ refetchMode: 'append' }) + expect(streamedArrayKeyOptions({ refetchMode: 'reset', maxChunks: 3 })).toEqual({ maxChunks: 3 }) + expect( + streamedArrayKeyOptions({ refetchMode: 'reset', maxChunks: Number.POSITIVE_INFINITY }), + ).toBeUndefined() + expect(streamedArrayKeyOptions({})).toBeUndefined() + expect(streamedArrayKeyOptions(undefined)).toBeUndefined() +}) diff --git a/packages/shared/src/stream-to-cache.ts b/packages/shared/src/stream-to-cache.ts new file mode 100644 index 000000000..fee85cb7f --- /dev/null +++ b/packages/shared/src/stream-to-cache.ts @@ -0,0 +1,151 @@ +import { throwIfAborted } from '@standard-server/shared' + +import { limitArraySize } from './array.ts' + +export interface PipeStreamToCacheOptions { + /** Once aborted, nothing more is written; the stream is closed and `signal.reason` is thrown. */ + signal: AbortSignal + /** Applies `updater` to the cached value synchronously. */ + update: (updater: (current: TData | undefined) => TData) => void + /** One write made when the stream connects, before the first chunk. Omit it to write nothing at connect. */ + onConnect?: ((current: TData | undefined) => TData) | undefined + onChunk: (current: TData | undefined, chunk: TChunk) => TData + /** Called right after the first write only. `true` closes the stream and resolves with the data so far. */ + settle?: (() => boolean) | undefined +} + +/** + * The shared core of streamed and live queries: writes nothing before the stream connects, then each write reads the cache live. + * @returns the last written data, or `undefined` if nothing was written + */ +export async function pipeStreamToCache( + stream: AsyncIterable, + { signal, update, onConnect, onChunk, settle }: PipeStreamToCacheOptions, +): Promise<{ data: TData } | undefined> { + let last: { data: TData } | undefined + + // true = stop + const write = (updater: (current: TData | undefined) => TData): boolean => { + const isFirst = last === undefined + update((current) => { + const data = updater(current) + last = { data } + return data + }) + return (isFirst && settle?.() === true) || signal.aborted + } + + const iterator = stream[Symbol.asyncIterator]() + // a stream that fails to close must not hide the result or the original error + const close = async (): Promise => { + try { + await iterator.return?.() + } catch {} + } + + try { + if (signal.aborted || (onConnect !== undefined && write(onConnect))) { + await close() + } else { + while (true) { + const result = await iterator.next() + + if (result.done) { + break + } + + if (signal.aborted || write((current) => onChunk(current, result.value))) { + await close() + break + } + } + } + } catch (error) { + await close() + throwIfAborted(signal) + throw error + } + + throwIfAborted(signal) + return last +} + +export interface StreamedArrayOptions { + /** + * What a fetch does with data already in the cache once its stream connects. + * Until then the data stays, so a query with data never returns to `pending`. + * + * - `'reset'`: replaced by `[]`, then each chunk is appended. + * - `'append'`: kept (trimmed to `maxChunks`), and each chunk is appended. + * + * @default 'reset' + */ + refetchMode?: 'reset' | 'append' | undefined + + /** + * Limits the number of data chunks stored in the query result. + * Older chunks are removed when the limit is reached. + * + * @default Number.POSITIVE_INFINITY (unlimited) + */ + maxChunks?: number | undefined +} + +/** + * The `onConnect` and `onChunk` writes of a streamed query, which collects chunks into an array. + */ +export function streamedArrayWrites({ + refetchMode = 'reset', + maxChunks = Number.POSITIVE_INFINITY, +}: StreamedArrayOptions = {}): Required< + Pick, 'onConnect' | 'onChunk'> +> { + return { + onConnect: (current) => + refetchMode === 'append' ? limitArraySize(current ?? [], maxChunks) : [], + onChunk: (current, chunk) => limitArraySize([...(current ?? []), chunk], maxChunks), + } +} + +/** + * The part of `StreamedArrayOptions` that belongs in the query key, so queries with different options get separate entries. + * Defaults are left out, so default options produce the same key as no options. + */ +export function streamedArrayKeyOptions( + options: StreamedArrayOptions | undefined, +): StreamedArrayOptions | undefined { + const { refetchMode = 'reset', maxChunks = Number.POSITIVE_INFINITY } = options ?? {} + + const keyOptions: StreamedArrayOptions = { + ...(refetchMode !== 'reset' ? { refetchMode } : {}), + ...(maxChunks !== Number.POSITIVE_INFINITY ? { maxChunks } : {}), + } + + return Object.keys(keyOptions).length > 0 ? keyOptions : undefined +} + +/** + * Runs `fn` with a signal that follows `signal` and is also aborted once `fn` ends, + * so the call that opened a stream is told to stop as soon as the stream is no longer read. + */ +export async function runWithOwnSignal( + signal: AbortSignal, + fn: (signal: AbortSignal) => Promise, +): Promise { + const controller = new AbortController() + + if (signal.aborted) { + controller.abort(signal.reason) + } else { + signal.addEventListener('abort', () => controller.abort(signal.reason), { + once: true, + signal: controller.signal, + }) + } + + try { + return await fn(controller.signal) + } finally { + controller.abort() + } +} diff --git a/packages/swr/src/procedure-utils.test-d.ts b/packages/swr/src/procedure-utils.test-d.ts index e7eb6cb29..ba6cc3ae4 100644 --- a/packages/swr/src/procedure-utils.test-d.ts +++ b/packages/swr/src/procedure-utils.test-d.ts @@ -146,6 +146,13 @@ describe('procedureUtils', () => { requiredUtils.subscriber({ context: { batch: 'invalid' } }) }) + it('should infer types for `refetchMode` correctly', () => { + optionalUtils.subscriber({ refetchMode: 'reset' }) + optionalUtils.subscriber({ refetchMode: 'append' }) + // @ts-expect-error - 'replace' is removed + optionalUtils.subscriber({ refetchMode: 'replace' }) + }) + it('should infer types for `output` and `error` correctly', () => { const subscription = useSWRSubscription( streamUtils.key(), diff --git a/packages/swr/src/procedure-utils.test.ts b/packages/swr/src/procedure-utils.test.ts index 33d1a39ad..7d3e9feca 100644 --- a/packages/swr/src/procedure-utils.test.ts +++ b/packages/swr/src/procedure-utils.test.ts @@ -70,14 +70,17 @@ describe('procedureUtils', () => { }) describe('.subscriber', async () => { - it('on success', async () => { + it.each([ + ['reset', []], + ['append', ['2', '3']], + ] as const)('on success refetchMode=%s', async (refetchMode, connected) => { client.mockImplementationOnce(async function* () { await new Promise((resolve) => setTimeout(resolve, 10)) yield '__event__1' yield '__event__2' yield '__event__3' }) - const subscriber = utils.subscriber({ context: { batch: true }, maxChunks: 2 }) + const subscriber = utils.subscriber({ context: { batch: true }, maxChunks: 2, refetchMode }) const next = vi.fn() const unsubscribe = subscriber(key, { next }) @@ -98,108 +101,78 @@ describe('procedureUtils', () => { expect(next).toHaveBeenCalledTimes(4) expect(next).toHaveBeenCalledWith(undefined, expect.any(Function)) - expect(next).toHaveBeenNthCalledWith(1, undefined, undefined) // reset mode + expect(next.mock.calls[0]![1](undefined)).toEqual([]) + expect(next.mock.calls[0]![1](['1', '2', '3'])).toEqual(connected) expect(next.mock.calls[1]![1](undefined)).toEqual(['__event__1']) expect(next.mock.calls[2]![1](['1'])).toEqual(['1', '__event__2']) // exceeds maxChunks, so it should only return the last 2 events expect(next.mock.calls[3]![1](['1', '2'])).toEqual(['2', '__event__3']) }) - it('on success refetchMode=append', async () => { - client.mockImplementationOnce(async function* () { - await new Promise((resolve) => setTimeout(resolve, 10)) - yield '__event__1' - yield '__event__2' - yield '__event__3' - }) - const subscriber = utils.subscriber({ - context: { batch: true }, - maxChunks: 2, - refetchMode: 'append', - }) - - const next = vi.fn() - const unsubscribe = subscriber(key, { next }) - - expect(unsubscribe).toBeInstanceOf(Function) - - await new Promise((resolve) => setTimeout(resolve, 20)) - - expect(client).toHaveBeenCalledTimes(1) - expect(client).toHaveBeenCalledWith( - { search: '__search__' }, - { - context: { batch: true, [OPERATION_CONTEXT_SYMBOL]: { key, type: 'subscriber' } }, - signal: expect.any(AbortSignal), - }, - ) - - expect(next).toHaveBeenCalledTimes(3) - expect(next).toHaveBeenCalledWith(undefined, expect.any(Function)) - - expect(next.mock.calls[0]![1](undefined)).toEqual(['__event__1']) - expect(next.mock.calls[1]![1](['1'])).toEqual(['1', '__event__2']) - // exceeds maxChunks, so it should only return the last 2 events - expect(next.mock.calls[2]![1](['1', '2'])).toEqual(['2', '__event__3']) - }) - - it('on success refetchMode=replace with previous data', async () => { - client.mockImplementationOnce(async function* () { + it.each(['reset', 'append'] as const)( + 'keeps existing data until connect, then appends to the live data (%s)', + async (refetchMode) => { + const connected = promiseWithResolvers() + const paused = promiseWithResolvers() + client.mockImplementationOnce(async () => { + await connected.promise + return (async function* () { + yield '__event__1' + await paused.promise + yield '__event__2' + })() + }) + const subscriber = utils.subscriber({ refetchMode }) + + // simulate SWR: updaters are invoked synchronously with the current data + let data: unknown = ['__previous__'] + const next = vi.fn((error, update?) => { + if (error === undefined) { + data = typeof update === 'function' ? update(data) : update + } + }) + + subscriber(key, { next }) await new Promise((resolve) => setTimeout(resolve, 10)) - yield '__event__1' - yield '__event__2' - yield '__event__3' - }) - const subscriber = utils.subscriber({ - context: { batch: true }, - maxChunks: 2, - refetchMode: 'replace', - }) - // simulate SWR: updaters are invoked synchronously with the current data - let data: unknown = ['__previous__'] - const next = vi.fn((error, update?) => { - if (error === undefined) { - data = typeof update === 'function' ? update(data) : update - } - }) + expect(next).not.toHaveBeenCalled() + expect(data).toEqual(['__previous__']) - subscriber(key, { next }) + connected.resolve() + await vi.waitFor(() => expect(next).toHaveBeenCalledTimes(2), { interval: 1 }) // connect + first event - await new Promise((resolve) => setTimeout(resolve, 20)) + const kept = refetchMode === 'append' ? ['__previous__'] : [] + expect(data).toEqual([...kept, '__event__1']) - // one probe call + one final replace call, no updates during the stream - expect(next).toHaveBeenCalledTimes(2) - expect(next).toHaveBeenNthCalledWith(1, undefined, expect.any(Function)) - // exceeds maxChunks, so it should only keep the last 2 events - expect(next).toHaveBeenNthCalledWith(2, undefined, ['__event__2', '__event__3']) - expect(data).toEqual(['__event__2', '__event__3']) - }) + data = [...(data as unknown[]), '__external__'] + paused.resolve() + await vi.waitFor(() => expect(next).toHaveBeenCalledTimes(3), { interval: 1 }) - it('on success refetchMode=replace without previous data', async () => { - client.mockImplementationOnce(async function* () { - await new Promise((resolve) => setTimeout(resolve, 10)) - yield '__event__1' - yield '__event__2' - yield '__event__3' - }) - const subscriber = utils.subscriber({ context: { batch: true }, refetchMode: 'replace' }) + expect(data).toEqual([...kept, '__event__1', '__external__', '__event__2']) + }, + ) - // simulate SWR: updaters are invoked synchronously with the current data - let data: unknown - const next = vi.fn((error, update?) => { - if (error === undefined) { - data = typeof update === 'function' ? update(data) : update - } + it('does not start the stream when unsubscribed before the call resolves', async () => { + const started = vi.fn() + const connected = promiseWithResolvers() + client.mockImplementationOnce(async () => { + await connected.promise + return (async function* () { + started() + yield '__event__1' + })() }) + const subscriber = utils.subscriber() - subscriber(key, { next }) + const next = vi.fn() + const unsubscribe = subscriber(key, { next }) - await new Promise((resolve) => setTimeout(resolve, 20)) + unsubscribe() + connected.resolve() + await new Promise((resolve) => setTimeout(resolve, 10)) - // one probe call + one update per event, streamed live like append mode - expect(next).toHaveBeenCalledTimes(4) - expect(data).toEqual(['__event__1', '__event__2', '__event__3']) + expect(started).not.toHaveBeenCalled() + expect(next).not.toHaveBeenCalled() }) it('on unsubscribe', async () => { @@ -272,7 +245,7 @@ describe('procedureUtils', () => { const next = vi.fn() const unsubscribe = subscriber(key, { next }) - await vi.waitFor(() => expect(next).toHaveBeenCalledTimes(2)) // reset + first event + await vi.waitFor(() => expect(next).toHaveBeenCalledTimes(2), { interval: 1 }) // reset + first event unsubscribe() paused.resolve() await new Promise((resolve) => setTimeout(resolve, 10)) @@ -281,34 +254,6 @@ describe('procedureUtils', () => { expect(cleanup).toHaveBeenCalledTimes(1) }) - it('does not replace data after unsubscribe in refetchMode=replace', async () => { - const paused = promiseWithResolvers() - client.mockImplementationOnce(async function* () { - yield '__event__1' - await paused.promise - yield '__event__2' - }) - const subscriber = utils.subscriber({ refetchMode: 'replace' }) - - // simulate SWR: updaters are invoked synchronously with the current data - let data: unknown = ['__previous__'] - const next = vi.fn((error, update?) => { - if (error === undefined) { - data = typeof update === 'function' ? update(data) : update - } - }) - - const unsubscribe = subscriber(key, { next }) - await vi.waitFor(() => expect(next).toHaveBeenCalledTimes(1)) // the probe - await new Promise((resolve) => setTimeout(resolve, 10)) // the first event is buffered - unsubscribe() - paused.resolve() - await new Promise((resolve) => setTimeout(resolve, 10)) - - expect(next).toHaveBeenCalledTimes(1) - expect(data).toEqual(['__previous__']) - }) - it('on error while yielding', async () => { client.mockImplementationOnce(async function* () { await new Promise((resolve) => setTimeout(resolve, 10)) @@ -384,9 +329,9 @@ describe('procedureUtils', () => { ) expect(next).toHaveBeenCalledTimes(3) - expect(next).toHaveBeenNthCalledWith(1, undefined, '__event__1') - expect(next).toHaveBeenNthCalledWith(2, undefined, '__event__2') - expect(next).toHaveBeenNthCalledWith(3, undefined, '__event__3') + expect(next.mock.calls[0]![1]).toEqual('__event__1') + expect(next.mock.calls[1]![1]).toEqual('__event__2') + expect(next.mock.calls[2]![1]).toEqual('__event__3') }) it('on unsubscribe', async () => { @@ -447,13 +392,13 @@ describe('procedureUtils', () => { const next = vi.fn() const unsubscribe = subscriber(key, { next }) - await vi.waitFor(() => expect(next).toHaveBeenCalledTimes(1)) + await vi.waitFor(() => expect(next).toHaveBeenCalledTimes(1), { interval: 1 }) unsubscribe() paused.resolve() await new Promise((resolve) => setTimeout(resolve, 10)) expect(next).toHaveBeenCalledTimes(1) - expect(next).toHaveBeenCalledWith(undefined, '__event__1') + expect(next.mock.calls[0]![1]).toEqual('__event__1') }) it('on error while yielding', async () => { diff --git a/packages/swr/src/procedure-utils.ts b/packages/swr/src/procedure-utils.ts index 483374f6e..0e72f3ad8 100644 --- a/packages/swr/src/procedure-utils.ts +++ b/packages/swr/src/procedure-utils.ts @@ -1,6 +1,12 @@ import type { Client, ClientContext } from '@orpc/client' -import type { MaybeOptionalOptions } from '@orpc/shared' -import { isAsyncIteratorObject, resolveMaybeOptionalOptions } from '@orpc/shared' +import type { MaybeOptionalOptions, PipeStreamToCacheOptions } from '@orpc/shared' +import { + isAsyncIteratorObject, + pipeStreamToCache, + resolveMaybeOptionalOptions, + streamedArrayWrites, +} from '@orpc/shared' +import type { SWRSubscriptionOptions } from 'swr/subscription' import type { OperationKey, OperationKeyOptions, OperationKeyPrefixOptions } from './key.ts' import { generateOperationKey } from './key.ts' @@ -89,95 +95,13 @@ export class ProcedureUtils< subscriber( ...rest: MaybeOptionalOptions> ): Subscriber, TError> { - const { maxChunks, refetchMode = 'reset', ...options } = resolveMaybeOptionalOptions(rest) + const { maxChunks, refetchMode, ...options } = resolveMaybeOptionalOptions(rest) - return (key, { next }) => { - const controller = new AbortController() - - void (async () => { - try { - /** - * Input can be omitted from the key when it is undefined, - * so we can safely cast it to TInput with an undefined fallback. - */ - const { input } = key[key.length - 1] as { input: TInput } - - const iterator = await this.call(input, { - context: { - [OPERATION_CONTEXT_SYMBOL]: { key, type: 'subscriber' }, - ...options.context, - } satisfies OperationContext as any, - signal: controller.signal, - }) - - if (!isAsyncIteratorObject(iterator)) { - throw new TypeError('.subscriber requires an AsyncIteratorObject output') - } - - if (controller.signal.aborted) { - await iterator.return?.() - return - } - - let hasPreviousData = false - - if (refetchMode === 'reset') { - next(undefined, undefined) - } else if (refetchMode === 'replace') { - /** - * The updater is invoked synchronously with the current data, - * so we can use it to detect whether previous data exists. - */ - next(undefined, (old) => { - hasPreviousData = old !== undefined - return old as InferSubscriberOutput - }) - } - - const shouldUpdateDataDuringStream = refetchMode !== 'replace' || !hasPreviousData - let buffer: unknown[] = [] - - for await (const event of iterator) { - if (controller.signal.aborted) { - break - } - - if (shouldUpdateDataDuringStream) { - next(undefined, (old) => { - const newData = Array.isArray(old) ? [...old, event] : [event] - - if (typeof maxChunks === 'number' && newData.length > maxChunks) { - return newData.slice(newData.length - maxChunks) as InferSubscriberOutput - } - - return newData as InferSubscriberOutput - }) - } else { - buffer = [...buffer, event] - - if (typeof maxChunks === 'number' && buffer.length > maxChunks) { - buffer = buffer.slice(buffer.length - maxChunks) - } - } - } - - if (!shouldUpdateDataDuringStream && !controller.signal.aborted) { - next(undefined, buffer as InferSubscriberOutput) - } - } catch (error) { - /** - * Only report the error if the controller is not aborted. - */ - if (!controller.signal.aborted) { - next(error as TError) - } - } - })() - - return () => { - controller.abort() - } - } + return this.streamSubscriber('subscriber', options.context, (next) => ({ + update: (updater) => + next(undefined, (current) => updater(Array.isArray(current) ? current : undefined) as any), + ...streamedArrayWrites({ refetchMode, maxChunks }), + })) } /** @@ -190,6 +114,20 @@ export class ProcedureUtils< ): Subscriber, TError> { const options = resolveMaybeOptionalOptions(rest) + return this.streamSubscriber('liveSubscriber', options.context, (next) => ({ + // a live write ignores the current data, so pass the event itself rather than a mutator + update: (updater) => next(undefined, updater(undefined)), + onChunk: (_, chunk) => chunk, + })) + } + + private streamSubscriber( + type: 'subscriber' | 'liveSubscriber', + context: TClientContext | undefined, + pipeOptions: ( + next: SWRSubscriptionOptions['next'], + ) => Omit, 'signal'>, + ): Subscriber { return (key, { next }) => { const controller = new AbortController() @@ -203,28 +141,17 @@ export class ProcedureUtils< const iterator = await this.call(input, { context: { - [OPERATION_CONTEXT_SYMBOL]: { key, type: 'liveSubscriber' }, - ...options.context, + [OPERATION_CONTEXT_SYMBOL]: { key, type }, + ...context, } satisfies OperationContext as any, signal: controller.signal, }) if (!isAsyncIteratorObject(iterator)) { - throw new TypeError('.liveSubscriber requires an AsyncIteratorObject output') - } - - if (controller.signal.aborted) { - await iterator.return?.() - return + throw new TypeError(`.${type} requires an AsyncIteratorObject output`) } - for await (const event of iterator) { - if (controller.signal.aborted) { - break - } - - next(undefined, event as InferLiveSubscriberOutput) - } + await pipeStreamToCache(iterator, { signal: controller.signal, ...pipeOptions(next) }) } catch (error) { /** * Only report the error if the controller is not aborted. diff --git a/packages/swr/src/types.ts b/packages/swr/src/types.ts index d0c554f96..aa613e2c8 100644 --- a/packages/swr/src/types.ts +++ b/packages/swr/src/types.ts @@ -1,5 +1,5 @@ import type { ClientContext } from '@orpc/client' -import type { PartialDeep } from '@orpc/shared' +import type { PartialDeep, StreamedArrayOptions } from '@orpc/shared' import type { SWRSubscriptionOptions } from 'swr/subscription' import type { OperationKey } from './key.ts' @@ -51,24 +51,7 @@ export type FetcherOptions = object extend : { context: TClientContext } export type SubscriberOptions = - FetcherOptions & { - /** - * Determines how data is handled when the subscription is subscribed again - * - * - `'reset'`: Clears existing data before streaming new data. - * - `'append'`: Adds new streamed chunks to the existing data. - * - `'replace'`: Buffers streamed data and replaces existing data after the stream completes. - * - * @default 'reset' - */ - refetchMode?: 'append' | 'reset' | 'replace' | undefined - - /** - * Maximum number of chunks to store. - * When exceeded, the oldest chunks will be removed. - */ - maxChunks?: number | undefined - } + FetcherOptions & StreamedArrayOptions export type MutatorOptions = FetcherOptions diff --git a/packages/swr/tests/e2e.test.tsx b/packages/swr/tests/e2e.test.tsx index 5efb07e8a..1bdfbfaa4 100644 --- a/packages/swr/tests/e2e.test.tsx +++ b/packages/swr/tests/e2e.test.tsx @@ -1,5 +1,6 @@ import { act, renderHook } from '@testing-library/react' -import useSWR, { mutate } from 'swr' +import type { ReactNode } from 'react' +import useSWR, { mutate, SWRConfig } from 'swr' import useSWRInfinite from 'swr/infinite' import useSWRMutation from 'swr/mutation' import useSWRSubscription from 'swr/subscription' @@ -104,45 +105,59 @@ it('case: useSWRSubscription & .subscriber', async () => { }) }) -it('case: useSWRSubscription & .subscriber refetchMode=replace', async () => { - const first = renderHook(() => - useSWRSubscription( - orpc.stream.key({ input: { input: 3 } }), - orpc.stream.subscriber({ refetchMode: 'replace' }), - ), - ) - - await act(async () => { - await vi.waitFor(() => - expect(first.result.current.data).toEqual([ - { output: '0' }, - { output: '1' }, - { output: '2' }, - ]), +it.each([ + ['reset', [{ output: '__new__' }]], + ['append', [{ output: '0' }, { output: '1' }, { output: '2' }, { output: '__new__' }]], +] as const)( + 'case: useSWRSubscription & .subscriber refetchMode=%s', + async (refetchMode, expected) => { + // subscription data is not cleared by the global mutate, so each case gets its own cache + const cache = new Map() + const wrapper = ({ children }: { children: ReactNode }) => ( + cache }}>{children} ) - }) - - first.unmount() - streamHandler.mockImplementationOnce(async function* () { - yield { output: '__replaced__' } - }) + const first = renderHook( + () => + useSWRSubscription( + orpc.stream.key({ input: { input: 3 } }), + orpc.stream.subscriber({ refetchMode }), + ), + { wrapper }, + ) - const second = renderHook(() => - useSWRSubscription( - orpc.stream.key({ input: { input: 3 } }), - orpc.stream.subscriber({ refetchMode: 'replace' }), - ), - ) + await act(async () => { + await vi.waitFor(() => + expect(first.result.current.data).toEqual([ + { output: '0' }, + { output: '1' }, + { output: '2' }, + ]), + ) + }) + + first.unmount() + + streamHandler.mockImplementationOnce(async function* () { + yield { output: '__new__' } + }) + + const second = renderHook( + () => + useSWRSubscription( + orpc.stream.key({ input: { input: 3 } }), + orpc.stream.subscriber({ refetchMode }), + ), + { wrapper }, + ) - // previous data is preserved while the new stream is buffering - expect(second.result.current.data).toEqual([{ output: '0' }, { output: '1' }, { output: '2' }]) + expect(second.result.current.data).toEqual([{ output: '0' }, { output: '1' }, { output: '2' }]) - await act(async () => { - // append mode would end with 4 events here, replace keeps only the new stream's events - await vi.waitFor(() => expect(second.result.current.data).toEqual([{ output: '__replaced__' }])) - }) -}) + await act(async () => { + await vi.waitFor(() => expect(second.result.current.data).toEqual(expected)) + }) + }, +) it('case: useSWRSubscription & .liveSubscriber', async () => { const { result } = renderHook(() => { diff --git a/packages/tanstack-query/package.json b/packages/tanstack-query/package.json index 6822a83f8..a1dde949e 100644 --- a/packages/tanstack-query/package.json +++ b/packages/tanstack-query/package.json @@ -68,6 +68,6 @@ "zod": "^4.6.7" }, "peerDependencies": { - "@tanstack/query-core": ">=5.89.0" + "@tanstack/query-core": ">=5.105.0" } } diff --git a/packages/tanstack-query/src/index.ts b/packages/tanstack-query/src/index.ts index dbc7f4d0b..369fa288d 100644 --- a/packages/tanstack-query/src/index.ts +++ b/packages/tanstack-query/src/index.ts @@ -26,6 +26,7 @@ export * from './router-utils.ts' export { createRouterUtils as createTanstackQueryUtils } from './router-utils.ts' export { SharedUtils } from './shared-utils.ts' export * from './ssr.ts' +export type { SettleOnFirstDataOptions } from './stream-cache.ts' export type { SerializableStreamedQueryOptions } from './stream-query.ts' export * from './structural-sharing.ts' export * from './types.ts' diff --git a/packages/tanstack-query/src/key.test.ts b/packages/tanstack-query/src/key.test.ts index 6af5c7b1f..52e3afe3c 100644 --- a/packages/tanstack-query/src/key.test.ts +++ b/packages/tanstack-query/src/key.test.ts @@ -18,12 +18,18 @@ it('generateOperationKey', () => { generateOperationKey(['planet', 'stream'], { type: 'streamed', input: { cursor: 0 }, - fnOptions: { refetchMode: 'append' }, + fnOptions: { refetchMode: 'append', maxChunks: 3 }, }), ).toEqual([ ['planet', 'stream'], - { type: 'streamed', input: { cursor: 0 }, fnOptions: { refetchMode: 'append' } }, + { type: 'streamed', input: { cursor: 0 }, fnOptions: { refetchMode: 'append', maxChunks: 3 } }, ]) + expect( + generateOperationKey(['planet', 'stream'], { + type: 'streamed', + fnOptions: { refetchMode: 'reset', maxChunks: Number.POSITIVE_INFINITY }, + }), + ).toEqual([['planet', 'stream'], { type: 'streamed' }]) expect(generateOperationKey(['path'], { prefix: '__prefix__' })).toEqual([ '__prefix__', diff --git a/packages/tanstack-query/src/key.ts b/packages/tanstack-query/src/key.ts index 5d8ff6723..7210b3f19 100644 --- a/packages/tanstack-query/src/key.ts +++ b/packages/tanstack-query/src/key.ts @@ -1,6 +1,6 @@ import { RPCJsonSerializer } from '@orpc/client' import type { PartialDeep, Public } from '@orpc/shared' -import { stringifyJSON } from '@orpc/shared' +import { streamedArrayKeyOptions, stringifyJSON } from '@orpc/shared' import { skipToken } from '@tanstack/query-core' import type { SerializableStreamedQueryOptions } from './stream-query.ts' @@ -67,6 +67,8 @@ export function generateOperationKey( path: string[], options: OperationKeyOptions = {}, ): OperationKey { + const fnOptions = streamedArrayKeyOptions(options.fnOptions) + return [ ...(options.prefix !== undefined ? [options.prefix] : []), path, @@ -77,7 +79,7 @@ export function generateOperationKey( ? serializeInput(options.serializer ?? defaultSerializer, options.input) : {}), ...(options.type !== undefined ? { type: options.type } : {}), - ...(options.fnOptions !== undefined ? { fnOptions: options.fnOptions } : {}), + ...(fnOptions !== undefined ? { fnOptions } : {}), }, ] as OperationKey } diff --git a/packages/tanstack-query/src/live-query.test.ts b/packages/tanstack-query/src/live-query.test.ts index f90d2c1f6..d03d300d4 100644 --- a/packages/tanstack-query/src/live-query.test.ts +++ b/packages/tanstack-query/src/live-query.test.ts @@ -1,4 +1,5 @@ -import { QueryClient } from '@tanstack/query-core' +import { sleep } from '@orpc/shared' +import { QueryClient, QueryObserver } from '@tanstack/query-core' import { liveQuery } from './live-query.ts' @@ -18,9 +19,9 @@ describe('liveQuery', async () => { it('works', async () => { const queryFn = liveQuery(async function* () { yield 1 - await new Promise((resolve) => setTimeout(resolve, 50)) + await sleep(50) yield 2 - await new Promise((resolve) => setTimeout(resolve, 50)) + await sleep(50) yield 3 }) @@ -45,15 +46,45 @@ describe('liveQuery', async () => { await resultPromise }) + it('keeps data written after the last event', async () => { + const queryFn = liveQuery(async function* () { + yield 1 + queryClient.setQueryData(['live-query'], 2) + }) + + await expect( + queryFn({ + queryKey: ['live-query'], + signal: new AbortController().signal, + client: queryClient, + } as any), + ).resolves.toEqual(2) + }) + + it('keeps null written after the last event', async () => { + const queryFn = liveQuery(async function* () { + yield 1 + queryClient.setQueryData(['live-query'], null) + }) + + await expect( + queryFn({ + queryKey: ['live-query'], + signal: new AbortController().signal, + client: queryClient, + } as any), + ).resolves.toBeNull() + }) + it('on abort signal', async () => { let cleanupCalled = false const queryFn = liveQuery(async function* () { try { yield 1 - await new Promise((resolve) => setTimeout(resolve, 50)) + await sleep(50) yield 2 - await new Promise((resolve) => setTimeout(resolve, 50)) + await sleep(50) yield 3 } finally { cleanupCalled = true @@ -80,32 +111,23 @@ describe('liveQuery', async () => { expect(cleanupCalled).toBe(true) }) - it('closes the stream without reading it when aborted before it resolves', async () => { - const started = vi.fn() - const controller = new AbortController() - - const queryFn = liveQuery(async () => { - controller.abort() - return (async function* () { - started() - yield 1 - })() - }) + it('throw if no data yielded', async () => { + const queryFn = liveQuery(async function* () {}) await expect( queryFn({ queryKey: ['live-query'], - signal: controller.signal, + signal: new AbortController().signal, client: queryClient, } as any), - ).rejects.toSatisfy((err) => err === controller.signal.reason) - - expect(started).not.toHaveBeenCalled() + ).rejects.toThrow( + 'Live query for ["live-query"] did not yield any data. Ensure the query function returns an AsyncIterable with at least one chunk.', + ) }) - it('throw if no data yielded', async () => { + it('rejects on an undefined event, which TanStack Query cannot store', async () => { const queryFn = liveQuery(async function* () { - // No yield + yield undefined }) await expect( @@ -114,8 +136,108 @@ describe('liveQuery', async () => { signal: new AbortController().signal, client: queryClient, } as any), - ).rejects.toThrow( - 'Live query for ["live-query"] did not yield any data. Ensure the query function returns an AsyncIterable with at least one chunk.', + ).rejects.toEqual( + new TypeError( + 'Live query for ["live-query"] yielded `undefined`, which TanStack Query cannot store.', + ), + ) + }) + + it("writes to the fetch's own query with a custom queryKeyHashFn", async () => { + await expect( + queryClient.fetchQuery({ + queryKey: ['live-query'], + queryKeyHashFn: () => 'custom', + queryFn: liveQuery( + async function* () { + yield 1 + yield 2 + }, + { settleOnFirstData: true }, + ), + }), + ).resolves.toBe(1) + + expect( + queryClient + .getQueryCache() + .getAll() + .map((query) => query.queryHash), + ).toEqual(['custom']) + }) + + it('keeps existing data until the first event', async () => { + const key = ['live-query'] + queryClient.setQueryData(key, 0) + const statuses: string[] = [] + onTestFinished( + queryClient.getQueryCache().subscribe(({ query }) => statuses.push(query.state.status)), + ) + + const promise = queryClient.fetchQuery({ + queryKey: key, + queryFn: liveQuery(async function* () { + await sleep(50) + yield 1 + }), + }) + + expect(queryClient.getQueryState(key)).toMatchObject({ + status: 'success', + fetchStatus: 'fetching', + data: 0, + }) + await expect(promise).resolves.toBe(1) + expect(statuses).not.toContain('pending') + }) + + it('settles on the first event marked stale when settleOnFirstData is true', async () => { + const key = ['live-query'] + const cleanup = vi.fn() + + await expect( + queryClient.fetchQuery({ + queryKey: key, + queryFn: liveQuery( + async function* () { + try { + yield 1 + await new Promise(() => {}) + } finally { + cleanup() + } + }, + { settleOnFirstData: true }, + ), + }), + ).resolves.toBe(1) + + await vi.waitFor(() => expect(cleanup).toHaveBeenCalledTimes(1)) + expect(queryClient.getQueryState(key)).toMatchObject({ + status: 'success', + fetchStatus: 'idle', + data: 1, + isInvalidated: true, + }) + }) + + it('follows the stream while an observer watches the query when settleOnFirstData is unobserved', async () => { + const observer = new QueryObserver(queryClient, { + queryKey: ['live-query'], + queryFn: liveQuery( + async function* () { + yield 1 + await sleep(10) + yield 2 + }, + { settleOnFirstData: 'unobserved' }, + ), + }) + onTestFinished(observer.subscribe(() => {})) + + await vi.waitFor(() => + expect(observer.getCurrentResult()).toMatchObject({ data: 2, fetchStatus: 'idle' }), ) + expect(queryClient.getQueryState(['live-query'])).toMatchObject({ isInvalidated: false }) }) }) diff --git a/packages/tanstack-query/src/live-query.ts b/packages/tanstack-query/src/live-query.ts index 753f8b912..fe0804a09 100644 --- a/packages/tanstack-query/src/live-query.ts +++ b/packages/tanstack-query/src/live-query.ts @@ -1,33 +1,39 @@ import type { Promisable } from '@orpc/shared' -import { stringifyJSON, throwIfAborted } from '@orpc/shared' +import { pipeStreamToCache, stringifyJSON } from '@orpc/shared' import type { QueryFunction, QueryFunctionContext, QueryKey } from '@tanstack/query-core' +import type { SettleOnFirstDataOptions } from './stream-cache.ts' +import { followStream } from './stream-cache.ts' + export function liveQuery( queryFn: (context: QueryFunctionContext) => Promisable>, + { settleOnFirstData }: SettleOnFirstDataOptions = {}, ): QueryFunction { return async (context) => { - const stream = await queryFn(context) - - if (context.signal.aborted) { - await stream[Symbol.asyncIterator]().return?.() - throw context.signal.reason - } + const result = await followStream, TQueryFnData>( + context, + settleOnFirstData, + async (callContext, writes) => + pipeStreamToCache(await queryFn(callContext), { + ...writes, + onChunk: (_, chunk) => { + if (chunk === undefined) { + throw new TypeError( + `Live query for ${stringifyJSON(context.queryKey)} yielded \`undefined\`, which TanStack Query cannot store.`, + ) + } - let last: { chunk: TQueryFnData } | undefined - - for await (const chunk of stream) { - throwIfAborted(context.signal) - - last = { chunk } - context.client.setQueryData(context.queryKey, chunk) - } + return chunk + }, + }), + ) - if (!last) { + if (!result) { throw new TypeError( `Live query for ${stringifyJSON(context.queryKey)} did not yield any data. Ensure the query function returns an AsyncIterable with at least one chunk.`, ) } - return last.chunk + return result.data } } diff --git a/packages/tanstack-query/src/procedure-utils.test-d.ts b/packages/tanstack-query/src/procedure-utils.test-d.ts index e7ce61267..0fd39debe 100644 --- a/packages/tanstack-query/src/procedure-utils.test-d.ts +++ b/packages/tanstack-query/src/procedure-utils.test-d.ts @@ -20,6 +20,7 @@ import * as z from 'zod' import type { ProcedureUtils, ProcedureUtilsOptions } from './procedure-utils.ts' import type { InfiniteOptionsIn, + LiveOptionsIn, MutationKeyOptions, MutationOptionsIn, QueryKeyOptions, @@ -294,10 +295,15 @@ describe('ProcedureUtils', () => { input: { search: 'search' }, queryFnOptions: { maxChunks: 1, - refetchMode: 'replace', + refetchMode: 'append', }, }), ).toExtend() + + // @ts-expect-error - 'replace' is no longer a refetch mode + optionalUtils.streamedKey({ queryFnOptions: { refetchMode: 'replace' } }) + // @ts-expect-error - settleOnFirstData is not part of the key + optionalUtils.streamedKey({ settleOnFirstData: true }) }) it('.getQueryState is typed correctly', () => { @@ -369,6 +375,25 @@ describe('ProcedureUtils', () => { streamUtils.streamedOptions({ initialData: () => 'invalid' }) }) + it('accepts settleOnFirstData and serializable queryFnOptions', () => { + streamUtils.streamedOptions({ settleOnFirstData: true }) + streamUtils.streamedOptions({ settleOnFirstData: false }) + streamUtils.streamedOptions({ settleOnFirstData: 'unobserved' }) + // @ts-expect-error - invalid settleOnFirstData + streamUtils.streamedOptions({ settleOnFirstData: 'always' }) + // @ts-expect-error - settleOnFirstData is only for streamed and live queries + streamUtils.queryOptions({ settleOnFirstData: true }) + streamUtils.infiniteOptions({ + input: () => ({}), + initialPageParam: 0, + getNextPageParam: () => 1, + // @ts-expect-error - settleOnFirstData is only for streamed and live queries + settleOnFirstData: true, + }) + + streamUtils.streamedOptions({ queryFnOptions: { refetchMode: 'reset', maxChunks: 10 } }) + }) + it('return valid streamed options', () => { expectTypeOf(streamUtils.streamedOptions()).toExtend< QueryObserverOptions< @@ -528,6 +553,14 @@ describe('ProcedureUtils', () => { streamUtils.liveOptions({ initialData: () => 'invalid' }) }) + it('accepts settleOnFirstData', () => { + streamUtils.liveOptions({ settleOnFirstData: true }) + streamUtils.liveOptions({ settleOnFirstData: false }) + streamUtils.liveOptions({ settleOnFirstData: 'unobserved' }) + // @ts-expect-error - invalid settleOnFirstData + streamUtils.liveOptions({ settleOnFirstData: 'always' }) + }) + it('return valid streamed options', () => { expectTypeOf(streamUtils.liveOptions()).toExtend< QueryObserverOptions< @@ -1134,6 +1167,7 @@ describe('ProcedureUtilsOptions', () => { streamedOptions: { input: { search: 'test' }, staleTime: 1000, + settleOnFirstData: 'unobserved', }, } @@ -1214,11 +1248,12 @@ describe('ProcedureUtilsOptions', () => { } }) - it('liveOptions should accept Partial | modifier function', () => { + it('liveOptions should accept Partial | modifier function', () => { const _defaults: TestOptions = { liveOptions: { input: { search: 'test' }, staleTime: 1000, + settleOnFirstData: 'unobserved', }, } @@ -1232,7 +1267,7 @@ describe('ProcedureUtilsOptions', () => { const _defaults2: TestOptions = { liveOptions: (options) => { expectTypeOf(options).toEqualTypeOf< - QueryOptionsIn + LiveOptionsIn >() return { diff --git a/packages/tanstack-query/src/procedure-utils.test.ts b/packages/tanstack-query/src/procedure-utils.test.ts index 5a7519f15..63937d563 100644 --- a/packages/tanstack-query/src/procedure-utils.test.ts +++ b/packages/tanstack-query/src/procedure-utils.test.ts @@ -185,7 +185,7 @@ describe('procedureUtils', () => { input: { search: '__search__' }, context: { batch: '__batch__' }, queryFnOptions: { - refetchMode: 'replace', + refetchMode: 'append', }, }) @@ -197,7 +197,7 @@ describe('procedureUtils', () => { type: 'streamed', input: { search: '__search__' }, fnOptions: { - refetchMode: 'replace', + refetchMode: 'append', }, }) expect(options.queryKey).toBe(generateOperationKeySpy.mock.results[0]!.value) @@ -212,7 +212,8 @@ describe('procedureUtils', () => { ).resolves.toEqual(['__1__', '__2__']) expect(streamedQuerySpy).toHaveBeenCalledTimes(1) expect(streamedQuerySpy).toHaveBeenCalledWith(expect.any(Function), { - refetchMode: 'replace', + refetchMode: 'append', + settleOnFirstData: undefined, }) expect(queryClient.getQueryData(options.queryKey)).toEqual(['__1__', '__2__']) @@ -220,7 +221,7 @@ describe('procedureUtils', () => { expect(client).toHaveBeenCalledWith( { search: '__search__' }, { - signal, + signal: expect.any(AbortSignal), context: { batch: '__batch__', [OPERATION_CONTEXT_SYMBOL]: { @@ -232,6 +233,30 @@ describe('procedureUtils', () => { ) }) + it('passes settleOnFirstData to the query function, not to the key or the query options', async () => { + const options = utils.streamedOptions({ + input: { search: '__search__' }, + queryFnOptions: { maxChunks: 2 }, + settleOnFirstData: 'unobserved', + }) + + expect(generateOperationKeySpy).toHaveBeenCalledWith(['ping'], { + type: 'streamed', + input: { search: '__search__' }, + fnOptions: { maxChunks: 2 }, + }) + expect(options).not.toHaveProperty('settleOnFirstData') + + client.mockImplementationOnce(async function* () { + yield '__1__' + }) + await options.queryFn!({ signal, client: queryClient, queryKey: options.queryKey } as any) + expect(streamedQuerySpy).toHaveBeenCalledWith(expect.any(Function), { + maxChunks: 2, + settleOnFirstData: 'unobserved', + }) + }) + it('with skipToken', () => { const options = utils.streamedOptions({ input: skipToken, context: { batch: '__batch__' } }) @@ -264,7 +289,7 @@ describe('procedureUtils', () => { expect(client).toHaveBeenCalledWith( { search: '__search__' }, { - signal, + signal: expect.any(AbortSignal), context: { batch: '__batch__', [OPERATION_CONTEXT_SYMBOL]: { @@ -333,14 +358,16 @@ describe('procedureUtils', () => { options.queryFn!({ signal, client: queryClient, queryKey: options.queryKey } as any), ).resolves.toEqual('__2__') expect(liveQuerySpy).toHaveBeenCalledTimes(1) - expect(liveQuerySpy).toHaveBeenCalledWith(expect.any(Function)) + expect(liveQuerySpy).toHaveBeenCalledWith(expect.any(Function), { + settleOnFirstData: undefined, + }) expect(queryClient.getQueryData(options.queryKey)).toEqual('__2__') expect(client).toHaveBeenCalledTimes(1) expect(client).toHaveBeenCalledWith( { search: '__search__' }, { - signal, + signal: expect.any(AbortSignal), context: { batch: '__batch__', [OPERATION_CONTEXT_SYMBOL]: { @@ -352,6 +379,25 @@ describe('procedureUtils', () => { ) }) + it('passes settleOnFirstData to the query function, not to the key or the query options', async () => { + const options = utils.liveOptions({ + input: { search: '__search__' }, + settleOnFirstData: true, + }) + + expect(generateOperationKeySpy).toHaveBeenCalledWith(['ping'], { + type: 'live', + input: { search: '__search__' }, + }) + expect(options).not.toHaveProperty('settleOnFirstData') + + client.mockImplementationOnce(async function* () { + yield '__1__' + }) + await options.queryFn!({ signal, client: queryClient, queryKey: options.queryKey } as any) + expect(liveQuerySpy).toHaveBeenCalledWith(expect.any(Function), { settleOnFirstData: true }) + }) + it('with skipToken', () => { const options = utils.liveOptions({ input: skipToken, context: { batch: '__batch__' } }) @@ -384,7 +430,7 @@ describe('procedureUtils', () => { expect(client).toHaveBeenCalledWith( { search: '__search__' }, { - signal, + signal: expect.any(AbortSignal), context: { batch: '__batch__', [OPERATION_CONTEXT_SYMBOL]: { @@ -1320,12 +1366,16 @@ describe('createProcedureUtils with interceptors', () => { type: 'streamed', }, }, - fnContext: { signal, client: queryClient, queryKey: options.queryKey }, + fnContext: { + signal: expect.any(AbortSignal), + client: queryClient, + queryKey: options.queryKey, + }, }) expect(client).toHaveBeenCalledWith( { search: '__intercepted__' }, { - signal, + signal: expect.any(AbortSignal), context: { batch: '__intercepted_batch__', [OPERATION_CONTEXT_SYMBOL]: { @@ -1378,12 +1428,16 @@ describe('createProcedureUtils with interceptors', () => { type: 'live', }, }, - fnContext: { signal, client: queryClient, queryKey: options.queryKey }, + fnContext: { + signal: expect.any(AbortSignal), + client: queryClient, + queryKey: options.queryKey, + }, }) expect(client).toHaveBeenCalledWith( { search: '__intercepted__' }, { - signal, + signal: expect.any(AbortSignal), context: { batch: '__intercepted_batch__', [OPERATION_CONTEXT_SYMBOL]: { diff --git a/packages/tanstack-query/src/procedure-utils.ts b/packages/tanstack-query/src/procedure-utils.ts index 87e90fc57..87f27432d 100644 --- a/packages/tanstack-query/src/procedure-utils.ts +++ b/packages/tanstack-query/src/procedure-utils.ts @@ -28,6 +28,7 @@ import type { InfiniteKeyOptions, InfiniteOptionsIn, InfiniteOptionsOut, + LiveOptionsIn, MutationKeyOptions, MutationOptionsIn, MutationOptionsOut, @@ -208,7 +209,7 @@ export interface ProcedureUtilsOptions< */ liveOptions?: | ProcedureUtilsModifier< - QueryOptionsIn< + LiveOptionsIn< TClientContext, TInput, InferLiveQueryOutput, @@ -510,6 +511,7 @@ export class ProcedureUtils< * Configure queries for [AsyncIteratorObject](https://orpc.dev/docs/async-iterator-object). * This is built on [TanStack Query streamedQuery](https://tanstack.com/query/latest/docs/reference/streamedQuery) * and works with hooks like `useQuery`, `useSuspenseQuery`, or `prefetchQuery`. + * The query never returns to pending once it has data. */ streamedOptions, UInitialData = undefined>( ...rest: MaybeOptionalOptions< @@ -534,9 +536,10 @@ export class ProcedureUtils< } const queryKey = this.streamedKey(optionsIn) + const { settleOnFirstData, ...queryOptions } = optionsIn return { - ...(optionsIn as any), + ...(queryOptions as any), queryKey, queryFn: (fnContext) => { return intercept( @@ -562,15 +565,18 @@ export class ProcedureUtils< return optionsIn.queryFn(fnContext) as any } - return serializableStreamedQuery(async (queryContext) => { - const output = await this.call(input, { signal: queryContext.signal, context }) + return serializableStreamedQuery( + async (queryContext) => { + const output = await this.call(input, { signal: queryContext.signal, context }) - if (!isAsyncIteratorObject(output)) { - throw new Error('streamedQuery requires an AsyncIteratorObject output') - } + if (!isAsyncIteratorObject(output)) { + throw new Error('streamedQuery requires an AsyncIteratorObject output') + } - return output - }, optionsIn.queryFnOptions)(fnContext) + return output + }, + { ...optionsIn.queryFnOptions, settleOnFirstData }, + )(fnContext) }, ) }, @@ -607,7 +613,7 @@ export class ProcedureUtils< */ liveOptions, UInitialData = undefined>( ...rest: MaybeOptionalOptions< - QueryOptionsIn< + LiveOptionsIn< TClientContext, TInput, InferLiveQueryOutput, @@ -626,10 +632,11 @@ export class ProcedureUtils< } const queryKey = this.liveKey(optionsIn) + const { settleOnFirstData, ...queryOptions } = optionsIn return { structuralSharing: createStructuralSharing(this.options.serializer), - ...(optionsIn as any), + ...(queryOptions as any), queryKey, queryFn: (fnContext) => { return intercept( @@ -655,18 +662,21 @@ export class ProcedureUtils< return optionsIn.queryFn(fnContext) } - return liveQuery(async (queryContext) => { - const output = await this.call(input, { - signal: queryContext.signal, - context, - }) + return liveQuery( + async (queryContext) => { + const output = await this.call(input, { + signal: queryContext.signal, + context, + }) - if (!isAsyncIteratorObject(output)) { - throw new Error('liveQuery requires an AsyncIteratorObject output') - } + if (!isAsyncIteratorObject(output)) { + throw new Error('liveQuery requires an AsyncIteratorObject output') + } - return output - })(fnContext) + return output + }, + { settleOnFirstData }, + )(fnContext) }, ) }, diff --git a/packages/tanstack-query/src/ssr.test.ts b/packages/tanstack-query/src/ssr.test.ts index aa59a8627..4aa7b4a84 100644 --- a/packages/tanstack-query/src/ssr.test.ts +++ b/packages/tanstack-query/src/ssr.test.ts @@ -1,8 +1,17 @@ import type { Client } from '@orpc/client' import { RPCJsonSerializer } from '@orpc/client' import { sleep } from '@orpc/shared' -import { dehydrate, hydrate, QueryCache, QueryClient } from '@tanstack/query-core' - +import type { QueryObserverOptions } from '@tanstack/query-core' +import { + defaultShouldDehydrateQuery, + dehydrate, + environmentManager, + hydrate, + QueryClient, + QueryObserver, +} from '@tanstack/query-core' + +import { distinctData } from '../tests/__shared__/utils.ts' import { createRouterUtils } from './router-utils.ts' import { createHydrationOptions, StreamQuerySSRUtilsPlugin } from './ssr.ts' @@ -64,22 +73,62 @@ describe('createHydrationOptions', () => { }) describe('streamQuerySSRUtilsPlugin', () => { + const started = vi.fn() const cleanup = vi.fn() const client = { - stream: vi.fn(async function* () { - try { - yield 1 - - while (true) { + stream: vi.fn(async (_input: undefined, { signal }: { signal: AbortSignal }) => { + await sleep(10) + return (async function* () { + started() + try { + yield 1 await sleep(10) yield 2 + await new Promise((_, reject) => + signal.addEventListener('abort', () => reject(signal.reason)), + ) + } finally { + cleanup() } - } finally { - cleanup() - } + })() }), } as unknown as { stream: Client, Error> } + const finiteClient = { + stream: vi.fn(async function* () { + yield 1 + yield 2 + }), + } as unknown as typeof client + + function callSignal(index = 0): AbortSignal { + return vi.mocked(client.stream).mock.calls[index]![1]!.signal! + } + + function createUtils(isServer: boolean, streamClient = client) { + return createRouterUtils(streamClient, { + plugins: [new StreamQuerySSRUtilsPlugin({ isServer })], + }) + } + + function observe( + queryClient: QueryClient, + options: QueryObserverOptions, + ) { + const observer = new QueryObserver(queryClient, options) + const results = [observer.getCurrentResult()] + onTestFinished(observer.subscribe((result) => results.push(result))) + return { observer, results } + } + + function createQueryClient( + options: ConstructorParameters[0] = {}, + ): QueryClient { + const queryClient = new QueryClient(options) + onTestFinished(() => queryClient.cancelQueries()) + return queryClient + } + beforeEach(() => { vi.clearAllMocks() }) @@ -88,162 +137,310 @@ describe('streamQuerySSRUtilsPlugin', () => { vi.unstubAllGlobals() }) - it('cancels live queries on their first event on the server', async () => { - const onSuccess = vi.fn() - const queryClient = new QueryClient({ queryCache: new QueryCache({ onSuccess }) }) - const utils = createRouterUtils(client, { - plugins: [new StreamQuerySSRUtilsPlugin({ isServer: true })], + describe('on the server', () => { + it('settles streamed queries once they connect', async () => { + const queryClient = createQueryClient() + const options = createUtils(true).stream.streamedOptions() + + await expect(queryClient.fetchQuery(options)).resolves.toEqual([]) + expect(queryClient.getQueryState(options.queryKey)).toMatchObject({ isInvalidated: true }) }) - await expect(queryClient.fetchQuery(utils.stream.liveOptions())).resolves.toBe(1) - expect(vi.mocked(client.stream).mock.calls[0]![1]!.signal!.aborted).toBe(true) + it('settles live queries on their first event', async () => { + const queryClient = createQueryClient() + const options = createUtils(true).stream.liveOptions() - // the stream ignores the signal, so it closes on its next event, which is dropped - await vi.waitFor(() => expect(cleanup).toHaveBeenCalledTimes(1)) - expect(queryClient.getQueryState(utils.stream.liveKey())).toMatchObject({ - status: 'success', - fetchStatus: 'idle', - data: 1, + await expect(queryClient.fetchQuery(options)).resolves.toBe(1) + expect(queryClient.getQueryState(options.queryKey)).toMatchObject({ isInvalidated: true }) }) - expect(onSuccess).not.toHaveBeenCalled() - }) - it('cancels streamed queries once they connect on the server', async () => { - const onSuccess = vi.fn() - const queryClient = new QueryClient({ queryCache: new QueryCache({ onSuccess }) }) - const utils = createRouterUtils(client, { - plugins: [new StreamQuerySSRUtilsPlugin({ isServer: true })], - }) + it('settles even when an observer is subscribed, as vue and svelte do on the server', async () => { + const queryClient = createQueryClient() + const { observer } = observe(queryClient, createUtils(true).stream.streamedOptions()) - await expect(queryClient.fetchQuery(utils.stream.streamedOptions())).resolves.toEqual([]) - expect(vi.mocked(client.stream).mock.calls[0]![1]!.signal!.aborted).toBe(true) + await vi.waitFor(() => + expect(observer.getCurrentResult()).toMatchObject({ + status: 'success', + fetchStatus: 'idle', + data: [], + }), + ) + expect(queryClient.getQueryState(observer.options.queryKey)).toMatchObject({ + isInvalidated: true, + }) - // the stream is closed before its generator starts - await expect(vi.mocked(client.stream).mock.results[0]!.value.next()).resolves.toEqual({ - done: true, - value: undefined, + await sleep(30) + expect(client.stream).toHaveBeenCalledTimes(1) + expect(started).not.toHaveBeenCalled() }) - expect(cleanup).not.toHaveBeenCalled() - expect(queryClient.getQueryState(utils.stream.streamedKey())).toMatchObject({ - status: 'success', - fetchStatus: 'idle', - data: [], + + it('waits for a finite stream to end with a per-call settleOnFirstData: false', async () => { + const queryClient = createQueryClient() + const utils = createUtils(true, finiteClient) + + await expect( + queryClient.fetchQuery(utils.stream.streamedOptions({ settleOnFirstData: false })), + ).resolves.toEqual([1, 2]) + await expect( + queryClient.fetchQuery(utils.stream.liveOptions({ settleOnFirstData: false })), + ).resolves.toBe(2) }) - expect(onSuccess).not.toHaveBeenCalled() - }) - it.each(['reset', 'append', 'replace'] as const)( - 'cancels streamed queries only after they connect on the server (%s)', - async (refetchMode) => { - const queryClient = new QueryClient() - const abortedOnCall: boolean[] = [] - const connectClient = { - stream: vi.fn((input, options) => { - abortedOnCall.push(options.signal.aborted) - return client.stream(input, options) - }), - } as unknown as typeof client - const utils = createRouterUtils(connectClient, { + it('lets utils-level options override settleOnFirstData', async () => { + const queryClient = createQueryClient() + const utils = createRouterUtils(finiteClient, { plugins: [new StreamQuerySSRUtilsPlugin({ isServer: true })], + scoped: { stream: { streamedOptions: { settleOnFirstData: false } } }, }) - await expect( - queryClient.fetchQuery(utils.stream.streamedOptions({ queryFnOptions: { refetchMode } })), - ).resolves.toEqual([]) - - expect(abortedOnCall).toEqual([false]) - expect(vi.mocked(connectClient.stream).mock.calls[0]![1]!.signal!.aborted).toBe(true) - }, - ) - - it('forwards query cancellation to the call on the server', async () => { - const queryClient = new QueryClient() - const pendingClient = { - stream: vi.fn(async function* (_input, { signal }) { - await new Promise((_, reject) => - signal.addEventListener('abort', () => reject(signal.reason)), - ) - }), - } as unknown as typeof client - const utils = createRouterUtils(pendingClient, { - plugins: [new StreamQuerySSRUtilsPlugin({ isServer: true })], + await expect(queryClient.fetchQuery(utils.stream.streamedOptions())).resolves.toEqual([1, 2]) + await expect(queryClient.fetchQuery(utils.stream.liveOptions())).resolves.toBe(1) }) - void queryClient.prefetchQuery(utils.stream.liveOptions()) - await vi.waitFor(() => expect(pendingClient.stream).toHaveBeenCalledTimes(1)) + it.each([ + ['streamed', []], + ['live', 1], + ] as const)( + 'resolves next() in an interceptor when the fetch settles (%s)', + async (type, firstData) => { + const interceptor = vi.fn((options: any) => options.next()) + const utils = createRouterUtils(client, { + plugins: [new StreamQuerySSRUtilsPlugin({ isServer: true })], + streamedInterceptors: [interceptor], + liveInterceptors: [interceptor], + }).stream + + await createQueryClient().fetchQuery( + (type === 'streamed' ? utils.streamedOptions() : utils.liveOptions()) as any, + ) + await expect(interceptor.mock.results[0]!.value).resolves.toEqual(firstData) + }, + ) + }) - const signal = vi.mocked(pendingClient.stream).mock.calls[0]![1]!.signal! - expect(signal.aborted).toBe(false) + describe('in the browser', () => { + it.each([ + ['streamed', [], [[], [1], [1, 2]]], + ['live', 1, [1, 2]], + ] as const)( + 'settles an unobserved suspense fetch at its first data, then reopens the stream on subscribe (%s)', + async (type, firstData, sequence) => { + const queryClient = createQueryClient() + const utils = createUtils(false).stream + const options: QueryObserverOptions = + type === 'streamed' ? utils.streamedOptions() : utils.liveOptions() + + await expect( + new QueryObserver(queryClient, options).fetchOptimistic(options), + ).resolves.toMatchObject({ data: firstData }) + await vi.waitFor(() => + expect(queryClient.getQueryState(options.queryKey)).toMatchObject({ + fetchStatus: 'idle', + isInvalidated: true, + }), + ) - await queryClient.cancelQueries() - expect(signal.aborted).toBe(true) - }) + const { observer, results } = observe(queryClient, options) - it('settles live queries whose stream ends before the first event on the server', async () => { - const queryClient = new QueryClient() - const emptyClient = { stream: vi.fn(async function* () {}) } as unknown as typeof client - const utils = createRouterUtils(emptyClient, { - plugins: [new StreamQuerySSRUtilsPlugin({ isServer: true })], + await vi.waitFor(() => + expect(observer.getCurrentResult()).toMatchObject({ + data: sequence.at(-1), + fetchStatus: 'fetching', + }), + ) + expect(results.map((result) => result.status)).not.toContain('pending') + expect(distinctData(results)).toEqual(sequence) + expect(client.stream).toHaveBeenCalledTimes(2) + }, + ) + + it('reopens a settled suspense fetch when a mounted suspense observer switches to its key', async () => { + const queryClient = createQueryClient() + const utils = createUtils(false).stream + // React clamps the staleTime of suspense queries to at least 1000ms + const optionsA = { + ...utils.streamedOptions({ queryKey: ['a'] }), + suspense: true, + staleTime: 1000, + } + const optionsB = { + ...utils.streamedOptions({ queryKey: ['b'] }), + suspense: true, + staleTime: 1000, + } + const { observer } = observe(queryClient, optionsA) + await vi.waitFor(() => expect(observer.getCurrentResult()).toMatchObject({ data: [1, 2] })) + + await expect(observer.fetchOptimistic(optionsB)).resolves.toMatchObject({ data: [] }) + await vi.waitFor(() => + expect(queryClient.getQueryState(['b'])).toMatchObject({ + fetchStatus: 'idle', + isInvalidated: true, + }), + ) + + observer.setOptions(optionsB) + await vi.waitFor(() => + expect(observer.getCurrentResult()).toMatchObject({ + data: [1, 2], + fetchStatus: 'fetching', + }), + ) + expect(client.stream).toHaveBeenCalledTimes(3) }) - await expect(queryClient.fetchQuery(utils.stream.liveOptions())).rejects.toThrow( - 'did not yield any data', + it.each([ + ['streamed', [1, 2]], + ['live', 2], + ] as const)( + 'reopens a stream cut off by unmount when it remounts within staleTime (%s)', + async (type, data) => { + const queryClient = createQueryClient({ + defaultOptions: { queries: { staleTime: 60_000 } }, + }) + const utils = createUtils(false).stream + const options: QueryObserverOptions = + type === 'streamed' ? utils.streamedOptions() : utils.liveOptions() + + const unsubscribe = new QueryObserver(queryClient, options).subscribe(() => {}) + await vi.waitFor(() => expect(queryClient.getQueryData(options.queryKey)).toEqual(data)) + unsubscribe() + + expect(callSignal().aborted).toBe(true) + expect(queryClient.getQueryState(options.queryKey)).toMatchObject({ + data, + fetchStatus: 'idle', + isInvalidated: true, + }) + + const { observer, results } = observe(queryClient, options) + + await vi.waitFor(() => + expect(observer.getCurrentResult()).toMatchObject({ fetchStatus: 'fetching' }), + ) + expect(results.map((result) => result.status)).not.toContain('pending') + expect(client.stream).toHaveBeenCalledTimes(2) + }, ) }) - it('follows streams in the browser', async () => { - const queryClient = new QueryClient() - const utils = createRouterUtils(client, { - plugins: [new StreamQuerySSRUtilsPlugin({ isServer: false })], - }) - - void queryClient.prefetchQuery(utils.stream.liveOptions()) + describe('end to end with hydration', () => { + function createSSRQueryClient(staleTime: number | 'static' = 60_000): QueryClient { + return createQueryClient({ + defaultOptions: { queries: { staleTime }, ...createHydrationOptions() }, + }) + } - await vi.waitFor(() => expect(queryClient.getQueryData(utils.stream.liveKey())).toBe(2)) - expect(vi.mocked(client.stream).mock.calls[0]![1]!.signal!.aborted).toBe(false) - await queryClient.cancelQueries() - }) + function getOptions( + isServer: boolean, + type: 'streamed' | 'live', + ): QueryObserverOptions { + const utils = createUtils(isServer).stream + return type === 'streamed' ? utils.streamedOptions() : utils.liveOptions() + } - it('detects the server by the absence of window', () => { - expect(new StreamQuerySSRUtilsPlugin().initProcedureOptions([], {})).toMatchObject({ - streamedInterceptors: [expect.any(Function)], - liveInterceptors: [expect.any(Function)], - }) + async function renderOnServer(type: 'streamed' | 'live', browser: QueryClient): Promise { + const server = createSSRQueryClient() + await server.prefetchQuery(getOptions(true, type)) + hydrate(browser, JSON.parse(JSON.stringify(dehydrate(server)))) + } - vi.stubGlobal('window', {}) - expect(new StreamQuerySSRUtilsPlugin().initProcedureOptions([], {})).not.toHaveProperty( - 'streamedInterceptors', + it.each([ + ['streamed', [], [[], [1], [1, 2]]], + ['live', 1, [1, 2]], + ] as const)( + 'reopens the stream on mount despite staleTime, keeping the server data on screen (%s)', + async (type, firstData, sequence) => { + const browser = createSSRQueryClient() + await renderOnServer(type, browser) + + const options = getOptions(false, type) + expect(browser.getQueryState(options.queryKey)).toMatchObject({ + status: 'success', + data: firstData, + isInvalidated: true, + }) + + const { observer, results } = observe(browser, options) + + await vi.waitFor(() => + expect(observer.getCurrentResult()).toMatchObject({ + data: sequence.at(-1), + fetchStatus: 'fetching', + }), + ) + expect(results.map((result) => result.status)).not.toContain('pending') + expect(distinctData(results)).toEqual(sequence) + }, ) - expect(new StreamQuerySSRUtilsPlugin().initProcedureOptions([], {})).not.toHaveProperty( - 'liveInterceptors', + + it.each([ + ['streamed', [], [[], [1], [1, 2]]], + ['live', 1, [1, 2]], + ] as const)( + 'hydrates a query dehydrated while pending with its first data, then reopens the stream on mount (%s)', + async (type, firstData, sequence) => { + const server = createSSRQueryClient() + const browser = createSSRQueryClient() + void server.prefetchQuery(getOptions(true, type)) + hydrate( + browser, + dehydrate(server, { + shouldDehydrateQuery: (query) => + defaultShouldDehydrateQuery(query) || query.state.status === 'pending', + }), + ) + + const options = getOptions(false, type) + await vi.waitFor(() => + expect(browser.getQueryState(options.queryKey)).toMatchObject({ fetchStatus: 'idle' }), + ) + expect(browser.getQueryState(options.queryKey)).toMatchObject({ + status: 'success', + data: firstData, + }) + + const { observer, results } = observe(browser, options) + + await vi.waitFor(() => + expect(observer.getCurrentResult()).toMatchObject({ + data: sequence.at(-1), + fetchStatus: 'fetching', + }), + ) + expect(distinctData(results)).toEqual(sequence) + expect(client.stream).toHaveBeenCalledTimes(2) + }, ) - }) - it('refetches streamed and live queries on mount', () => { - const utils = createRouterUtils(client, { plugins: [new StreamQuerySSRUtilsPlugin()] }) + it('does not reopen the stream with staleTime: static', async () => { + const browser = createSSRQueryClient('static') + await renderOnServer('streamed', browser) + + const { observer } = observe(browser, getOptions(false, 'streamed')) - expect(utils.stream.streamedOptions()).toMatchObject({ - initialData: [], - refetchOnMount: 'always', + await sleep(30) + expect(observer.getCurrentResult()).toMatchObject({ data: [], fetchStatus: 'idle' }) + expect(client.stream).toHaveBeenCalledTimes(1) }) - expect(utils.stream.liveOptions()).toMatchObject({ refetchOnMount: 'always' }) - expect(utils.stream.liveOptions()).not.toHaveProperty('initialData') - expect(utils.stream.queryOptions()).not.toHaveProperty('refetchOnMount') }) - it('lets utils and call options override its defaults', () => { - const utils = createRouterUtils(client, { - plugins: [new StreamQuerySSRUtilsPlugin()], - scoped: { - stream: { - streamedOptions: (options) => ({ ...options, refetchOnMount: false }), - }, - }, - }) + it("defaults to TanStack Query's server detection", () => { + const isServer = environmentManager.isServer() + try { + environmentManager.setIsServer(() => true) + expect(new StreamQuerySSRUtilsPlugin().initProcedureOptions([], {})).toEqual({ + streamedOptions: { settleOnFirstData: true, refetchOnMount: false }, + liveOptions: { settleOnFirstData: true, refetchOnMount: false }, + }) - expect(utils.stream.streamedOptions()).toMatchObject({ initialData: [], refetchOnMount: false }) - expect(utils.stream.liveOptions({ refetchOnMount: true })).toMatchObject({ - refetchOnMount: true, - }) + environmentManager.setIsServer(() => false) + expect(new StreamQuerySSRUtilsPlugin().initProcedureOptions([], {})).toEqual({ + streamedOptions: { settleOnFirstData: 'unobserved', refetchOnMount: 'always' }, + liveOptions: { settleOnFirstData: 'unobserved', refetchOnMount: 'always' }, + }) + } finally { + environmentManager.setIsServer(() => isServer) + } }) }) diff --git a/packages/tanstack-query/src/ssr.ts b/packages/tanstack-query/src/ssr.ts index 84be4c33f..43ddc6426 100644 --- a/packages/tanstack-query/src/ssr.ts +++ b/packages/tanstack-query/src/ssr.ts @@ -6,9 +6,10 @@ import type { } from '@orpc/client' import { RPCJsonSerializer } from '@orpc/client' import type { Public } from '@orpc/shared' +import { environmentManager } from '@tanstack/query-core' import type { RouterUtilsPlugin } from './plugin.ts' -import type { ProcedureUtilsOptions, ProcedureUtilsQueryInterceptor } from './procedure-utils.ts' +import type { ProcedureUtilsOptions } from './procedure-utils.ts' import { mergeProcedureUtilsOptions } from './procedure-utils.ts' /** @@ -39,40 +40,15 @@ export function createHydrationOptions( export interface StreamQuerySSRUtilsPluginOptions { /** - * Whether utils run on the server, where streamed and live queries are cancelled once they have data. - * Defaults to `true` when there is no global `window`. + * Whether utils run on the server (`settleOnFirstData: true`) or in the browser (`settleOnFirstData: 'unobserved'`). + * Defaults to TanStack Query's `environmentManager.isServer()`. */ isServer?: boolean | undefined } /** - * Cancels the query once it has data, keeping that data. The cancel aborts the call's `signal`. - */ -const cancelOnFirstData: ProcedureUtilsQueryInterceptor = async ({ - next, - fnContext, -}) => { - const cache = fnContext.client.getQueryCache() - const query = cache.find({ queryKey: fnContext.queryKey, exact: true }) - - const unsubscribe = cache.subscribe((event) => { - if (event.query === query && event.type === 'updated' && event.action.type === 'success') { - unsubscribe() - void event.query.cancel({ revert: true }) - } - }) - - try { - return await next() - } finally { - unsubscribe() - } -} - -/** - * Router utils plugin for SSR of streamed and live queries. On the server, streamed queries are - * cancelled once they connect and live queries on their first event, keeping the data received so far; - * in the browser, they refetch on mount so streams reopen after hydration. + * Router utils plugin that sets the `settleOnFirstData` and `refetchOnMount` options of `.streamedOptions` and `.liveOptions`, + * so streamed and live queries can be server-rendered and reopen their stream in the browser. * * @see {@link https://orpc.dev/docs/integrations/tanstack-query#streamed-and-live-queries | TanStack Query Integration - Streamed and Live Queries} */ @@ -81,10 +57,16 @@ export class StreamQuerySSRUtilsPlugin< > implements RouterUtilsPlugin { readonly name = '~stream-query-ssr' - private readonly isServer: boolean + private readonly settleOnFirstData: true | 'unobserved' + // browser: hydrated data (even dehydrated while pending) reopens its stream on mount, despite a `staleTime` + // server: an observer (vue, svelte) must not reopen a stream the render already settled + private readonly refetchOnMount: boolean | 'always' constructor(options: StreamQuerySSRUtilsPluginOptions = {}) { - this.isServer = options.isServer ?? !('window' in globalThis) + // the server always settles: vue and svelte observe queries there too + const isServer = options.isServer ?? environmentManager.isServer() + this.settleOnFirstData = isServer ? true : 'unobserved' + this.refetchOnMount = isServer ? false : 'always' } initProcedureOptions( @@ -93,11 +75,14 @@ export class StreamQuerySSRUtilsPlugin< ): ProcedureUtilsOptions, any, any, InferClientError> { return mergeProcedureUtilsOptions( { - ...(this.isServer - ? { streamedInterceptors: [cancelOnFirstData], liveInterceptors: [cancelOnFirstData] } - : {}), - streamedOptions: { initialData: [], refetchOnMount: 'always' }, - liveOptions: { refetchOnMount: 'always' }, + streamedOptions: { + settleOnFirstData: this.settleOnFirstData, + refetchOnMount: this.refetchOnMount, + }, + liveOptions: { + settleOnFirstData: this.settleOnFirstData, + refetchOnMount: this.refetchOnMount, + }, }, options, ) diff --git a/packages/tanstack-query/src/stream-cache.ts b/packages/tanstack-query/src/stream-cache.ts new file mode 100644 index 000000000..16fcb11e2 --- /dev/null +++ b/packages/tanstack-query/src/stream-cache.ts @@ -0,0 +1,123 @@ +import type { PipeStreamToCacheOptions } from '@orpc/shared' +import { runWithOwnSignal } from '@orpc/shared' +import type { Query, QueryCache, QueryFunctionContext } from '@tanstack/query-core' + +export interface SettleOnFirstDataOptions { + /** + * Settles the fetch at its first data (streamed: when the stream connects; live: on the first event) + * instead of following the stream to its end. The stream is closed and the data kept, marked stale, + * so the stream reopens once the query is observed (e.g. refetch on mount). + * - `true`: always + * - `'unobserved'`: only if no observer watches the query at that moment + * - `false`: never + * @default false + */ + settleOnFirstData?: boolean | 'unobserved' | undefined +} + +// retries call the query function again with the same signal +const trackedSignals = /* @__PURE__ */ new WeakSet() + +export type StreamWrites = Required< + Pick, 'signal' | 'update'> +> & + Pick, 'settle'> + +/** + * Runs a fetch's stream against the fetch's own query (so a custom `queryKeyHashFn` is respected), + * giving each call its own signal, aborted once its stream is no longer read. + * - A stream cut short (cancelled, or its last observer left) marks the query stale. + * - A settled fetch is marked stale once stored, so its stream reopens when observed. + * @returns the data to resolve the fetch with, or `undefined` if nothing was written + */ +export async function followStream, TData>( + context: TContext, + settleOnFirstData: SettleOnFirstDataOptions['settleOnFirstData'], + run: (callContext: TContext, writes: StreamWrites) => Promise<{ data: TData } | undefined>, +): Promise<{ data: TData } | undefined> { + const queryCache = context.client.getQueryCache() + let query: Query | undefined + const findQuery = () => (query ??= queryCache.find({ queryKey: context.queryKey, exact: true })) + + if (!trackedSignals.has(context.signal)) { + trackedSignals.add(context.signal) + context.signal.addEventListener('abort', () => findQuery()?.invalidate(), { once: true }) + } + + const update: StreamWrites['update'] = (updater) => { + const query = findQuery() + + if (!query) { + context.client.setQueryData(context.queryKey, updater as any) + return + } + + const current = query.state.data + const data = updater(current) + + // an unchanged write (e.g. 'append' at connect) is not new data + if (data !== current) { + query.setData(data, { manual: true }) + } + } + + let settled = false + const settle = settleOnFirstData + ? () => { + const query = findQuery() + settled = !!query && (settleOnFirstData === true || query.getObserversCount() === 0) + return settled + } + : undefined + + const result = await runWithOwnSignal(context.signal, (signal) => + run({ ...context, signal }, { signal: context.signal, update, settle }), + ) + + if (settled) { + invalidateOnceStored( + queryCache, + findQuery()!, + context.signal, + settleOnFirstData === 'unobserved', + ) + } + + if (!result) { + return undefined + } + + // the returned data is written again, so return the cache to keep later writes + const cached = findQuery()?.state.data as TData | undefined + return { data: cached === undefined ? result.data : cached } +} + +function invalidateOnceStored( + queryCache: QueryCache, + query: Query, + signal: AbortSignal, + reopenIfObserved: boolean, +): void { + const unsubscribe = queryCache.subscribe((event) => { + if (event.query !== query || event.type !== 'updated') { + return + } + + if (event.action.type === 'error' || event.action.type === 'failed') { + unsubscribe() + } else if (event.action.type === 'success' && !event.action.manual) { + unsubscribe() + query.invalidate() + + // an observer that arrived while the stream closed joined this fetch, so reopen the stream for it + if (reopenIfObserved && query.getObserversCount() > 0) { + queueMicrotask(() => { + void query.fetch().catch(() => {}) + }) + } + } + }) + + // a fetch cancelled before it is stored (e.g. by a refetch) must not invalidate the next one + signal.addEventListener('abort', unsubscribe, { once: true }) +} diff --git a/packages/tanstack-query/src/stream-query.test.ts b/packages/tanstack-query/src/stream-query.test.ts index 56ca6f1f9..387ccfe69 100644 --- a/packages/tanstack-query/src/stream-query.test.ts +++ b/packages/tanstack-query/src/stream-query.test.ts @@ -1,6 +1,15 @@ -import { sleep } from '@orpc/shared' -import { QueryClient } from '@tanstack/query-core' - +import { promiseWithResolvers, sleep } from '@orpc/shared' +import type { QueryKey, QueryState } from '@tanstack/query-core' +import { + dehydrate, + hashKey, + hydrate, + QueryCache, + QueryClient, + QueryObserver, +} from '@tanstack/query-core' + +import { distinctData } from '../tests/__shared__/utils.ts' import { serializableStreamedQuery } from './stream-query.ts' describe('serializableStreamedQuery', () => { @@ -16,474 +25,213 @@ describe('serializableStreamedQuery', () => { queryClient.clear() }) - describe('refetchMode option', () => { - it('works with reset refetch mode (default)', async () => { - const key = ['refetch-test'] - - const queryFn1 = serializableStreamedQuery(async () => { - await sleep(50) - return (async function* () { - await sleep(50) - yield 'initial1' - await sleep(50) - yield 'initial2' - })() - }) - - const promise1 = expect( - queryFn1({ - queryKey: key, - signal: new AbortController().signal, - client: queryClient, - } as any), - ).resolves.toEqual(['initial1', 'initial2']) - - // not define anything on initial - await sleep(0) - expect(queryClient.getQueryState(key)).toBeUndefined() - expect(queryClient.getQueryData(key)).toBeUndefined() - - await sleep(50) - expect(queryClient.getQueryState(key)?.status).toBe('success') - expect(queryClient.getQueryData(key)).toEqual([]) - - await sleep(50) - expect(queryClient.getQueryState(key)?.status).toBe('success') - expect(queryClient.getQueryData(key)).toEqual(['initial1']) - - await sleep(50) - expect(queryClient.getQueryState(key)?.status).toBe('success') - expect(queryClient.getQueryData(key)).toEqual(['initial1', 'initial2']) - - await promise1 - - // Refetch with reset mode (default) - const queryFn2 = serializableStreamedQuery(async () => { - await sleep(50) - return (async function* () { - await sleep(50) - yield 'refetch1' - await sleep(50) - yield 'refetch2' - })() - }) - - const resultPromise = expect( - queryFn2({ - queryKey: key, - signal: new AbortController().signal, - client: queryClient, - } as any), - ).resolves.toEqual(['refetch1', 'refetch2']) - - // should reset status to pending - await sleep(0) - expect(queryClient.getQueryState(key)?.status).toBe('pending') - expect(queryClient.getQueryData(key)).toBeUndefined() - - await sleep(50) - expect(queryClient.getQueryState(key)?.status).toBe('success') - expect(queryClient.getQueryData(key)).toEqual([]) - - await sleep(50) - expect(queryClient.getQueryState(key)?.status).toBe('success') - expect(queryClient.getQueryData(key)).toEqual(['refetch1']) - - await resultPromise - }) - - it('works with append refetch mode', async () => { - const key = ['append-test'] - - // First query - const queryFn1 = serializableStreamedQuery(async () => { - await sleep(50) - return (async function* () { - await sleep(50) - yield 'initial1' - await sleep(50) - yield 'initial2' - })() - }) - - const promise1 = expect( - queryFn1({ - queryKey: key, - signal: new AbortController().signal, - client: queryClient, - } as any), - ).resolves.toEqual(['initial1', 'initial2']) - - await sleep(0) - expect(queryClient.getQueryState(key)).toBeUndefined() - expect(queryClient.getQueryData(key)).toBeUndefined() - - await sleep(50) - expect(queryClient.getQueryState(key)?.status).toBe('success') - expect(queryClient.getQueryData(key)).toEqual([]) - - await sleep(50) - expect(queryClient.getQueryState(key)?.status).toBe('success') - expect(queryClient.getQueryData(key)).toEqual(['initial1']) - - await sleep(50) - expect(queryClient.getQueryState(key)?.status).toBe('success') - expect(queryClient.getQueryData(key)).toEqual(['initial1', 'initial2']) - - await promise1 - - // Refetch with append mode - const queryFn2 = serializableStreamedQuery( - async () => { - await sleep(50) - return (async function* () { - await sleep(50) - yield 'append1' - await sleep(50) - yield 'append2' - })() - }, - { refetchMode: 'append' }, - ) - - const promise2 = expect( - queryFn2({ - queryKey: key, - signal: new AbortController().signal, - client: queryClient, - } as any), - ).resolves.toEqual(['initial1', 'initial2', 'append1', 'append2']) - - await sleep(0) - expect(queryClient.getQueryState(key)?.status).toBe('success') - expect(queryClient.getQueryData(key)).toEqual(['initial1', 'initial2']) - - await sleep(50) - expect(queryClient.getQueryState(key)?.status).toBe('success') - expect(queryClient.getQueryData(key)).toEqual(['initial1', 'initial2']) - - await sleep(50) - expect(queryClient.getQueryState(key)?.status).toBe('success') - expect(queryClient.getQueryData(key)).toEqual(['initial1', 'initial2', 'append1']) - - await sleep(50) - expect(queryClient.getQueryState(key)?.status).toBe('success') - expect(queryClient.getQueryData(key)).toEqual(['initial1', 'initial2', 'append1', 'append2']) - - await promise2 + function stream(chunks: T[], delay = 10) { + return async () => { + await sleep(delay) + return (async function* () { + for (const chunk of chunks) { + await sleep(delay) + yield chunk + } + })() + } + } + + function recordStates(key: QueryKey, client = queryClient) { + const states: Pick[] = [] + const unsubscribe = client.getQueryCache().subscribe(({ query }) => { + if (query.queryHash === hashKey(key)) { + const { status, fetchStatus, data } = query.state + states.push({ status, fetchStatus, data }) + } }) + onTestFinished(unsubscribe) + return states + } - it('works with replace refetch mode', async () => { - const key = ['replace-test'] - - const queryFn1 = serializableStreamedQuery(async () => { - await sleep(50) - return (async function* () { - await sleep(50) - yield 'initial1' - await sleep(50) - yield 'initial2' - })() - }) - - const promise1 = expect( - queryFn1({ - queryKey: key, - signal: new AbortController().signal, - client: queryClient, - } as any), - ).resolves.toEqual(['initial1', 'initial2']) - - await sleep(0) - expect(queryClient.getQueryState(key)).toBeUndefined() - expect(queryClient.getQueryData(key)).toBeUndefined() - - await sleep(50) - expect(queryClient.getQueryState(key)?.status).toBe('success') - expect(queryClient.getQueryData(key)).toEqual([]) - - await sleep(50) - expect(queryClient.getQueryState(key)?.status).toBe('success') - expect(queryClient.getQueryData(key)).toEqual(['initial1']) - - await sleep(50) - expect(queryClient.getQueryState(key)?.status).toBe('success') - expect(queryClient.getQueryData(key)).toEqual(['initial1', 'initial2']) - - await promise1 - - // Refetch with replace mode - const queryFn2 = serializableStreamedQuery( - async () => { - await sleep(50) - return (async function* () { - await sleep(50) - yield 'replace1' - await sleep(50) - yield 'replace2' - })() - }, - { refetchMode: 'replace' }, - ) - - const promise2 = expect( - queryFn2({ + describe('refetchMode option', () => { + const origins: Record Promise<{ initialData?: string[] }>> = { + 'initialData': async () => ({ initialData: ['x1', 'x2'] }), + 'setQueryData': async (key) => { + queryClient.setQueryData(key, ['x1', 'x2']) + return {} + }, + 'a previous fetch': async (key) => { + await queryClient.fetchQuery({ queryKey: key, - signal: new AbortController().signal, - client: queryClient, - } as any), - ).resolves.toEqual(['replace1', 'replace2']) - - await sleep(0) - expect(queryClient.getQueryState(key)?.status).toBe('success') - expect(queryClient.getQueryData(key)).toEqual(['initial1', 'initial2']) - - await sleep(50) - expect(queryClient.getQueryState(key)?.status).toBe('success') - expect(queryClient.getQueryData(key)).toEqual(['initial1', 'initial2']) - - await sleep(50) - expect(queryClient.getQueryState(key)?.status).toBe('success') - expect(queryClient.getQueryData(key)).toEqual(['initial1', 'initial2']) - - await sleep(50) - expect(queryClient.getQueryState(key)?.status).toBe('success') - expect(queryClient.getQueryData(key)).toEqual(['replace1', 'replace2']) - - await promise2 - }) + queryFn: serializableStreamedQuery(stream(['x1', 'x2'])), + }) + return {} + }, + 'hydrate': async (key) => { + const server = new QueryClient() + server.setQueryData(key, ['x1', 'x2']) + hydrate(queryClient, dehydrate(server)) + return {} + }, + } + + it.each( + Object.keys(origins).flatMap( + (origin) => + [ + [origin, 'reset', [['x1', 'x2'], [], ['c1'], ['c1', 'c2']]], + [ + origin, + 'append', + [ + ['x1', 'x2'], + ['x1', 'x2', 'c1'], + ['x1', 'x2', 'c1', 'c2'], + ], + ], + ] as const, + ), + )( + 'keeps data from %s until the stream connects, never pending (%s)', + async (origin, refetchMode, expected) => { + const key = ['refetch-test', origin, refetchMode] + const options = await origins[origin]!(key) + const states = recordStates(key) + + await expect( + queryClient.fetchQuery({ + queryKey: key, + queryFn: serializableStreamedQuery(stream(['c1', 'c2']), { refetchMode }), + ...options, + }), + ).resolves.toEqual(expected.at(-1)) + + expect(distinctData(states)).toEqual(expected) + expect(states.map((state) => state.status)).not.toContain('pending') + expect(states.find((state) => state.fetchStatus === 'fetching')).toEqual({ + status: 'success', + fetchStatus: 'fetching', + data: ['x1', 'x2'], + }) + expect(states.at(-1)).toMatchObject({ status: 'success', fetchStatus: 'idle' }) + }, + ) }) describe('maxChunks option', () => { - it('works with append & parallel', async () => { - const queryFn = serializableStreamedQuery( - async () => { - await sleep(50) - - return (async function* () { - await sleep(50) - yield 'chunk1' - await sleep(50) - yield 'chunk2' + it.each([ + [ + 'append', + [ + ['pre5', 'pre6'], + ['pre6', 'chunk1'], + ], + ], + ['reset', [[], ['chunk1']]], + ] as const)( + 'works with %s & parallel', + async (refetchMode, [afterConnect, afterFirstChunk]) => { + const key = ['max-chunks-test'] + const queryFn = serializableStreamedQuery( + async () => { await sleep(50) - yield 'chunk3' - await sleep(50) - yield 'chunk4' - await sleep(50) - })() - }, - { maxChunks: 2, refetchMode: 'append' }, - ) - - const controller = new AbortController() - - // exceed maxChunks with pre-existing data - should keep only the latest 2 chunks - queryClient.setQueryData(['max-chunks-test'], ['pre1', 'pre2', 'pre3']) - - const resultPromise = expect( - queryFn({ - queryKey: ['max-chunks-test'], - signal: controller.signal, - client: queryClient, - } as any), - ).resolves.toEqual(['parallel5', 'parallel6']) - - await sleep(0) - expect(queryClient.getQueryData(['max-chunks-test'])).toEqual(['pre2', 'pre3']) - - // exceed while still resolve streaming - queryClient.setQueryData(['max-chunks-test'], ['pre4', 'pre5', 'pre6']) - await sleep(50) - expect(queryClient.getQueryData(['max-chunks-test'])).toEqual(['pre5', 'pre6']) - - await sleep(50) - expect(queryClient.getQueryData(['max-chunks-test'])).toEqual(['pre6', 'chunk1']) - await sleep(50) - expect(queryClient.getQueryData(['max-chunks-test'])).toEqual(['chunk1', 'chunk2']) + return (async function* () { + await sleep(50) + yield 'chunk1' + await sleep(50) + yield 'chunk2' + await sleep(50) + yield 'chunk3' + await sleep(50) + yield 'chunk4' + await sleep(50) + })() + }, + { maxChunks: 2, refetchMode }, + ) + + queryClient.setQueryData(key, ['pre1', 'pre2', 'pre3']) + + const resultPromise = expect( + queryFn({ + queryKey: key, + signal: new AbortController().signal, + client: queryClient, + } as any), + ).resolves.toEqual(['parallel5', 'parallel6']) + + await sleep(0) + expect(queryClient.getQueryData(key)).toEqual(['pre1', 'pre2', 'pre3']) + + queryClient.setQueryData(key, ['pre4', 'pre5', 'pre6']) + await sleep(50) + expect(queryClient.getQueryData(key)).toEqual(afterConnect) - await sleep(50) - expect(queryClient.getQueryData(['max-chunks-test'])).toEqual(['chunk2', 'chunk3']) + await sleep(50) + expect(queryClient.getQueryData(key)).toEqual(afterFirstChunk) - // parallel during streaming - queryClient.setQueryData(['max-chunks-test'], ['parallel1', 'parallel2', 'parallel3']) + await sleep(50) + expect(queryClient.getQueryData(key)).toEqual(['chunk1', 'chunk2']) - await sleep(50) - expect(queryClient.getQueryData(['max-chunks-test'])).toEqual(['parallel3', 'chunk4']) + await sleep(50) + expect(queryClient.getQueryData(key)).toEqual(['chunk2', 'chunk3']) - // parallel during streaming - queryClient.setQueryData(['max-chunks-test'], ['parallel4', 'parallel5', 'parallel6']) + queryClient.setQueryData(key, ['parallel1', 'parallel2', 'parallel3']) - await resultPromise - }) + await sleep(50) + expect(queryClient.getQueryData(key)).toEqual(['parallel3', 'chunk4']) - it('works with reset & parallel', async () => { - const queryFn = serializableStreamedQuery( - async () => { - await sleep(50) + queryClient.setQueryData(key, ['parallel4', 'parallel5', 'parallel6']) - return (async function* () { - await sleep(50) - yield 'chunk1' - await sleep(50) - yield 'chunk2' - await sleep(50) - yield 'chunk3' - await sleep(50) - yield 'chunk4' - await sleep(50) - })() - }, - { maxChunks: 2, refetchMode: 'reset' }, - ) + await resultPromise + }, + ) + }) + describe('abort signal handling', () => { + it('marks the query stale when aborted before the stream connects', async () => { + const key = ['abort-before-connect'] + queryClient.setQueryData(key, ['old']) const controller = new AbortController() + const queryFn = serializableStreamedQuery(async () => { + controller.abort() + return stream(['new'])() + }) - // exceed maxChunks with pre-existing data - should keep only the latest 2 chunks - queryClient.setQueryData(['max-chunks-test'], ['pre1', 'pre2', 'pre3']) - - const resultPromise = expect( - queryFn({ - queryKey: ['max-chunks-test'], - signal: controller.signal, - client: queryClient, - } as any), - ).resolves.toEqual(['parallel5', 'parallel6']) - - await sleep(0) - expect(queryClient.getQueryData(['max-chunks-test'])).toEqual(undefined) - - // exceed while still resolve streaming - queryClient.setQueryData(['max-chunks-test'], ['pre4', 'pre5', 'pre6']) - await sleep(50) - expect(queryClient.getQueryData(['max-chunks-test'])).toEqual(['pre5', 'pre6']) - - await sleep(50) - expect(queryClient.getQueryData(['max-chunks-test'])).toEqual(['pre6', 'chunk1']) - - await sleep(50) - expect(queryClient.getQueryData(['max-chunks-test'])).toEqual(['chunk1', 'chunk2']) - - await sleep(50) - expect(queryClient.getQueryData(['max-chunks-test'])).toEqual(['chunk2', 'chunk3']) - - // parallel during streaming - queryClient.setQueryData(['max-chunks-test'], ['parallel1', 'parallel2', 'parallel3']) - - await sleep(50) - expect(queryClient.getQueryData(['max-chunks-test'])).toEqual(['parallel3', 'chunk4']) - - // parallel during streaming - queryClient.setQueryData(['max-chunks-test'], ['parallel4', 'parallel5', 'parallel6']) - - await resultPromise + await expect( + queryFn({ queryKey: key, signal: controller.signal, client: queryClient } as any), + ).rejects.toBe(controller.signal.reason) + expect(queryClient.getQueryState(key)).toMatchObject({ data: ['old'], isInvalidated: true }) }) - it('works with replace & parallel', async () => { - const queryFn = serializableStreamedQuery( - async () => { - await sleep(50) - - return (async function* () { - await sleep(50) - yield 'chunk1' - await sleep(50) - yield 'chunk2' - await sleep(50) - yield 'chunk3' - await sleep(50) - yield 'chunk4' - await sleep(50) - })() - }, - { maxChunks: 2, refetchMode: 'replace' }, - ) - + it('adds one abort listener per signal across retries', async () => { const controller = new AbortController() - - // exceed maxChunks with pre-existing data - should keep only the latest 2 chunks - queryClient.setQueryData(['max-chunks-test'], ['pre1', 'pre2', 'pre3']) - - const resultPromise = expect( - queryFn({ - queryKey: ['max-chunks-test'], - signal: controller.signal, - client: queryClient, - } as any), - ).resolves.toEqual(['chunk3', 'chunk4']) - - await sleep(0) - expect(queryClient.getQueryData(['max-chunks-test'])).toEqual(['pre2', 'pre3']) - - // exceed while still resolve streaming - queryClient.setQueryData(['max-chunks-test'], ['pre4', 'pre5', 'pre6']) - await sleep(50) - expect(queryClient.getQueryData(['max-chunks-test'])).toEqual(['pre5', 'pre6']) - - await sleep(50) - expect(queryClient.getQueryData(['max-chunks-test'])).toEqual(['pre5', 'pre6']) - - await sleep(50) - expect(queryClient.getQueryData(['max-chunks-test'])).toEqual(['pre5', 'pre6']) - - await sleep(50) - expect(queryClient.getQueryData(['max-chunks-test'])).toEqual(['pre5', 'pre6']) - - // parallel during streaming - queryClient.setQueryData(['max-chunks-test'], ['parallel1', 'parallel2', 'parallel3']) - - await sleep(50) - expect(queryClient.getQueryData(['max-chunks-test'])).toEqual([ - 'parallel1', - 'parallel2', - 'parallel3', - ]) - - // parallel during streaming - queryClient.setQueryData(['max-chunks-test'], ['parallel4', 'parallel5', 'parallel6']) - - await resultPromise - }) - - it('with value Infinite (unlimited)', async () => { - const queryFn = serializableStreamedQuery( - async function* () { - yield 'chunk1' - yield 'chunk2' - yield 'chunk3' - yield 'chunk4' - yield 'chunk5' - }, - { maxChunks: Number.POSITIVE_INFINITY }, - ) - - const result = await queryFn({ - queryKey: ['unlimited-chunks'], - signal: new AbortController().signal, + const addEventListener = vi.spyOn(controller.signal, 'addEventListener') + const queryFn = serializableStreamedQuery(async function* () { + throw new Error('disconnected') + }) + const context = { + queryKey: ['retry-listeners'], + signal: controller.signal, client: queryClient, - } as any) - - expect(result).toEqual(['chunk1', 'chunk2', 'chunk3', 'chunk4', 'chunk5']) - expect(queryClient.getQueryData(['unlimited-chunks'])).toEqual([ - 'chunk1', - 'chunk2', - 'chunk3', - 'chunk4', - 'chunk5', - ]) + } as any + + await expect(queryFn(context)).rejects.toThrow('disconnected') + await expect(queryFn(context)).rejects.toThrow('disconnected') + // listeners removed once their call ends pass their own `signal` + expect( + addEventListener.mock.calls.filter( + ([type, , options]) => type === 'abort' && !(options as AddEventListenerOptions)?.signal, + ), + ).toHaveLength(1) }) - }) - describe('abort signal handling', () => { it('handles abort signal during streaming', async () => { let cleanupCalled = false const yieldFn = vi.fn((v) => v) const queryFn = serializableStreamedQuery(async function* () { try { yield yieldFn('chunk1') - await new Promise((resolve) => setTimeout(resolve, 50)) + await sleep(50) yield yieldFn('chunk2') - await new Promise((resolve) => setTimeout(resolve, 50)) + await sleep(50) yield yieldFn('chunk3') } finally { cleanupCalled = true @@ -512,48 +260,58 @@ describe('serializableStreamedQuery', () => { expect(queryClient.getQueryData(['abort-test'])).toEqual(['chunk1']) }) - it('handles abort signal with replace refetch mode', async () => { - // First query - const queryFn1 = serializableStreamedQuery(async function* () { - yield 'initial1' - yield 'initial2' - }) - - await queryFn1({ - queryKey: ['abort-replace-test'], - signal: new AbortController().signal, - client: queryClient, - } as any) - - expect(queryClient.getQueryData(['abort-replace-test'])).toEqual(['initial1', 'initial2']) - - // Refetch with replace mode and abort - const queryFn2 = serializableStreamedQuery( - async function* () { - yield 'replace1' - await new Promise((resolve) => setTimeout(resolve, 50)) - yield 'replace2' - }, - { refetchMode: 'replace' }, - ) - + it('does not write or start the stream when aborted before it connects', async () => { + const key = ['abort-before-connect-test'] + queryClient.setQueryData(key, ['x']) + const started = vi.fn() const controller = new AbortController() - const resultPromise = expect( - queryFn2({ - queryKey: ['abort-replace-test'], + await expect( + serializableStreamedQuery(async () => { + controller.abort() + return (async function* () { + started() + yield 'chunk1' + })() + })({ + queryKey: key, signal: controller.signal, client: queryClient, } as any), ).rejects.toSatisfy((err) => err === controller.signal.reason) - await new Promise((resolve) => setTimeout(resolve, 25)) - controller.abort() + expect(started).not.toHaveBeenCalled() + expect(queryClient.getQueryState(key)).toMatchObject({ data: ['x'], dataUpdateCount: 1 }) + }) - await resultPromise + it('keeps the chunks received before cancelQueries and writes no more', async () => { + const key = ['cancel-test'] + const cleanup = vi.fn() + const gate = promiseWithResolvers() + const promise = queryClient.fetchQuery({ + queryKey: key, + queryFn: serializableStreamedQuery(async function* () { + try { + yield 'c1' + await gate.promise + yield 'c2' + } finally { + cleanup() + } + }), + }) + + await vi.waitFor(() => expect(queryClient.getQueryData(key)).toEqual(['c1'])) + await queryClient.cancelQueries({ queryKey: key }) + await expect(promise).resolves.toEqual(['c1']) - // Should not replace cache when aborted - expect(queryClient.getQueryData(['abort-replace-test'])).toEqual(['initial1', 'initial2']) + gate.resolve() + await vi.waitFor(() => expect(cleanup).toHaveBeenCalledTimes(1)) + expect(queryClient.getQueryState(key)).toMatchObject({ + status: 'success', + fetchStatus: 'idle', + data: ['c1'], + }) }) }) @@ -577,7 +335,6 @@ describe('serializableStreamedQuery', () => { await sleep(0) expect(queryClient.getQueryData(key)).toEqual(['chunk1']) - // reset data before stream ends void queryClient.resetQueries({ queryKey: key, exact: true }) await promise1 @@ -586,7 +343,7 @@ describe('serializableStreamedQuery', () => { it('handles error in stream', async () => { const queryFn = serializableStreamedQuery(async function* () { yield 'chunk1' - await new Promise((resolve) => setTimeout(resolve, 50)) + await sleep(50) throw new Error('Stream error') }) @@ -598,8 +355,267 @@ describe('serializableStreamedQuery', () => { } as any), ).rejects.toThrow('Stream error') - // Should still have the chunk that was yielded before error expect(queryClient.getQueryData(['error-stream'])).toEqual(['chunk1']) }) }) + + describe('settleOnFirstData option', () => { + it('invalidates only the settled query', async () => { + const closed = promiseWithResolvers() + const settling = serializableStreamedQuery( + async () => ({ + [Symbol.asyncIterator]: () => ({ + next: async () => ({ done: false as const, value: 'x' }), + return: async () => { + await closed.promise + return { done: true as const, value: undefined } + }, + }), + }), + { settleOnFirstData: true }, + ) + + const promise = queryClient.fetchQuery({ queryKey: ['settled'], queryFn: settling }) + await vi.waitFor(() => expect(queryClient.getQueryData(['settled'])).toEqual([]), { + interval: 1, + }) + + await queryClient.fetchQuery({ queryKey: ['other'], queryFn: () => 'other' }) + closed.resolve() + await promise + + expect(queryClient.getQueryState(['settled'])).toMatchObject({ isInvalidated: true }) + expect(queryClient.getQueryState(['other'])).toMatchObject({ isInvalidated: false }) + }) + + it('does not invalidate a later fetch once a settled fetch is cancelled', async () => { + const key = ['settle-then-cancel'] + const closed = promiseWithResolvers() + const settling = serializableStreamedQuery( + async () => ({ + [Symbol.asyncIterator]: () => ({ + next: async () => ({ done: false as const, value: 'x' }), + return: async () => { + await closed.promise + return { done: true as const, value: undefined } + }, + }), + }), + { settleOnFirstData: true }, + ) + + void queryClient.fetchQuery({ queryKey: key, queryFn: settling }).catch(() => {}) + await vi.waitFor(() => expect(queryClient.getQueryData(key)).toEqual([]), { interval: 1 }) + + queryClient + .getQueryCache() + .find({ queryKey: key, exact: true })! + .setOptions({ queryKey: key, queryFn: serializableStreamedQuery(stream(['a'], 1)) }) + await queryClient.refetchQueries({ queryKey: key }, { cancelRefetch: true }) + closed.resolve() + + expect(queryClient.getQueryState(key)).toMatchObject({ data: ['a'], isInvalidated: false }) + }) + + it.each([ + ['reset', []], + ['append', ['x2', 'x3']], + ] as const)( + 'settles at connect, marked stale, when true (%s)', + async (refetchMode, expected) => { + const onSuccess = vi.fn() + const client = new QueryClient({ queryCache: new QueryCache({ onSuccess }) }) + const key = ['settle-test'] + client.setQueryData(key, ['x1', 'x2', 'x3']) + const started = vi.fn() + let signal: AbortSignal | undefined + + await expect( + client.fetchQuery({ + queryKey: key, + queryFn: serializableStreamedQuery( + async (context) => { + signal = context.signal + return (async function* () { + started() + yield 'c1' + })() + }, + { refetchMode, maxChunks: 2, settleOnFirstData: true }, + ), + }), + ).resolves.toEqual(expected) + + expect(client.getQueryState(key)).toMatchObject({ + status: 'success', + fetchStatus: 'idle', + isInvalidated: true, + data: expected, + }) + expect(started).not.toHaveBeenCalled() + // the call that opened the stream is told to stop + expect(signal!.aborted).toBe(true) + expect(onSuccess).toHaveBeenCalledTimes(1) + }, + ) + + it('settles only when no observer watches the query when unobserved', async () => { + const key = ['unobserved-test'] + const options = { + queryKey: key, + queryFn: serializableStreamedQuery(stream([1, 2]), { settleOnFirstData: 'unobserved' }), + } + + await expect(queryClient.fetchQuery(options)).resolves.toEqual([]) + expect(queryClient.getQueryState(key)).toMatchObject({ + fetchStatus: 'idle', + isInvalidated: true, + }) + + const observer = new QueryObserver(queryClient, options) + onTestFinished(observer.subscribe(() => {})) + + await vi.waitFor(() => + expect(observer.getCurrentResult()).toMatchObject({ + data: [1, 2], + fetchStatus: 'idle', + isStale: true, + }), + ) + expect(queryClient.getQueryState(key)).toMatchObject({ isInvalidated: false }) + }) + + it.each([ + [false, [1, 2]], + [true, []], + ] as const)( + "writes to the fetch's own query with a custom queryKeyHashFn (%s)", + async (settleOnFirstData, expected) => { + await expect( + queryClient.fetchQuery({ + queryKey: ['hash-fn-test'], + queryKeyHashFn: () => 'custom', + queryFn: serializableStreamedQuery(stream([1, 2]), { settleOnFirstData }), + }), + ).resolves.toEqual(expected) + + expect( + queryClient + .getQueryCache() + .getAll() + .map((query) => query.queryHash), + ).toEqual(['custom']) + }, + ) + + it('reopens the stream when an observer subscribes while the settled stream closes', async () => { + const key = ['closing-test'] + const closing = promiseWithResolvers() + const fresh = stream([1, 2]) + const connect = vi.fn(async () => + connect.mock.calls.length > 1 + ? fresh() + : { + [Symbol.asyncIterator]: () => ({ + next: async () => ({ done: false, value: 0 }), + return: async () => { + await closing.promise + return { done: true, value: undefined } + }, + }), + }, + ) + const options = { + queryKey: key, + queryFn: serializableStreamedQuery(connect as any, { settleOnFirstData: 'unobserved' }), + } + + void queryClient.prefetchQuery(options) + await vi.waitFor(() => expect(queryClient.getQueryData(key)).toEqual([])) + const observer = new QueryObserver(queryClient, options) + onTestFinished(observer.subscribe(() => {})) + closing.resolve() + + await vi.waitFor(() => + expect(observer.getCurrentResult()).toMatchObject({ data: [1, 2], fetchStatus: 'idle' }), + ) + expect(connect).toHaveBeenCalledTimes(2) + expect(queryClient.getQueryState(key)).toMatchObject({ isInvalidated: false }) + }) + + it('reopens the stream when an observer joins the settled fetch before it is stored', async () => { + const key = ['join-test'] + const connect = vi.fn(stream([1, 2])) + const streamed = serializableStreamedQuery(connect, { settleOnFirstData: 'unobserved' }) + const observer = new QueryObserver(queryClient, { queryKey: key, enabled: false }) + + void queryClient.prefetchQuery({ + queryKey: key, + queryFn: async (context) => { + const data = await streamed(context) + onTestFinished(observer.subscribe(() => {})) + return data + }, + }) + + await vi.waitFor(() => + expect(observer.getCurrentResult()).toMatchObject({ data: [1, 2], fetchStatus: 'idle' }), + ) + expect(connect).toHaveBeenCalledTimes(2) + }) + + it('does not invalidate a later fetch once a settled fetch fails', async () => { + const key = ['settled-error-test'] + const streamed = serializableStreamedQuery(stream([1]), { settleOnFirstData: true }) + + await expect( + queryClient.fetchQuery({ + queryKey: key, + queryFn: async (context) => { + await streamed(context) + throw new Error('after settle') + }, + }), + ).rejects.toThrow('after settle') + + await queryClient.fetchQuery({ queryKey: key, queryFn: () => [2] }) + expect(queryClient.getQueryState(key)).toMatchObject({ data: [2], isInvalidated: false }) + }) + + it('does not write unchanged data at connect (append)', async () => { + const key = ['unchanged-test'] + queryClient.setQueryData(key, ['a']) + const setData = vi.spyOn(queryClient.getQueryCache().find({ queryKey: key })!, 'setData') + + await expect( + queryClient.fetchQuery({ + queryKey: key, + queryFn: serializableStreamedQuery(stream([1]), { + refetchMode: 'append', + settleOnFirstData: true, + }), + }), + ).resolves.toEqual(['a']) + + expect(setData.mock.calls.filter(([, options]) => options?.manual)).toEqual([]) + }) + + it('follows the stream when an observer subscribes before it connects', async () => { + const key = ['unobserved-prefetch-test'] + const connect = vi.fn(stream([1, 2])) + const options = { + queryKey: key, + queryFn: serializableStreamedQuery(connect, { settleOnFirstData: 'unobserved' }), + } + + void queryClient.prefetchQuery(options) + const observer = new QueryObserver(queryClient, options) + onTestFinished(observer.subscribe(() => {})) + + await vi.waitFor(() => + expect(observer.getCurrentResult()).toMatchObject({ data: [1, 2], fetchStatus: 'idle' }), + ) + expect(connect).toHaveBeenCalledTimes(1) + }) + }) }) diff --git a/packages/tanstack-query/src/stream-query.ts b/packages/tanstack-query/src/stream-query.ts index f3148777a..7ae19c460 100644 --- a/packages/tanstack-query/src/stream-query.ts +++ b/packages/tanstack-query/src/stream-query.ts @@ -1,26 +1,11 @@ -import type { Promisable } from '@orpc/shared' +import type { Promisable, StreamedArrayOptions } from '@orpc/shared' +import { limitArraySize, pipeStreamToCache, streamedArrayWrites } from '@orpc/shared' import type { QueryFunction, QueryFunctionContext, QueryKey } from '@tanstack/query-core' -export interface SerializableStreamedQueryOptions { - /** - * Determines how data is handled when the query is fetched again - * - * - `'reset'`: Clears existing data and returns the query to a pending state. - * - `'append'`: Adds new streamed chunks to the existing data. - * - `'replace'`: Buffers streamed data and replaces the cache after the stream completes. - * - * @default 'reset' - */ - refetchMode?: 'append' | 'reset' | 'replace' | undefined +import type { SettleOnFirstDataOptions } from './stream-cache.ts' +import { followStream } from './stream-cache.ts' - /** - * Limits the number of data chunks stored in the query result. - * Older chunks are removed when the limit is reached. - * - * @default Number.POSITIVE_INFINITY (unlimited) - */ - maxChunks?: number | undefined -} +export type SerializableStreamedQueryOptions = StreamedArrayOptions /** * A variant of `streamedQuery` where: @@ -33,79 +18,21 @@ export function serializableStreamedQuery< >( queryFn: (context: QueryFunctionContext) => Promisable>, { - refetchMode = 'reset', - maxChunks = Number.POSITIVE_INFINITY, - }: SerializableStreamedQueryOptions = {}, + settleOnFirstData, + ...options + }: SerializableStreamedQueryOptions & SettleOnFirstDataOptions = {}, ): QueryFunction { - /** - * below code is inspired from old `streamedQuery`: https://github.com/TanStack/query/blob/9973e0f1d82acd6ccf5a49d9a0b5e7e401fc5489/packages/query-core/src/streamedQuery.ts#L19 - */ return async (context) => { - const query = context.client.getQueryCache().find({ queryKey: context.queryKey, exact: true }) - const hasPreviousData = !!query && query.state.data !== undefined - - if (hasPreviousData) { - if (refetchMode === 'reset') { - query.setState({ - status: 'pending', - data: undefined, - error: null, - fetchStatus: 'fetching', - }) - } else { - // not `setQueryData`: trimming is not new data, so it should not report a success - query.setState({ - data: limitArraySize(query.state.data as Array, maxChunks), - }) - } - } - - let result: Array = [] - const stream = await queryFn(context) - const shouldUpdateCacheDuringStream = !hasPreviousData || refetchMode !== 'replace' - - // after resolve stream successfully we can treat result as empty array - context.client.setQueryData>(context.queryKey, (prev = []) => - limitArraySize(prev, maxChunks), + const result = await followStream, TQueryFnData[]>( + context, + settleOnFirstData, + async (callContext, writes) => + pipeStreamToCache(await queryFn(callContext), { + ...writes, + ...streamedArrayWrites(options), + }), ) - // close without reading, so a generator the query no longer wants never starts - if (context.signal.aborted) { - await stream[Symbol.asyncIterator]().return?.() - throw context.signal.reason - } - - for await (const chunk of stream) { - if (context.signal.aborted) { - throw context.signal.reason - } - - result.push(chunk) - result = limitArraySize(result, maxChunks) - - if (shouldUpdateCacheDuringStream) { - context.client.setQueryData>(context.queryKey, (prev = []) => - limitArraySize([...prev, chunk], maxChunks), - ) - } - } - - if (!shouldUpdateCacheDuringStream) { - context.client.setQueryData>(context.queryKey, result) - } - - const cachedData = context.client.getQueryData>(context.queryKey) - if (cachedData) { - return limitArraySize(cachedData, maxChunks) - } - - return result - } -} - -function limitArraySize(items: Array, maxSize: number): Array { - if (items.length <= maxSize) { - return items + return limitArraySize(result!.data, options.maxChunks ?? Number.POSITIVE_INFINITY) } - return items.slice(items.length - maxSize) } diff --git a/packages/tanstack-query/src/types.ts b/packages/tanstack-query/src/types.ts index 795b4c79e..3b6a95369 100644 --- a/packages/tanstack-query/src/types.ts +++ b/packages/tanstack-query/src/types.ts @@ -10,6 +10,7 @@ import type { SkipToken, } from '@tanstack/query-core' +import type { SettleOnFirstDataOptions } from './stream-cache.ts' import type { SerializableStreamedQueryOptions } from './stream-query.ts' export type InferStreamedQueryOutput = TOutput extends AsyncIterable ? U[] : never @@ -83,7 +84,7 @@ export type StreamedOptionsIn< TInitialData, > = QueryOptionsIn & { queryFnOptions?: SerializableStreamedQueryOptions | undefined -} +} & SettleOnFirstDataOptions export type StreamedOptionsOut = QueryOptionsOut< TOutput, @@ -92,6 +93,16 @@ export type StreamedOptionsOut = Que TInitialData > +export type LiveOptionsIn< + TClientContext extends ClientContext, + TInput, + TOutput, + TError, + TSelectData, + TInitialData, +> = QueryOptionsIn & + SettleOnFirstDataOptions + export type InfiniteKeyOptions = | Pick< InfiniteOptionsIn, diff --git a/packages/tanstack-query/tests/__shared__/utils.ts b/packages/tanstack-query/tests/__shared__/utils.ts new file mode 100644 index 000000000..9fa37bd68 --- /dev/null +++ b/packages/tanstack-query/tests/__shared__/utils.ts @@ -0,0 +1,7 @@ +import { stringifyJSON } from '@orpc/shared' + +export function distinctData(items: { data: unknown }[]): unknown[] { + return items + .map((item) => item.data) + .filter((data, i, all) => i === 0 || stringifyJSON(data) !== stringifyJSON(all[i - 1])) +}