fix: correct query .set and .refresh behavior in commands

.changeset/swift-fans-check.md

@@ -0,0 +1,5 @@
+---
+'@sveltejs/kit': patch
+---
+
+fix: correct query `.set` and `.refresh` behavior in commands

packages/kit/src/runtime/app/server/remote/query.js

@@ -72,46 +72,21 @@ export function query(validate_or_fn, maybe_fn) {
 
 		const { event, state } = get_request_store();
 
+		const get_remote_function_result = () =>
+			run_remote_function(event, state, false, arg, validate, fn);
+
 		/** @type {Promise<any> & Partial<RemoteQuery<any>>} */
-		const promise = get_response(__, arg, state, () =>
-			run_remote_function(event, state, false, arg, validate, fn)
-		);
+		const promise = get_response(__, arg, state, get_remote_function_result);
 
 		promise.catch(() => {});
 
-		/** @param {Output} value */
-		promise.set = (value) => {
-			const { state } = get_request_store();
-			const refreshes = state.refreshes;
-
-			if (!refreshes) {
-				throw new Error(
-					`Cannot call set on query '${__.name}' because it is not executed in the context of a command/form remote function`
-				);
-			}
-
-			if (__.id) {
-				const cache = get_cache(__, state);
-				const key = stringify_remote_arg(arg, state.transport);
-				refreshes[create_remote_cache_key(__.id, key)] = cache[key] = Promise.resolve(value);
-			}
-		};
+		promise.set = (value) => update_refresh_value(get_refresh_context(__, 'set', arg), value);
 
 		promise.refresh = () => {
-			const { state } = get_request_store();
-			const refreshes = state.refreshes;
-
-			if (!refreshes) {
-				throw new Error(
-					`Cannot call refresh on query '${__.name}' because it is not executed in the context of a command/form remote function`
-				);
-			}
-
-			const cache_key = create_remote_cache_key(__.id, stringify_remote_arg(arg, state.transport));
-			refreshes[cache_key] = promise;
-
-			// TODO we could probably just return promise here, but would need to update the types
-			return promise.then(() => {});
+			const refresh_context = get_refresh_context(__, 'refresh', arg);
+			const is_immediate_refresh = !refresh_context.cache[refresh_context.key];
+			const value = is_immediate_refresh ? promise : get_remote_function_result();
+			return update_refresh_value(refresh_context, value, is_immediate_refresh);
 		};
 
 		promise.withOverride = () => {
@@ -200,8 +175,7 @@ function batch(validate_or_fn, maybe_fn) {
 
 		const { event, state } = get_request_store();
 
-		/** @type {Promise<any> & Partial<RemoteQuery<any>>} */
-		const promise = get_response(__, arg, state, () => {
+		const get_remote_function_result = () => {
 			// Collect all the calls to the same query in the same macrotask,
 			// then execute them as one backend request.
 			return new Promise((resolve, reject) => {
@@ -239,22 +213,20 @@ function batch(validate_or_fn, maybe_fn) {
 					}
 				}, 0);
 			});
-		});
+		};
 
-		promise.catch(() => {});
+		/** @type {Promise<any> & Partial<RemoteQuery<any>>} */
+		const promise = get_response(__, arg, state, get_remote_function_result);
 
-		promise.refresh = async () => {
-			const { state } = get_request_store();
-			const refreshes = state.refreshes;
+		promise.catch(() => {});
 
-			if (!refreshes) {
-				throw new Error(
-					`Cannot call refresh on query.batch '${__.name}' because it is not executed in the context of a command/form remote function`
-				);
-			}
+		promise.set = (value) => update_refresh_value(get_refresh_context(__, 'set', arg), value);
 
-			const cache_key = create_remote_cache_key(__.id, stringify_remote_arg(arg, state.transport));
-			refreshes[cache_key] = await /** @type {Promise<any>} */ (promise);
+		promise.refresh = () => {
+			const refresh_context = get_refresh_context(__, 'refresh', arg);
+			const is_immediate_refresh = !refresh_context.cache[refresh_context.key];
+			const value = is_immediate_refresh ? promise : get_remote_function_result();
+			return update_refresh_value(refresh_context, value, is_immediate_refresh);
 		};
 
 		promise.withOverride = () => {
@@ -271,3 +243,57 @@ function batch(validate_or_fn, maybe_fn) {
 
 // Add batch as a property to the query function
 Object.defineProperty(query, 'batch', { value: batch, enumerable: true });
+
+/**
+ * @param {RemoteInfo} __
+ * @param {'set' | 'refresh'} action
+ * @param {any} [arg]
+ * @returns {{ __: RemoteInfo; state: any; refreshes: Record<string, Promise<any>>; cache: Record<string, Promise<any>>; cache_key: string; key: string }}
+ */
+function get_refresh_context(__, action, arg) {
+	const { state } = get_request_store();
+	const { refreshes } = state;
+
+	if (!refreshes) {
+		throw new Error(
+			`Cannot call ${action} on ${format_remote_name(__)} because it is not executed in the context of a command/form remote function`
+		);
+	}
+
+	const key = stringify_remote_arg(arg, state.transport);
+	const cache_key = create_remote_cache_key(__.id, key);
+	const cache = get_cache(__, state);
+
+	return { __, state, refreshes, cache, cache_key, key };
+}
+
+/**
+ * @param {{ __: RemoteInfo; refreshes: Record<string, Promise<any>>; cache: Record<string, Promise<any>>; cache_key: string; key: string }} context
+ * @param {any} value
+ * @param {boolean} [is_immediate_refresh=false]
+ * @returns {Promise<void>}
+ */
+function update_refresh_value(
+	{ __, refreshes, cache, cache_key, key },
+	value,
+	is_immediate_refresh = false
+) {
+	const promise = Promise.resolve(value);
+
+	if (!is_immediate_refresh) {
+		cache[key] = promise;
+	}
+
+	if (__.id) {
+		refreshes[cache_key] = promise;
+	}
+
+	return promise.then(() => {});
+}
+
+/**
+ * @param {RemoteInfo} __
+ */
+function format_remote_name(__) {
+	return __.type === 'query_batch' ? `query.batch '${__.name}'` : `query '${__.name}'`;
+}

packages/kit/src/runtime/app/server/remote/shared.js

@@ -69,7 +69,7 @@ export function create_validator(validate_or_fn, maybe_fn) {
  * @returns {Promise<T>}
  */
 export async function get_response(info, arg, state, get_result) {
-	// wait a beat, in case `myQuery().set(...)` is immediately called
+	// wait a beat, in case `myQuery().set(...)` or `myQuery().refresh()` is immediately called
 	// eslint-disable-next-line @typescript-eslint/await-thenable
 	await 0;
 

packages/kit/test/apps/basics/src/routes/remote/+page.svelte

@@ -6,6 +6,7 @@
 		get_count,
 		set_count,
 		set_count_server_refresh,
+		set_count_server_refresh_after_read,
 		set_count_server_set,
 		resolve_deferreds
 	} from './query-command.remote.js';
@@ -67,6 +68,14 @@
 >
 	command (query server refresh)
 </button>
+<button
+	onclick={async () => {
+		command_result = await set_count_server_refresh_after_read(6);
+	}}
+	id="multiply-server-refresh-after-read-btn"
+>
+	command (query server refresh after read)
+</button>
 <button
 	onclick={async () => {
 		// slow, else test will not be able to see the override

packages/kit/test/apps/basics/src/routes/remote/batch/+page.svelte

@@ -1,15 +1,21 @@
 <script>
-	import { get_todo } from './batch.remote.js';
+	import {
+		get_todo,
+		set_todo_title_server_refresh,
+		reset_todos,
+		set_todo_title
+	} from './batch.remote.js';
 
 	const todoIds = ['1', '2', '1', 'error'];
+	const todos = todoIds.map((id) => ({ id, promise: get_todo(id) }));
 </script>
 
 <h1>Query Batch Test</h1>
 
 <ul>
-	{#each todoIds as id, idx}
+	{#each todos as { id, promise }, idx (idx)}
 		<li>
-			{#await get_todo(id)}
+			{#await promise}
 				<span id="batch-result-{idx + 1}">Loading todo {id}...</span>
 			{:then todo}
 				<span id="batch-result-{idx + 1}">{todo.title}</span>
@@ -21,3 +27,27 @@
 </ul>
 
 <button onclick={() => todoIds.forEach((id) => get_todo(id).refresh())}>refresh</button>
+<button
+	onclick={async () => {
+		await set_todo_title({ id: '1', title: 'Buy cat food' });
+	}}
+	id="batch-set-btn"
+>
+	set first todo
+</button>
+<button
+	onclick={async () => {
+		await set_todo_title_server_refresh({ id: '2', title: 'Walk the dog (refreshed)' });
+	}}
+	id="batch-refresh-btn"
+>
+	refresh second todo via command
+</button>
+<button
+	onclick={async () => {
+		await reset_todos();
+	}}
+	id="batch-reset-btn"
+>
+	reset todos
+</button>

packages/kit/test/apps/basics/src/routes/remote/batch/batch.remote.js

@@ -1,15 +1,43 @@
-import { query } from '$app/server';
+import { command, query } from '$app/server';
 import { error } from '@sveltejs/kit';
 
+/** @type {Array<[string, { id: string; title: string }]>} **/
+const INITIAL_TODOS = [
+	['1', { id: '1', title: 'Buy groceries' }],
+	['2', { id: '2', title: 'Walk the dog' }]
+];
+
+let todos = new Map(INITIAL_TODOS);
+
 export const get_todo = query.batch('unchecked', (ids) => {
-	if (JSON.stringify(ids) !== JSON.stringify(['1', '2', 'error'])) {
-		throw new Error(`Expected 3 IDs (deduplicated), got ${JSON.stringify(ids)}`);
+	if (new Set(ids).size !== ids.length) {
+		throw new Error(`batch queries must be deduplicated, but got ${JSON.stringify(ids)}`);
 	}
 
-	return (id) =>
-		id === '1'
-			? { id: '1', title: 'Buy groceries' }
-			: id === '2'
-				? { id: '2', title: 'Walk the dog' }
-				: error(404, { message: 'Not found' });
+	return (id) => {
+		if (id === 'error') return error(404, { message: 'Not found' });
+
+		return todos.get(id);
+	};
+});
+
+export const set_todo_title = command('unchecked', ({ id, title }) => {
+	const todo = { id, title };
+	todos.set(id, todo);
+	get_todo(id).set(todo);
+	return todo;
+});
+
+export const set_todo_title_server_refresh = command('unchecked', ({ id, title }) => {
+	const todo = { id, title };
+	todos.set(id, todo);
+	get_todo(id).refresh();
+	return todo;
+});
+
+export const reset_todos = command(() => {
+	todos = new Map(INITIAL_TODOS);
+	for (const [id, todo] of todos) {
+		get_todo(id).set({ ...todo });
+	}
 });

packages/kit/test/apps/basics/src/routes/remote/query-command.remote.js

@@ -36,6 +36,13 @@ export const set_count_server_refresh = command('unchecked', (c) => {
 	return c;
 });
 
+export const set_count_server_refresh_after_read = command('unchecked', async (c) => {
+	await get_count();
+	count = c;
+	await get_count().refresh();
+	return c;
+});
+
 export const set_count_server_set = command('unchecked', async (c) => {
 	get_count_called = false;
 	count = c;

packages/kit/test/apps/basics/test/client.test.js

@@ -1860,6 +1860,20 @@ test.describe('remote functions', () => {
 		expect(request_count).toBe(1); // no query refreshes, since that happens as part of the command response
 	});
 
+	test('command refresh after reading query reruns the query', async ({ page }) => {
+		await page.goto('/remote');
+		await expect(page.locator('#count-result')).toHaveText('0 / 0 (false)');
+
+		let request_count = 0;
+		page.on('request', (r) => (request_count += r.url().includes('/_app/remote') ? 1 : 0));
+
+		await page.click('#multiply-server-refresh-after-read-btn');
+		await expect(page.locator('#command-result')).toHaveText('6');
+		await expect(page.locator('#count-result')).toHaveText('6 / 6 (false)');
+		await page.waitForTimeout(100); // allow all requests to finish (in case there are query refreshes which shouldn't happen)
+		expect(request_count).toBe(1);
+	});
+
 	test('command does server-initiated single flight mutation (set)', async ({ page }) => {
 		await page.goto('/remote');
 		await expect(page.locator('#count-result')).toHaveText('0 / 0 (false)');
@@ -2027,6 +2041,36 @@ test.describe('remote functions', () => {
 		expect(request_count).toBe(1);
 	});
 
+	test('query.batch set updates cache without extra request', async ({ page }) => {
+		await page.goto('/remote/batch');
+		await page.click('#batch-reset-btn');
+		await expect(page.locator('#batch-result-1')).toHaveText('Buy groceries');
+
+		let request_count = 0;
+		const handler = (r) => (request_count += r.url().includes('/_app/remote') ? 1 : 0);
+		page.on('request', handler);
+
+		await page.click('#batch-set-btn');
+		await expect(page.locator('#batch-result-1')).toHaveText('Buy cat food');
+		await page.waitForTimeout(100); // allow all requests to finish
+		expect(request_count).toBe(1); // only the command request
+	});
+
+	test('query.batch refresh in command reuses single flight', async ({ page }) => {
+		await page.goto('/remote/batch');
+		await page.click('#batch-reset-btn');
+		await expect(page.locator('#batch-result-2')).toHaveText('Walk the dog');
+
+		let request_count = 0;
+		const handler = (r) => (request_count += r.url().includes('/_app/remote') ? 1 : 0);
+		page.on('request', handler);
+
+		await page.click('#batch-refresh-btn');
+		await expect(page.locator('#batch-result-2')).toHaveText('Walk the dog (refreshed)');
+		await page.waitForTimeout(100); // allow all requests to finish
+		expect(request_count).toBe(1); // only the command request
+	});
+
 	// TODO ditto
 	test('query works with transport', async ({ page }) => {
 		await page.goto('/remote/transport');