Fixed bug

aram356 · aram356 · commit fd5b91d2d20f · 2025-12-17T01:18:19.000-08:00
diff --git a/crates/common/src/integrations/nextjs/html_post_process.rs b/crates/common/src/integrations/nextjs/html_post_process.rs
@@ -30,38 +30,62 @@ impl IntegrationHtmlPostProcessor for NextJsHtmlPostProcessor {
     }
 
     fn should_process(&self, html: &str, ctx: &IntegrationHtmlContext<'_>) -> bool {
-        let _ = html;
         if !self.config.enabled || self.config.rewrite_attributes.is_empty() {
             return false;
         }
 
-        let Some(state) = ctx
+        // Check if we have captured placeholders from streaming
+        if let Some(state) = ctx
             .document_state
             .get::<Mutex<NextJsRscPostProcessState>>(NEXTJS_INTEGRATION_ID)
-        else {
-            return false;
-        };
+        {
+            let guard = state.lock().unwrap_or_else(|e| e.into_inner());
+            if !guard.payloads.is_empty() {
+                return true;
+            }
+        }
 
-        let guard = state.lock().unwrap_or_else(|e| e.into_inner());
-        !guard.payloads.is_empty()
+        // Also check if HTML contains RSC scripts that weren't captured during streaming
+        // (e.g., fragmented scripts that we skipped during the streaming pass)
+        html.contains("__next_f.push") && html.contains(ctx.origin_host)
     }
 
     fn post_process(&self, html: &mut String, ctx: &IntegrationHtmlContext<'_>) -> bool {
-        let Some(state) = ctx
+        // Try to get payloads captured during streaming (placeholder approach)
+        let payloads = ctx
             .document_state
             .get::<Mutex<NextJsRscPostProcessState>>(NEXTJS_INTEGRATION_ID)
-        else {
-            return false;
-        };
-
-        let payloads = {
-            let mut guard = state.lock().unwrap_or_else(|e| e.into_inner());
-            guard.take_payloads()
-        };
-        if payloads.is_empty() {
-            return false;
+            .map(|state| {
+                let mut guard = state.lock().unwrap_or_else(|e| e.into_inner());
+                guard.take_payloads()
+            })
+            .unwrap_or_default();
+
+        if !payloads.is_empty() {
+            // Placeholder approach: substitute placeholders with rewritten payloads
+            return self.substitute_placeholders(html, ctx, payloads);
         }
 
+        // Fallback: re-parse HTML to find RSC scripts that weren't captured during streaming
+        // (e.g., fragmented scripts that we skipped during the streaming pass)
+        post_process_rsc_html_in_place_with_limit(
+            html,
+            ctx.origin_host,
+            ctx.request_host,
+            ctx.request_scheme,
+            self.config.max_combined_payload_bytes,
+        )
+    }
+}
+
+impl NextJsHtmlPostProcessor {
+    /// Substitute placeholders with rewritten payloads (fast path for unfragmented scripts).
+    fn substitute_placeholders(
+        &self,
+        html: &mut String,
+        ctx: &IntegrationHtmlContext<'_>,
+        payloads: Vec<String>,
+    ) -> bool {
         let payload_refs: Vec<&str> = payloads.iter().map(String::as_str).collect();
         let mut rewritten_payloads = rewrite_rsc_scripts_combined_with_limit(
             payload_refs.as_slice(),
diff --git a/crates/common/src/integrations/nextjs/mod.rs b/crates/common/src/integrations/nextjs/mod.rs
@@ -419,4 +419,73 @@ mod tests {
             final_html
         );
     }
+
+    #[test]
+    fn html_processor_preserves_non_rsc_scripts_with_chunked_streaming() {
+        // Regression test: ensure non-RSC scripts are preserved when streamed alongside RSC scripts.
+        // With small chunk sizes, scripts get fragmented and the buffering logic must correctly
+        // handle non-RSC scripts without corrupting them.
+        let html = r#"<html><body>
+<script>console.log("hello world");</script>
+<script>self.__next_f.push([1,'{"url":"https://origin.example.com/page"}'])</script>
+<script>window.analytics = { track: function(e) { console.log(e); } };</script>
+</body></html>"#;
+
+        let mut settings = create_test_settings();
+        settings
+            .integrations
+            .insert_config(
+                "nextjs",
+                &json!({
+                    "enabled": true,
+                    "rewrite_attributes": ["url"],
+                }),
+            )
+            .expect("should update nextjs config");
+        let registry = IntegrationRegistry::new(&settings);
+        let config = config_from_settings(&settings, &registry);
+        let processor = create_html_processor(config);
+        // Use small chunk size to force fragmentation
+        let pipeline_config = PipelineConfig {
+            input_compression: Compression::None,
+            output_compression: Compression::None,
+            chunk_size: 16,
+        };
+        let mut pipeline = StreamingPipeline::new(pipeline_config, processor);
+
+        let mut output = Vec::new();
+        pipeline
+            .process(Cursor::new(html.as_bytes()), &mut output)
+            .unwrap();
+        let final_html = String::from_utf8_lossy(&output);
+
+        // Non-RSC scripts should be preserved
+        assert!(
+            final_html.contains(r#"console.log("hello world");"#),
+            "First non-RSC script should be preserved intact. Output: {}",
+            final_html
+        );
+        assert!(
+            final_html.contains("window.analytics"),
+            "Third non-RSC script should be preserved. Output: {}",
+            final_html
+        );
+        assert!(
+            final_html.contains("track: function(e)"),
+            "Third non-RSC script content should be intact. Output: {}",
+            final_html
+        );
+
+        // RSC scripts should be rewritten
+        assert!(
+            final_html.contains("test.example.com"),
+            "RSC URL should be rewritten. Output: {}",
+            final_html
+        );
+        assert!(
+            !final_html.contains(RSC_PAYLOAD_PLACEHOLDER_PREFIX),
+            "No placeholders should remain. Output: {}",
+            final_html
+        );
+    }
 }
diff --git a/crates/common/src/integrations/nextjs/rsc_placeholders.rs b/crates/common/src/integrations/nextjs/rsc_placeholders.rs
@@ -1,4 +1,3 @@
-use std::borrow::Cow;
 use std::sync::{Arc, Mutex};
 
 use crate::integrations::{
@@ -11,36 +10,18 @@ use super::{NextJsIntegrationConfig, NEXTJS_INTEGRATION_ID};
 pub(super) const RSC_PAYLOAD_PLACEHOLDER_PREFIX: &str = "__ts_rsc_payload_";
 pub(super) const RSC_PAYLOAD_PLACEHOLDER_SUFFIX: &str = "__";
 
+/// State for RSC placeholder-based rewriting.
+///
+/// Stores RSC payloads extracted during streaming for later rewriting during post-processing.
+/// Only unfragmented RSC scripts are processed during streaming; fragmented scripts are
+/// handled by the post-processor which re-parses the final HTML.
 #[derive(Default)]
 pub(super) struct NextJsRscPostProcessState {
     pub(super) payloads: Vec<String>,
-    buffer: String,
-    buffering: bool,
 }
 
 impl NextJsRscPostProcessState {
-    fn buffer_chunk(&mut self, chunk: &str) {
-        if !self.buffering {
-            self.buffering = true;
-            self.buffer.clear();
-        }
-        self.buffer.push_str(chunk);
-    }
-
-    /// Returns the complete script content, either borrowed from input or owned from buffer.
-    fn take_script_or_borrow<'a>(&mut self, chunk: &'a str) -> Cow<'a, str> {
-        if self.buffering {
-            self.buffer.push_str(chunk);
-            self.buffering = false;
-            Cow::Owned(std::mem::take(&mut self.buffer))
-        } else {
-            Cow::Borrowed(chunk)
-        }
-    }
-
     pub(super) fn take_payloads(&mut self) -> Vec<String> {
-        self.buffer.clear();
-        self.buffering = false;
         std::mem::take(&mut self.payloads)
     }
 }
@@ -73,84 +54,48 @@ impl IntegrationScriptRewriter for NextJsRscPlaceholderRewriter {
             return ScriptRewriteAction::keep();
         }
 
+        // Only process complete (unfragmented) scripts during streaming.
+        // Fragmented scripts are handled by the post-processor which re-parses the final HTML.
+        // This avoids corrupting non-RSC scripts that happen to be fragmented during streaming.
         if !ctx.is_last_in_text_node {
-            if let Some(existing) = ctx
-                .document_state
-                .get::<Mutex<NextJsRscPostProcessState>>(NEXTJS_INTEGRATION_ID)
-            {
-                let mut guard = existing.lock().unwrap_or_else(|e| e.into_inner());
-                if guard.buffering {
-                    guard.buffer_chunk(content);
-                    return ScriptRewriteAction::remove_node();
-                }
-            }
-
-            let trimmed = content.trim_start();
-            if trimmed.starts_with('{') || trimmed.starts_with('[') {
-                // Avoid interfering with other inline JSON scripts (e.g. `__NEXT_DATA__`, JSON-LD).
-                return ScriptRewriteAction::keep();
-            }
-
-            let state = ctx
-                .document_state
-                .get_or_insert_with(NEXTJS_INTEGRATION_ID, || {
-                    Mutex::new(NextJsRscPostProcessState::default())
-                });
-            let mut guard = state.lock().unwrap_or_else(|e| e.into_inner());
-            guard.buffer_chunk(content);
-            return ScriptRewriteAction::remove_node();
-        }
-
-        if !content.contains("__next_f")
-            && ctx
-                .document_state
-                .get::<Mutex<NextJsRscPostProcessState>>(NEXTJS_INTEGRATION_ID)
-                .is_none()
-        {
+            // Script is fragmented - skip placeholder processing.
+            // The post-processor will handle RSC scripts at end-of-document.
             return ScriptRewriteAction::keep();
         }
 
-        let state = ctx
-            .document_state
-            .get_or_insert_with(NEXTJS_INTEGRATION_ID, || {
-                Mutex::new(NextJsRscPostProcessState::default())
-            });
-        let mut guard = state.lock().unwrap_or_else(|e| e.into_inner());
-        let script = guard.take_script_or_borrow(content);
-        let was_buffered = matches!(script, Cow::Owned(_));
-
-        if !script.contains("__next_f") {
-            if was_buffered {
-                return ScriptRewriteAction::replace(script.into_owned());
-            }
+        // Quick check: skip scripts that can't be RSC payloads
+        if !content.contains("__next_f") {
             return ScriptRewriteAction::keep();
         }
 
-        let Some((payload_start, payload_end)) = find_rsc_push_payload_range(&script) else {
-            if was_buffered {
-                return ScriptRewriteAction::replace(script.into_owned());
-            }
+        let Some((payload_start, payload_end)) = find_rsc_push_payload_range(content) else {
+            // Contains __next_f but doesn't match RSC push pattern - leave unchanged
             return ScriptRewriteAction::keep();
         };
 
         if payload_start > payload_end
-            || payload_end > script.len()
-            || !script.is_char_boundary(payload_start)
-            || !script.is_char_boundary(payload_end)
+            || payload_end > content.len()
+            || !content.is_char_boundary(payload_start)
+            || !content.is_char_boundary(payload_end)
         {
-            if was_buffered {
-                return ScriptRewriteAction::replace(script.into_owned());
-            }
             return ScriptRewriteAction::keep();
         }
 
+        // Insert placeholder for this RSC payload and store original for post-processing
+        let state = ctx
+            .document_state
+            .get_or_insert_with(NEXTJS_INTEGRATION_ID, || {
+                Mutex::new(NextJsRscPostProcessState::default())
+            });
+        let mut guard = state.lock().unwrap_or_else(|e| e.into_inner());
+
         let placeholder_index = guard.payloads.len();
         let placeholder = rsc_payload_placeholder(placeholder_index);
         guard
             .payloads
-            .push(script[payload_start..payload_end].to_string());
+            .push(content[payload_start..payload_end].to_string());
 
-        let mut rewritten = script.into_owned();
+        let mut rewritten = content.to_string();
         rewritten.replace_range(payload_start..payload_end, &placeholder);
         ScriptRewriteAction::replace(rewritten)
     }
@@ -211,32 +156,52 @@ mod tests {
     }
 
     #[test]
-    fn buffers_fragmented_scripts_and_emits_single_replacement() {
+    fn skips_fragmented_scripts_for_post_processor_handling() {
+        // Fragmented scripts are not processed during streaming - they're passed through
+        // unchanged and handled by the post-processor which re-parses the final HTML.
         let state = IntegrationDocumentState::default();
         let rewriter = NextJsRscPlaceholderRewriter::new(test_config());
 
         let first = "self.__next_f.push([1,\"https://origin.example.com";
         let second = "/page\"])";
 
+        // Intermediate chunk should be kept (not processed)
         let action_first = rewriter.rewrite(first, &ctx(false, &state));
         assert_eq!(
             action_first,
-            ScriptRewriteAction::RemoveNode,
-            "Intermediate chunk should be removed"
+            ScriptRewriteAction::Keep,
+            "Intermediate chunk should be kept unchanged"
         );
 
+        // Final chunk should also be kept since it doesn't contain the full RSC pattern
         let action_second = rewriter.rewrite(second, &ctx(true, &state));
-        let ScriptRewriteAction::Replace(rewritten) = action_second else {
-            panic!("Final chunk should be replaced with combined output");
-        };
+        assert_eq!(
+            action_second,
+            ScriptRewriteAction::Keep,
+            "Final chunk of fragmented script should be kept"
+        );
 
+        // No payloads should be stored - post-processor will handle this
         assert!(
-            rewritten.contains(RSC_PAYLOAD_PLACEHOLDER_PREFIX),
-            "Combined output should include placeholder. Got: {rewritten}"
+            state
+                .get::<Mutex<NextJsRscPostProcessState>>(NEXTJS_INTEGRATION_ID)
+                .is_none(),
+            "No RSC state should be created for fragmented scripts"
         );
-        assert!(
-            rewritten.contains("self.__next_f.push"),
-            "Combined output should keep the push call. Got: {rewritten}"
+    }
+
+    #[test]
+    fn skips_non_rsc_scripts() {
+        let state = IntegrationDocumentState::default();
+        let rewriter = NextJsRscPlaceholderRewriter::new(test_config());
+
+        let script = r#"console.log("hello world");"#;
+        let action = rewriter.rewrite(script, &ctx(true, &state));
+
+        assert_eq!(
+            action,
+            ScriptRewriteAction::Keep,
+            "Non-RSC scripts should be kept unchanged"
         );
     }
 }
