fix: preserve SSR context when block expressions contain `await` (#16791)

* DRY out

* fix: preserve SSR context when block expressions contain `await`

.changeset/friendly-months-prove.md

@@ -0,0 +1,5 @@
+---
+'svelte': patch
+---
+
+fix: preserve SSR context when block expressions contain `await`

packages/svelte/src/compiler/phases/3-transform/client/utils.js

@@ -4,7 +4,7 @@
 /** @import { Analysis } from '../../types.js' */
 /** @import { Scope } from '../../scope.js' */
 import * as b from '#compiler/builders';
-import { is_simple_expression } from '../../../utils/ast.js';
+import { is_simple_expression, save } from '../../../utils/ast.js';
 import {
 	PROPS_IS_LAZY_INITIAL,
 	PROPS_IS_IMMUTABLE,
@@ -296,7 +296,7 @@ export function create_derived(state, expression, async = false) {
 	const thunk = b.thunk(expression, async);
 
 	if (async) {
-		return b.call(b.await(b.call('$.save', b.call('$.async_derived', thunk))));
+		return save(b.call('$.async_derived', thunk));
 	} else {
 		return b.call(state.analysis.runes ? '$.derived' : '$.derived_safe_equal', thunk);
 	}

packages/svelte/src/compiler/phases/3-transform/client/visitors/AwaitExpression.js

@@ -1,6 +1,7 @@
 /** @import { AwaitExpression, Expression } from 'estree' */
 /** @import { Context } from '../types' */
 import { dev, is_ignored } from '../../../../state.js';
+import { save } from '../../../../utils/ast.js';
 import * as b from '../../../../utils/builders.js';
 
 /**
@@ -11,7 +12,7 @@ export function AwaitExpression(node, context) {
 	const argument = /** @type {Expression} */ (context.visit(node.argument));
 
 	if (context.state.analysis.pickled_awaits.has(node)) {
-		return b.call(b.await(b.call('$.save', argument)));
+		return save(argument);
 	}
 
 	// in dev, note which values are read inside a reactive expression,

packages/svelte/src/compiler/phases/3-transform/client/visitors/VariableDeclaration.js

@@ -2,7 +2,7 @@
 /** @import { Binding } from '#compiler' */
 /** @import { ComponentContext } from '../types' */
 import { dev, is_ignored, locate_node } from '../../../../state.js';
-import { extract_paths } from '../../../../utils/ast.js';
+import { extract_paths, save } from '../../../../utils/ast.js';
 import * as b from '#compiler/builders';
 import * as assert from '../../../../utils/assert.js';
 import { get_rune } from '../../../scope.js';
@@ -212,7 +212,7 @@ export function VariableDeclaration(node, context) {
 							location ? b.literal(location) : undefined
 						);
 
-						call = b.call(b.await(b.call('$.save', call)));
+						call = save(call);
 						if (dev) call = b.call('$.tag', call, b.literal(declarator.id.name));
 
 						declarations.push(b.declarator(declarator.id, call));
@@ -243,7 +243,7 @@ export function VariableDeclaration(node, context) {
 								b.thunk(expression, true),
 								location ? b.literal(location) : undefined
 							);
-							call = b.call(b.await(b.call('$.save', call)));
+							call = save(call);
 						}
 
 						if (dev) {

packages/svelte/src/compiler/phases/3-transform/server/visitors/AwaitExpression.js

@@ -1,6 +1,6 @@
 /** @import { AwaitExpression, Expression } from 'estree' */
 /** @import { Context } from '../types' */
-import * as b from '../../../../utils/builders.js';
+import { save } from '../../../../utils/ast.js';
 
 /**
  * @param {AwaitExpression} node
@@ -10,7 +10,30 @@ export function AwaitExpression(node, context) {
 	const argument = /** @type {Expression} */ (context.visit(node.argument));
 
 	if (context.state.analysis.pickled_awaits.has(node)) {
-		return b.call(b.await(b.call('$.save', argument)));
+		return save(argument);
+	}
+
+	// we also need to restore context after block expressions
+	let i = context.path.length;
+	while (i--) {
+		const parent = context.path[i];
+
+		if (
+			parent.type === 'ArrowFunctionExpression' ||
+			parent.type === 'FunctionExpression' ||
+			parent.type === 'FunctionDeclaration'
+		) {
+			break;
+		}
+
+		// @ts-ignore
+		if (parent.metadata) {
+			if (parent.type !== 'ExpressionTag' && parent.type !== 'Fragment') {
+				return save(argument);
+			}
+
+			break;
+		}
 	}
 
 	return argument === node.argument ? node : { ...node, argument };

packages/svelte/src/compiler/utils/ast.js

@@ -625,3 +625,11 @@ export function has_await(expression) {
 
 	return has_await;
 }
+
+/**
+ * Turns `await ...` to `(await $.save(...))()`
+ * @param {ESTree.Expression} expression
+ */
+export function save(expression) {
+	return b.call(b.await(b.call('$.save', expression)));
+}

packages/svelte/tests/runtime-runes/samples/async-push-element-dev/_config.js

@@ -0,0 +1,15 @@
+import { test } from '../../test';
+
+export default test({
+	mode: ['async-server'],
+
+	compileOptions: {
+		// include `push_element` calls, so that we can check they
+		// run with the correct ssr_context
+		dev: true
+	},
+
+	html: `
+		<h1>hello!</h1>
+	`
+});

packages/svelte/tests/runtime-runes/samples/async-push-element-dev/main.svelte

@@ -0,0 +1,3 @@
+{#if await true}
+	<h1>hello!</h1>
+{/if}

packages/svelte/tests/snapshot/samples/async-each-fallback-hoisting/_expected/server/index.svelte.js

@@ -3,7 +3,7 @@ import * as $ from 'svelte/internal/server';
 
 export default function Async_each_fallback_hoisting($$renderer) {
 	 $$renderer.child(async ($$renderer) => {
-		const each_array = $.ensure_array_like(await Promise.resolve([]));
+		const each_array = $.ensure_array_like((await $.save(Promise.resolve([])))());
 
 		if (each_array.length !== 0) {
 			 $$renderer.push('<!--[-->');

packages/svelte/tests/snapshot/samples/async-each-hoisting/_expected/server/index.svelte.js

@@ -9,7 +9,7 @@ export default function Async_each_hoisting($$renderer) {
 	 $$renderer.push(`<!--[-->`);
 
 	 $$renderer.child(async ($$renderer) => {
-		const each_array = $.ensure_array_like(await Promise.resolve([first, second, third]));
+		const each_array = $.ensure_array_like((await $.save(Promise.resolve([first, second, third])))());
 
 		for (let $$index = 0, $$length = each_array.length; $$index < $$length; $$index++) {
 			let item = each_array[$$index];

packages/svelte/tests/snapshot/samples/async-if-alternate-hoisting/_expected/server/index.svelte.js

@@ -3,7 +3,7 @@ import * as $ from 'svelte/internal/server';
 
 export default function Async_if_alternate_hoisting($$renderer) {
 	 $$renderer.child(async ($$renderer) => {
-		if (await Promise.resolve(false)) {
+		if ((await $.save(Promise.resolve(false)))()) {
 			 $$renderer.push('<!--[-->');
 			 $$renderer.push(async () => $.escape(await Promise.reject('no no no')));
 		} else {

packages/svelte/tests/snapshot/samples/async-if-hoisting/_expected/server/index.svelte.js

@@ -3,7 +3,7 @@ import * as $ from 'svelte/internal/server';
 
 export default function Async_if_hoisting($$renderer) {
 	 $$renderer.child(async ($$renderer) => {
-		if (await Promise.resolve(true)) {
+		if ((await $.save(Promise.resolve(true)))()) {
 			 $$renderer.push('<!--[-->');
 			 $$renderer.push(async () => $.escape(await Promise.resolve('yes yes yes')));
 		} else {