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

packages/fullstack-hooks/src/document.rs

@@ -0,0 +1,57 @@
+//! 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::*;
+
+fn head_element_written_on_server() -> bool {
+    dioxus_fullstack_protocol::head_element_hydration_entry()
+        .get()
+        .ok()
+        .unwrap_or_default()
+}
+
+/// A document provider for fullstack web clients
+#[derive(Clone)]
+pub struct FullstackWebDocument<D> {
+    document: D,
+}
+
+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 {
+        self.document.eval(js)
+    }
+
+    /// Set the title of the document
+    fn set_title(&self, title: String) {
+        self.document.set_title(title);
+    }
+
+    /// Create a new meta tag in the head
+    fn create_meta(&self, props: MetaProps) {
+        self.document.create_meta(props);
+    }
+
+    /// Create a new script tag in the head
+    fn create_script(&self, props: ScriptProps) {
+        self.document.create_script(props);
+    }
+
+    /// Create a new style tag in the head
+    fn create_style(&self, props: StyleProps) {
+        self.document.create_style(props);
+    }
+
+    /// Create a new link tag in the head
+    fn create_link(&self, props: LinkProps) {
+        self.document.create_link(props);
+    }
+
+    fn create_head_component(&self) -> bool {
+        !head_element_written_on_server()
+    }
+}

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()
     }
 }
 

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;

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();
     }