fix: restore hydration state after await in <script> (#16806)

* WIP fix more hydration stuff

* fix

* tidy up

* changeset

* fix

* fix

Author: Rich-Harris

.changeset/sweet-doors-brush.md

@@ -0,0 +1,5 @@
+---
+'svelte': patch
+---
+
+fix: restore hydration state after `await` in `<script>`

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

@@ -40,11 +40,7 @@ import { TitleElement } from './visitors/TitleElement.js';
 import { UpdateExpression } from './visitors/UpdateExpression.js';
 import { VariableDeclaration } from './visitors/VariableDeclaration.js';
 import { SvelteBoundary } from './visitors/SvelteBoundary.js';
-import {
-	create_child_block,
-	call_component_renderer,
-	create_async_block
-} from './visitors/shared/utils.js';
+import { call_component_renderer, create_async_block } from './visitors/shared/utils.js';
 
 /** @type {Visitors} */
 const global_visitors = {
@@ -248,7 +244,7 @@ export function server_component(analysis, options) {
 	]);
 
 	if (analysis.instance.has_await) {
-		component_block = b.block([create_child_block(component_block, true)]);
+		component_block = b.block([create_async_block(component_block)]);
 	}
 
 	// trick esrap into including comments

packages/svelte/src/internal/client/reactivity/async.js

@@ -20,6 +20,14 @@ import {
 	set_from_async_derived
 } from './deriveds.js';
 import { aborted } from './effects.js';
+import {
+	hydrate_next,
+	hydrate_node,
+	hydrating,
+	set_hydrate_node,
+	set_hydrating,
+	skip_nodes
+} from '../dom/hydration.js';
 
 /**
  *
@@ -40,6 +48,8 @@ export function flatten(sync, async, fn) {
 
 	var restore = capture();
 
+	var was_hydrating = hydrating;
+
 	Promise.all(async.map((expression) => async_derived(expression)))
 		.then((result) => {
 			batch?.activate();
@@ -55,6 +65,10 @@ export function flatten(sync, async, fn) {
 				}
 			}
 
+			if (was_hydrating) {
+				set_hydrating(false);
+			}
+
 			batch?.deactivate();
 			unset_context();
 		})
@@ -74,12 +88,23 @@ function capture() {
 	var previous_component_context = component_context;
 	var previous_batch = current_batch;
 
+	var was_hydrating = hydrating;
+
+	if (was_hydrating) {
+		var previous_hydrate_node = hydrate_node;
+	}
+
 	return function restore() {
 		set_active_effect(previous_effect);
 		set_active_reaction(previous_reaction);
 		set_component_context(previous_component_context);
 		previous_batch?.activate();
 
+		if (was_hydrating) {
+			set_hydrating(true);
+			set_hydrate_node(previous_hydrate_node);
+		}
+
 		if (DEV) {
 			set_from_async_derived(null);
 		}
@@ -186,13 +211,34 @@ export async function async_body(fn) {
 
 	var active = /** @type {Effect} */ (active_effect);
 
+	var was_hydrating = hydrating;
+	var next_hydrate_node = undefined;
+
+	if (was_hydrating) {
+		hydrate_next();
+		next_hydrate_node = skip_nodes(false);
+	}
+
 	try {
-		await fn();
+		var promise = fn();
+	} finally {
+		if (next_hydrate_node) {
+			set_hydrate_node(next_hydrate_node);
+			hydrate_next();
+		}
+	}
+
+	try {
+		await promise;
 	} catch (error) {
 		if (!aborted(active)) {
 			invoke_error_boundary(error, active);
 		}
 	} finally {
+		if (was_hydrating) {
+			set_hydrating(false);
+		}
+
 		boundary.update_pending_count(-1);
 
 		if (pending) {

packages/svelte/tests/runtime-runes/samples/async-no-pending-await-in-script-2/Child.svelte

@@ -0,0 +1,3 @@
+<script lang="ts">
+	await 1;
+</script>

packages/svelte/tests/runtime-runes/samples/async-no-pending-await-in-script-2/_config.js

@@ -0,0 +1,13 @@
+import { tick } from 'svelte';
+import { test } from '../../test';
+
+export default test({
+	skip_mode: ['server'],
+
+	ssrHtml: '<p>hello</p>',
+
+	async test({ assert, target }) {
+		await tick();
+		assert.htmlEqual(target.innerHTML, '<p>hello</p>');
+	}
+});

packages/svelte/tests/runtime-runes/samples/async-no-pending-await-in-script-2/main.svelte

@@ -0,0 +1,9 @@
+<script>
+	import Child from './Child.svelte';
+</script>
+
+<Child />
+
+{#if true}
+	<p>hello</p>
+{/if}

packages/svelte/tests/runtime-runes/samples/async-no-pending-await-in-script/_config.js

@@ -0,0 +1,13 @@
+import { tick } from 'svelte';
+import { test } from '../../test';
+
+export default test({
+	skip_mode: ['server'],
+
+	ssrHtml: '<p>hello</p>',
+
+	async test({ assert, target }) {
+		await tick();
+		assert.htmlEqual(target.innerHTML, '<p>hello</p>');
+	}
+});

packages/svelte/tests/runtime-runes/samples/async-no-pending-await-in-script/main.svelte

@@ -0,0 +1,7 @@
+<script lang="ts">
+	await 1;
+</script>
+
+{#if true}
+	<p>hello</p>
+{/if}