From 021565c1711ede7024e8066cb18083b30a599ca4 Mon Sep 17 00:00:00 2001
From: Evan Almloff <evanalmloff@gmail.com>
Date: Tue, 19 Aug 2025 14:23:45 -0500
Subject: [PATCH] Fix hydration of link element (#4561)

* Fix hydration of link element

* revert some logging changes

* clean up fullstack history provider

* fix router tests

* fix document context type
---
 Cargo.lock                                    |  2 +
 packages/dioxus/src/launch.rs                 | 36 -------
 packages/fullstack-hooks/Cargo.toml           |  1 +
 .../src/document.rs}                          | 27 +++--
 packages/fullstack-hooks/src/history.rs       | 98 ++++++++++++-------
 packages/fullstack-hooks/src/lib.rs           |  1 +
 packages/fullstack-hooks/src/streaming.rs     |  1 +
 packages/fullstack/src/lib.rs                 |  8 +-
 packages/server/src/render.rs                 | 18 ++--
 packages/web/Cargo.toml                       |  3 +-
 packages/web/src/document.rs                  | 48 +++++++--
 packages/web/src/lib.rs                       |  7 +-
 12 files changed, 141 insertions(+), 109 deletions(-)
 rename packages/{fullstack/src/web.rs => fullstack-hooks/src/document.rs} (58%)

