fix: correct query .set and .refresh behavior in commands (#14877)

Closes #14832
Closes #14602

---------

Co-authored-by: Simon H <5968653+dummdidumm@users.noreply.github.com>
Co-authored-by: Simon Holthausen <simon.holthausen@vercel.com>

Author: phi-bre

.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

@@ -2,7 +2,7 @@
 /** @import { RemoteInfo, MaybePromise } from 'types' */
 /** @import { StandardSchemaV1 } from '@standard-schema/spec' */
 import { get_request_store } from '@sveltejs/kit/internal/server';
-import { create_remote_cache_key, stringify_remote_arg } from '../../../shared.js';
+import { create_remote_key, stringify_remote_arg } from '../../../shared.js';
 import { prerendering } from '__sveltekit/environment';
 import { create_validator, get_cache, get_response, run_remote_function } from './shared.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.cache_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.cache_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,51 @@ 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>>; refreshes_key: string; cache_key: string }}
+ */
+function get_refresh_context(__, action, arg) {
+	const { state } = get_request_store();
+	const { refreshes } = state;
+
+	if (!refreshes) {
+		const name = __.type === 'query_batch' ? `query.batch '${__.name}'` : `query '${__.name}'`;
+		throw new Error(
+			`Cannot call ${action} on ${name} because it is not executed in the context of a command/form remote function`
+		);
+	}
+
+	const cache = get_cache(__, state);
+	const cache_key = stringify_remote_arg(arg, state.transport);
+	const refreshes_key = create_remote_key(__.id, cache_key);
+
+	return { __, state, refreshes, refreshes_key, cache, cache_key };
+}
+
+/**
+ * @param {{ __: RemoteInfo; refreshes: Record<string, Promise<any>>; cache: Record<string, Promise<any>>; refreshes_key: string; cache_key: string }} context
+ * @param {any} value
+ * @param {boolean} [is_immediate_refresh=false]
+ * @returns {Promise<void>}
+ */
+function update_refresh_value(
+	{ __, refreshes, refreshes_key, cache, cache_key },
+	value,
+	is_immediate_refresh = false
+) {
+	const promise = Promise.resolve(value);
+
+	if (!is_immediate_refresh) {
+		cache[cache_key] = promise;
+	}
+
+	if (__.id) {
+		refreshes[refreshes_key] = promise;
+	}
+
+	return promise.then(() => {});
+}

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

@@ -68,7 +68,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/src/runtime/client/remote-functions/shared.svelte.js

@@ -5,7 +5,7 @@ import * as devalue from 'devalue';
 import { app, goto, query_map, remote_responses } from '../client.js';
 import { HttpError, Redirect } from '@sveltejs/kit/internal';
 import { tick } from 'svelte';
-import { create_remote_cache_key, stringify_remote_arg } from '../../shared.js';
+import { create_remote_key, stringify_remote_arg } from '../../shared.js';
 
 /**
  *
@@ -48,7 +48,7 @@ export async function remote_request(url) {
 export function create_remote_function(id, create) {
 	return (/** @type {any} */ arg) => {
 		const payload = stringify_remote_arg(arg, app.hooks.transport);
-		const cache_key = create_remote_cache_key(id, payload);
+		const cache_key = create_remote_key(id, payload);
 		let entry = query_map.get(cache_key);
 
 		let tracking = true;

packages/kit/src/runtime/server/page/render.js

@@ -16,7 +16,7 @@ import { add_resolution_suffix } from '../../pathname.js';
 import { try_get_request_store, with_request_store } from '@sveltejs/kit/internal/server';
 import { text_encoder } from '../../utils.js';
 import { get_global_name } from '../utils.js';
-import { create_remote_cache_key } from '../../shared.js';
+import { create_remote_key } from '../../shared.js';
 
 // TODO rename this function/module
 
@@ -493,7 +493,7 @@ export async function render_response({
 				if (!info.id) continue;
 
 				for (const key in cache) {
-					remote[create_remote_cache_key(info.id, key)] = await cache[key];
+					remote[create_remote_key(info.id, key)] = await cache[key];
 				}
 			}
 

packages/kit/src/runtime/shared.js

@@ -88,6 +88,6 @@ export function parse_remote_arg(string, transport) {
  * @param {string} id
  * @param {string} payload
  */
-export function create_remote_cache_key(id, payload) {
+export function create_remote_key(id, payload) {
 	return id + '/' + payload;
 }

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,23 @@
 <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'];
+	// Need to write this outside at the top level to ensure tests succeed in non-async-mode
+	// Else updates are not coming through properly because of state-created-inside-effects-not-updating logic in non-async mode
+	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 +29,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;