diff --git a/Cargo.lock b/Cargo.lock
index 4dbc1ff42..1ba4945bc 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -5533,6 +5533,7 @@ version = "0.7.0-rc.0"
 dependencies = [
  "dioxus",
  "dioxus-core",
+ "dioxus-document",
  "dioxus-fullstack",
  "dioxus-fullstack-protocol",
  "dioxus-history",
@@ -6024,6 +6025,7 @@ dependencies = [
  "dioxus-core-types",
  "dioxus-devtools",
  "dioxus-document",
+ "dioxus-fullstack-hooks",
  "dioxus-fullstack-protocol",
  "dioxus-history",
  "dioxus-html",
diff --git a/packages/dioxus/src/launch.rs b/packages/dioxus/src/launch.rs
index 4b6464ad7..6eb24641d 100644
--- a/packages/dioxus/src/launch.rs
+++ b/packages/dioxus/src/launch.rs
@@ -347,42 +347,6 @@ impl LaunchBuilder {
 
         #[cfg(feature = "web")]
         if matches!(platform, KnownPlatform::Web) {
-            // If the server feature is enabled, launch the client with hydration enabled
-            #[cfg(feature = "fullstack")]
-            {
-                let platform_config = configs
-                    .into_iter()
-                    .find_map(|cfg| cfg.downcast::<dioxus_web::Config>().ok())
-                    .unwrap_or_default()
-                    .hydrate(true);
-
-                let mut vdom = dioxus_core::VirtualDom::new(app);
-
-                #[cfg(feature = "document")]
-                {
-                    use dioxus_fullstack::FullstackWebDocument;
-                    let document = std::rc::Rc::new(FullstackWebDocument)
-                        as std::rc::Rc<dyn crate::prelude::document::Document>;
-                    vdom.provide_root_context(document);
-                }
-
-                #[cfg(feature = "document")]
-                {
-                    use dioxus_fullstack::FullstackHistory;
-                    let history =
-                        std::rc::Rc::new(FullstackHistory::new(dioxus_web::WebHistory::default()))
-                            as std::rc::Rc<dyn crate::prelude::History>;
-                    vdom.provide_root_context(history);
-                }
-
-                for context in contexts {
-                    vdom.insert_any_root_context(context());
-                }
-
-                return dioxus_web::launch::launch_virtual_dom(vdom, platform_config);
-            }
-
-            #[cfg(not(any(feature = "fullstack")))]
             return dioxus_web::launch::launch(app, contexts, configs);
         }
 
diff --git a/packages/fullstack-hooks/Cargo.toml b/packages/fullstack-hooks/Cargo.toml
index 31269a58d..6e659767d 100644
--- a/packages/fullstack-hooks/Cargo.toml
+++ b/packages/fullstack-hooks/Cargo.toml
@@ -18,6 +18,7 @@ dioxus-fullstack-protocol = { workspace = true }
 futures-channel = { workspace = true }
 serde = { workspace = true }
 dioxus-history.workspace = true
+dioxus-document.workspace = true
 
 [dev-dependencies]
 dioxus-fullstack = { workspace = true }
diff --git a/packages/fullstack/src/web.rs b/packages/fullstack-hooks/src/document.rs
similarity index 58%
rename from packages/fullstack/src/web.rs
rename to packages/fullstack-hooks/src/document.rs
index c1479cf6f..570877b12 100644
--- a/packages/fullstack/src/web.rs
+++ b/packages/fullstack-hooks/src/document.rs
@@ -1,7 +1,6 @@
-//! On the client, we use the [`WebDocument`] implementation to render the head for any elements that were not rendered on the server.
+//! On the client, we use the [`FullstackWebDocument`] implementation to render the head for any elements that were not rendered on the server.
 
 use dioxus_document::*;
-use dioxus_web::WebDocument;
 
 fn head_element_written_on_server() -> bool {
     dioxus_fullstack_protocol::head_element_hydration_entry()
@@ -12,36 +11,44 @@ fn head_element_written_on_server() -> bool {
 
 /// A document provider for fullstack web clients
 #[derive(Clone)]
-pub struct FullstackWebDocument;
+pub struct FullstackWebDocument<D> {
+    document: D,
+}
 
-impl Document for FullstackWebDocument {
+impl<D> From<D> for FullstackWebDocument<D> {
+    fn from(document: D) -> Self {
+        Self { document }
+    }
+}
+
+impl<D: Document> Document for FullstackWebDocument<D> {
     fn eval(&self, js: String) -> Eval {
-        WebDocument.eval(js)
+        self.document.eval(js)
     }
 
     /// Set the title of the document
     fn set_title(&self, title: String) {
-        WebDocument.set_title(title);
+        self.document.set_title(title);
     }
 
     /// Create a new meta tag in the head
     fn create_meta(&self, props: MetaProps) {
-        WebDocument.create_meta(props);
+        self.document.create_meta(props);
     }
 
     /// Create a new script tag in the head
     fn create_script(&self, props: ScriptProps) {
-        WebDocument.create_script(props);
+        self.document.create_script(props);
     }
 
     /// Create a new style tag in the head
     fn create_style(&self, props: StyleProps) {
-        WebDocument.create_style(props);
+        self.document.create_style(props);
     }
 
     /// Create a new link tag in the head
     fn create_link(&self, props: LinkProps) {
-        WebDocument.create_link(props);
+        self.document.create_link(props);
     }
 
     fn create_head_component(&self) -> bool {
diff --git a/packages/fullstack-hooks/src/history.rs b/packages/fullstack-hooks/src/history.rs
index fcfa50967..4354094ed 100644
--- a/packages/fullstack-hooks/src/history.rs
+++ b/packages/fullstack-hooks/src/history.rs
@@ -1,10 +1,10 @@
 //! A history provider for fullstack apps that is compatible with hydration.
 
-use std::cell::OnceCell;
+use std::{cell::RefCell, rc::Rc};
 
-use dioxus_core::{queue_effect, schedule_update};
-use dioxus_fullstack_protocol::is_hydrating;
-use dioxus_history::History;
+use dioxus_core::{provide_context, queue_effect, schedule_update, try_consume_context};
+use dioxus_fullstack_protocol::{is_hydrating, SerializeContextEntry};
+use dioxus_history::{history, provide_history_context, History};
 
 // If we are currently in a scope and this is the first run then queue a rerender
 // for after hydration
@@ -21,51 +21,75 @@ fn match_hydration<O>(
     }
 }
 
+#[derive(Debug, Clone, PartialEq)]
+pub(crate) struct ResolvedRouteContext {
+    route: String,
+}
+
+pub(crate) fn finalize_route() {
+    // This may run in tests without the full hydration context set up, if it does, then just
+    // return without modifying the context
+    let Some(entry) = try_consume_context::<RouteEntry>() else {
+        return;
+    };
+    let entry = entry
+        .entry
+        .borrow_mut()
+        .take()
+        .expect("Failed to get initial route from hydration context");
+    if cfg!(feature = "server") {
+        let history = history();
+        let initial_route = history.current_route();
+        entry.insert(&initial_route, std::panic::Location::caller());
+        provide_context(ResolvedRouteContext {
+            route: initial_route,
+        });
+    } else if cfg!(feature = "web") {
+        let initial_route = entry
+            .get()
+            .expect("Failed to get initial route from hydration context");
+        provide_context(ResolvedRouteContext {
+            route: initial_route,
+        });
+    }
+}
+
+/// Provide the fullstack history context. This interacts with the hydration context so it must
+/// be called in the same order on the client and server after the hydration context is created
+pub fn provide_fullstack_history_context<H: History + 'static>(history: H) {
+    let entry = dioxus_fullstack_protocol::serialize_context().create_entry();
+    provide_context(RouteEntry {
+        entry: Rc::new(RefCell::new(Some(entry.clone()))),
+    });
+    provide_history_context(Rc::new(FullstackHistory::new(history)));
+}
+
+#[derive(Clone)]
+struct RouteEntry {
+    entry: Rc<RefCell<Option<SerializeContextEntry<String>>>>,
+}
+
 /// A history provider for fullstack apps that is compatible with hydration.
 #[derive(Clone)]
-pub struct FullstackHistory<H> {
-    initial_route: OnceCell<String>,
-    #[cfg(feature = "server")]
-    in_hydration_context: std::cell::Cell<bool>,
+struct FullstackHistory<H> {
     history: H,
 }
 
 impl<H> FullstackHistory<H> {
     /// Create a new `FullstackHistory` with the given history.
     pub fn new(history: H) -> Self {
-        Self {
-            initial_route: OnceCell::new(),
-            #[cfg(feature = "server")]
-            in_hydration_context: std::cell::Cell::new(false),
-            history,
-        }
-    }
-
-    /// Create a new `FullstackHistory` with the given history and initial route.
-    pub fn new_server(history: H) -> Self
-    where
-        H: History,
-    {
-        let initial_route = history.current_route();
-        let history = Self::new(history);
-        history.initial_route.set(initial_route).unwrap();
-        history
+        Self { history }
     }
 
     /// Get the initial route of the history.
-    fn initial_route(&self) -> String {
-        let entry = dioxus_fullstack_protocol::serialize_context().create_entry();
-        let route = self.initial_route.get_or_init(|| {
-            entry
-                .get()
-                .expect("Failed to get initial route from hydration context")
-        });
-        #[cfg(feature = "server")]
-        if !self.in_hydration_context.get() {
-            entry.insert(route, std::panic::Location::caller());
-            self.in_hydration_context.set(true);
+    fn initial_route(&self) -> String
+    where
+        H: History,
+    {
+        match try_consume_context::<ResolvedRouteContext>() {
+            Some(context) => context.route,
+            None => self.history.current_route(),
         }
-        route.clone()
     }
 }
 
diff --git a/packages/fullstack-hooks/src/lib.rs b/packages/fullstack-hooks/src/lib.rs
index 83637de73..f5da3d6b5 100644
--- a/packages/fullstack-hooks/src/lib.rs
+++ b/packages/fullstack-hooks/src/lib.rs
@@ -1,6 +1,7 @@
 #![warn(missing_docs)]
 #![doc = include_str!("../README.md")]
 
+pub mod document;
 pub mod history;
 mod hooks;
 mod streaming;
diff --git a/packages/fullstack-hooks/src/streaming.rs b/packages/fullstack-hooks/src/streaming.rs
index 5f7177bc4..147569ed4 100644
--- a/packages/fullstack-hooks/src/streaming.rs
+++ b/packages/fullstack-hooks/src/streaming.rs
@@ -69,6 +69,7 @@ impl StreamingContext {
 /// }
 /// ```
 pub fn commit_initial_chunk() {
+    crate::history::finalize_route();
     if let Some(mut streaming) = try_consume_context::<StreamingContext>() {
         streaming.commit_initial_chunk();
     }
diff --git a/packages/fullstack/src/lib.rs b/packages/fullstack/src/lib.rs
index 9156d839c..47f56854a 100644
--- a/packages/fullstack/src/lib.rs
+++ b/packages/fullstack/src/lib.rs
@@ -4,17 +4,11 @@
 #![warn(missing_docs)]
 #![cfg_attr(docsrs, feature(doc_cfg))]
 
-#[cfg(all(feature = "web", feature = "document"))]
-mod web;
-
 mod error;
 #[doc(hidden)]
 pub mod mock_client;
 
-#[cfg(all(feature = "web", feature = "document"))]
-pub use web::FullstackWebDocument;
-
-pub use dioxus_fullstack_hooks::history::FullstackHistory;
+pub use dioxus_fullstack_hooks::history::provide_fullstack_history_context;
 
 pub use crate::error::{ServerFnError, ServerFnResult};
 #[doc(inline)]
diff --git a/packages/server/src/render.rs b/packages/server/src/render.rs
index ffbac12dc..99fb6fa87 100644
--- a/packages/server/src/render.rs
+++ b/packages/server/src/render.rs
@@ -9,7 +9,7 @@ use dioxus_core::{
     has_context, provide_error_boundary, DynamicNode, ErrorContext, ScopeId, SuspenseContext,
     VNode, VirtualDom,
 };
-use dioxus_fullstack_hooks::history::FullstackHistory;
+use dioxus_fullstack_hooks::history::provide_fullstack_history_context;
 use dioxus_fullstack_hooks::{StreamingContext, StreamingStatus};
 use dioxus_fullstack_protocol::{HydrationContext, SerializedHydrationData};
 use dioxus_isrg::{CachedRender, IncrementalRendererError, RenderFreshness};
@@ -192,24 +192,22 @@ impl SsrRendererPool {
             virtual_dom.provide_root_context(document.clone());
             // If there is a base path, trim the base path from the route and add the base path formatting to the
             // history provider
-            let history;
-            if let Some(base_path) = base_path() {
+            let history = if let Some(base_path) = base_path() {
                 let base_path = base_path.trim_matches('/');
                 let base_path = format!("/{base_path}");
                 let route = route.strip_prefix(&base_path).unwrap_or(&route);
-                history =
-                    dioxus_history::MemoryHistory::with_initial_path(route).with_prefix(base_path);
+                dioxus_history::MemoryHistory::with_initial_path(route).with_prefix(base_path)
             } else {
-                history = dioxus_history::MemoryHistory::with_initial_path(&route);
-            }
-            // Wrap the memory history in a fullstack history provider to provide the initial route for hydration
-            let history = FullstackHistory::new_server(history);
+                dioxus_history::MemoryHistory::with_initial_path(&route)
+            };
 
             let streaming_context = in_root_scope(&virtual_dom, StreamingContext::new);
-            virtual_dom.provide_root_context(Rc::new(history) as Rc<dyn dioxus_history::History>);
             virtual_dom.provide_root_context(document.clone() as Rc<dyn dioxus_document::Document>);
             virtual_dom.provide_root_context(streaming_context);
 
+            // Wrap the memory history in a fullstack history provider to provide the initial route for hydration
+            in_root_scope(&virtual_dom, || provide_fullstack_history_context(history));
+
             // rebuild the virtual dom
             virtual_dom.rebuild_in_place();
 
diff --git a/packages/web/Cargo.toml b/packages/web/Cargo.toml
index 644c3f79c..8f694471c 100644
--- a/packages/web/Cargo.toml
+++ b/packages/web/Cargo.toml
@@ -23,6 +23,7 @@ dioxus-interpreter-js = { workspace = true, features = [
     "webonly",
 ] }
 dioxus-fullstack-protocol = { workspace = true, features = ["web"], optional = true }
+dioxus-fullstack-hooks = { workspace = true, features = ["web"], optional = true }
 generational-box = { workspace = true }
 
 js-sys = { workspace = true }
@@ -86,7 +87,7 @@ lazy-js-bundle = { workspace = true }
 
 [features]
 default = ["mounted", "file_engine", "devtools", "document"]
-hydrate = ["web-sys/Comment", "dep:serde", "dep:dioxus-fullstack-protocol"]
+hydrate = ["web-sys/Comment", "dep:serde", "dep:dioxus-fullstack-protocol", "dep:dioxus-fullstack-hooks"]
 mounted = [
     "web-sys/Element",
     "dioxus-html/mounted",
diff --git a/packages/web/src/document.rs b/packages/web/src/document.rs
index 750eabc0c..4de16e274 100644
--- a/packages/web/src/document.rs
+++ b/packages/web/src/document.rs
@@ -1,3 +1,4 @@
+use dioxus_core::provide_context;
 use dioxus_core::queue_effect;
 use dioxus_core::ScopeId;
 use dioxus_document::{
@@ -66,18 +67,51 @@ extern "C" {
 
 }
 
+fn init_document_with(document: impl FnOnce(), history: impl FnOnce()) {
+    use dioxus_core::has_context;
+
+    ScopeId::ROOT.in_runtime(|| {
+        if has_context::<Rc<dyn Document>>().is_none() {
+            document();
+        }
+        if has_context::<Rc<dyn History>>().is_none() {
+            history();
+        }
+    })
+}
+
 /// Provides the Document through [`ScopeId::provide_context`].
 pub fn init_document() {
-    let provider: Rc<dyn Document> = Rc::new(WebDocument);
-    if ScopeId::ROOT.has_context::<Rc<dyn Document>>().is_none() {
-        ScopeId::ROOT.provide_context(provider);
-    }
-    let history_provider: Rc<dyn History> = Rc::new(WebHistory::default());
-    if ScopeId::ROOT.has_context::<Rc<dyn History>>().is_none() {
-        ScopeId::ROOT.provide_context(history_provider);
+    // If hydrate is enabled, we add the FullstackWebDocument with the initial hydration data
+    #[cfg(not(feature = "hydrate"))]
+    {
+        use dioxus_history::provide_history_context;
+
+        init_document_with(
+            || {
+                provide_context(Rc::new(WebDocument) as Rc<dyn Document>);
+            },
+            || {
+                provide_history_context(Rc::new(WebHistory::default()));
+            },
+        );
     }
 }
 
+#[cfg(feature = "hydrate")]
+pub fn init_fullstack_document() {
+    use dioxus_fullstack_hooks::{
+        document::FullstackWebDocument, history::provide_fullstack_history_context,
+    };
+
+    init_document_with(
+        || {
+            provide_context(Rc::new(FullstackWebDocument::from(WebDocument)) as Rc<dyn Document>);
+        },
+        || provide_fullstack_history_context(WebHistory::default()),
+    );
+}
+
 /// The web-target's document provider.
 #[derive(Clone)]
 pub struct WebDocument;
diff --git a/packages/web/src/lib.rs b/packages/web/src/lib.rs
index d851f35e0..d63d48a21 100644
--- a/packages/web/src/lib.rs
+++ b/packages/web/src/lib.rs
@@ -64,7 +64,8 @@ pub async fn run(mut virtual_dom: VirtualDom, web_config: Config) -> ! {
     #[cfg(all(feature = "devtools", debug_assertions))]
     let mut hotreload_rx = devtools::init();
 
-    let should_hydrate = web_config.hydrate;
+    // If the hydrate feature is enabled, launch the client with hydration enabled
+    let should_hydrate = web_config.hydrate || cfg!(feature = "hydrate");
 
     let mut websys_dom = WebsysDom::new(web_config, runtime);
 
@@ -134,6 +135,10 @@ pub async fn run(mut virtual_dom: VirtualDom, web_config: Config) -> ! {
                 virtual_dom.in_runtime(|| dioxus_core::ScopeId::APP.throw_error(error));
             }
             server_data.in_context(|| {
+                #[cfg(feature = "document")]
+                virtual_dom.in_runtime(|| {
+                    document::init_fullstack_document();
+                });
                 virtual_dom.rebuild(&mut websys_dom);
             });
             websys_dom.skip_mutations = false;