Improvements and docs

aram356 · aram356 · commit d5e6ff1cf2bf · 2026-02-09T13:24:53.000-08:00
diff --git a/crates/common/src/html_processor.rs b/crates/common/src/html_processor.rs
@@ -111,7 +111,7 @@ impl HtmlProcessorConfig {
     }
 }
 
-/// Create an HTML processor with URL replacement and optional Prebid injection
+/// Create an HTML processor with URL replacement and integration hooks
 #[must_use]
 pub fn create_html_processor(config: HtmlProcessorConfig) -> impl StreamProcessor {
     let post_processors = config.integrations.html_post_processors();
@@ -474,6 +474,7 @@ mod tests {
     use super::*;
     use crate::integrations::{
         AttributeRewriteAction, IntegrationAttributeContext, IntegrationAttributeRewriter,
+        IntegrationHeadInjector, IntegrationHtmlContext,
     };
     use crate::streaming_processor::{Compression, PipelineConfig, StreamingPipeline};
     use crate::test_support::tests::create_test_settings;
@@ -544,6 +545,77 @@ mod tests {
         assert!(!processed.contains("remove-me"));
     }
 
+    #[test]
+    fn integration_head_injector_prepends_after_tsjs_once() {
+        struct TestHeadInjector;
+
+        impl IntegrationHeadInjector for TestHeadInjector {
+            fn integration_id(&self) -> &'static str {
+                "test"
+            }
+
+            fn head_inserts(&self, _ctx: &IntegrationHtmlContext<'_>) -> Vec<String> {
+                vec![r#"<script>window.__testHeadInjector=true;</script>"#.to_string()]
+            }
+        }
+
+        let html = r#"<html><head><title>Test</title></head><body></body></html>"#;
+
+        let mut config = create_test_config();
+        config.integrations = IntegrationRegistry::from_rewriters_with_head_injectors(
+            Vec::new(),
+            Vec::new(),
+            vec![Arc::new(TestHeadInjector)],
+        );
+
+        let processor = create_html_processor(config);
+        let pipeline_config = PipelineConfig {
+            input_compression: Compression::None,
+            output_compression: Compression::None,
+            chunk_size: 8192,
+        };
+        let mut pipeline = StreamingPipeline::new(pipeline_config, processor);
+
+        let mut output = Vec::new();
+        pipeline
+            .process(Cursor::new(html.as_bytes()), &mut output)
+            .expect("pipeline should process HTML");
+        let processed = String::from_utf8(output).expect("output should be valid UTF-8");
+
+        let tsjs_marker = "id=\"trustedserver-js\"";
+        let head_marker = "window.__testHeadInjector=true";
+
+        assert_eq!(
+            processed.matches(tsjs_marker).count(),
+            1,
+            "should inject unified tsjs tag once"
+        );
+        assert_eq!(
+            processed.matches(head_marker).count(),
+            1,
+            "should inject head snippet once"
+        );
+
+        let tsjs_index = processed
+            .find(tsjs_marker)
+            .expect("should include unified tsjs tag");
+        let head_index = processed
+            .find(head_marker)
+            .expect("should include head snippet");
+        let title_index = processed
+            .find("<title>")
+            .expect("should keep existing head content");
+
+        assert!(
+            tsjs_index < head_index,
+            "should inject head snippet after tsjs tag"
+        );
+        assert!(
+            head_index < title_index,
+            "should prepend head snippet before existing head content"
+        );
+    }
+
     #[test]
     fn test_create_html_processor_url_replacement() {
         let config = create_test_config();
diff --git a/crates/common/src/integrations/mod.rs b/crates/common/src/integrations/mod.rs
@@ -15,9 +15,9 @@ pub mod testlight;
 pub use registry::{
     AttributeRewriteAction, AttributeRewriteOutcome, IntegrationAttributeContext,
     IntegrationAttributeRewriter, IntegrationDocumentState, IntegrationEndpoint,
-    IntegrationHtmlContext, IntegrationHtmlPostProcessor, IntegrationMetadata, IntegrationProxy,
-    IntegrationRegistration, IntegrationRegistrationBuilder, IntegrationRegistry,
-    IntegrationScriptContext, IntegrationScriptRewriter, ScriptRewriteAction,
+    IntegrationHeadInjector, IntegrationHtmlContext, IntegrationHtmlPostProcessor,
+    IntegrationMetadata, IntegrationProxy, IntegrationRegistration, IntegrationRegistrationBuilder,
+    IntegrationRegistry, IntegrationScriptContext, IntegrationScriptRewriter, ScriptRewriteAction,
 };
 
 type IntegrationBuilder = fn(&Settings) -> Option<IntegrationRegistration>;
diff --git a/crates/common/src/integrations/registry.rs b/crates/common/src/integrations/registry.rs
@@ -503,6 +503,7 @@ pub struct IntegrationMetadata {
     pub routes: Vec<IntegrationEndpoint>,
     pub attribute_rewriters: usize,
     pub script_selectors: Vec<&'static str>,
+    pub head_injectors: usize,
 }
 
 impl IntegrationMetadata {
@@ -512,6 +513,7 @@ impl IntegrationMetadata {
             routes: Vec::new(),
             attribute_rewriters: 0,
             script_selectors: Vec::new(),
+            head_injectors: 0,
         }
     }
 }
@@ -725,6 +727,13 @@ impl IntegrationRegistry {
             entry.script_selectors.push(rewriter.selector());
         }
 
+        for injector in &self.inner.head_injectors {
+            let entry = map
+                .entry(injector.integration_id())
+                .or_insert_with(|| IntegrationMetadata::new(injector.integration_id()));
+            entry.head_injectors += 1;
+        }
+
         map.into_values().collect()
     }
 
@@ -751,7 +760,7 @@ impl IntegrationRegistry {
     }
 
     #[cfg(test)]
-    #[must_use] 
+    #[must_use]
     pub fn from_rewriters_with_head_injectors(
         attribute_rewriters: Vec<Arc<dyn IntegrationAttributeRewriter>>,
         script_rewriters: Vec<Arc<dyn IntegrationScriptRewriter>>,
diff --git a/docs/guide/creative-processing.md b/docs/guide/creative-processing.md
@@ -666,6 +666,32 @@ impl IntegrationScriptRewriter for NextJsIntegration {
 - `replace(content)` - Replace script content
 - `remove_node()` - Delete script element
 
+### Head Injectors
+
+Integrations can inject HTML snippets at the start of `<head>`, immediately after the unified TSJS bundle:
+
+**Example**: An integration injects configuration that runs after the TSJS API is available
+
+```rust
+impl IntegrationHeadInjector for MyIntegration {
+    fn integration_id(&self) -> &'static str { "my_integration" }
+
+    fn head_inserts(&self, ctx: &IntegrationHtmlContext<'_>) -> Vec<String> {
+        vec![format!(
+            r#"<script>tsjs.setConfig({{ host: "{}" }});</script>"#,
+            ctx.request_host
+        )]
+    }
+}
+```
+
+**Behavior**:
+
+- Snippets are prepended into `<head>` after the TSJS bundle tag
+- Called once per HTML response
+- Multiple integrations can each contribute snippets
+- If no snippets are returned, no extra markup is added
+
 See [Integration Guide](/guide/integration-guide) for creating custom rewriters.
 
 ## TSJS Injection
diff --git a/docs/guide/integration-guide.md b/docs/guide/integration-guide.md
@@ -6,7 +6,7 @@ This document explains how to integrate a new integration module with the Truste
 
 | Component                                                  | Purpose                                                                                                                                                                                                                                                      |
 | ---------------------------------------------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------ |
-| `crates/common/src/integrations/registry.rs`               | Defines the `IntegrationProxy`, `IntegrationAttributeRewriter`, and `IntegrationScriptRewriter` traits and hosts the `IntegrationRegistry`, which drives proxy routing and HTML/text rewrites.                                                               |
+| `crates/common/src/integrations/registry.rs`               | Defines the `IntegrationProxy`, `IntegrationAttributeRewriter`, `IntegrationScriptRewriter`, and `IntegrationHeadInjector` traits and hosts the `IntegrationRegistry`, which drives proxy routing, HTML/text rewrites, and head injection.                   |
 | `Settings::integrations` (`crates/common/src/settings.rs`) | Free-form JSON blob keyed by integration ID. Use `IntegrationSettings::insert_config` to seed configs; each module deserializes and validates (`validator::Validate`) its own config and exposes an `enabled` flag so the core settings schema stays stable. |
 | Fastly entrypoint (`crates/fastly/src/main.rs`)            | Instantiates the registry once per request, routes `/integrations/<id>/…` requests to the appropriate proxy, and passes the registry to the publisher origin proxy so HTML rewriting remains integration-aware.                                              |
 | `html_processor.rs`                                        | Applies first-party URL rewrites, injects the Trusted Server JS shim, and lets integrations override attribute values (for example to swap script URLs).                                                                                                     |
@@ -79,14 +79,15 @@ pub fn register(settings: &Settings) -> Option<IntegrationRegistration> {
         IntegrationRegistration::builder("my_integration")
             .with_proxy(integration.clone())
             .with_attribute_rewriter(integration.clone())
-            .with_script_rewriter(integration)
+            .with_script_rewriter(integration.clone())
+            .with_head_injector(integration)
             .with_asset("my_integration")
             .build(),
     )
 }
 ```
 
-Any combination of the three vectors may be populated. Modules that only need HTML rewrites can skip the `proxies` field altogether, and vice versa. The registry automatically iterates over the static builder list in `crates/common/src/integrations/mod.rs`, so adding the new `register` function is enough to make the integration discoverable.
+Any combination of the vectors may be populated. Modules that only need HTML rewrites can skip the `proxies` field altogether, and vice versa. The registry automatically iterates over the static builder list in `crates/common/src/integrations/mod.rs`, so adding the new `register` function is enough to make the integration discoverable.
 
 ### 4. Implement IntegrationProxy for Endpoints
 
@@ -214,6 +215,29 @@ impl IntegrationScriptRewriter for MyIntegration {
 Returning `AttributeRewriteAction::remove_element()` (or `ScriptRewriteAction::RemoveNode` for inline content) removes the element entirely, so integrations can drop publisher-provided markup when the Trusted Server already injects a safe alternative. Prebid, for example, simply removes `prebid.js` because the unified TSJS bundle is injected automatically at the start of `<head>`.
 :::
 
+### 5b. Implement Head Injection (Optional)
+
+If the integration needs to inject HTML snippets at the start of `<head>` (for example, configuration scripts that run after the unified TSJS bundle), implement `IntegrationHeadInjector`. Snippets returned by this trait are prepended into `<head>` immediately after the TSJS bundle tag, so the `tsjs` API is available.
+
+```rust
+impl IntegrationHeadInjector for MyIntegration {
+    fn integration_id(&self) -> &'static str { "my_integration" }
+
+    fn head_inserts(&self, ctx: &IntegrationHtmlContext<'_>) -> Vec<String> {
+        vec![format!(
+            r#"<script>tsjs.setConfig({{ mode: "my_integration", host: "{}" }});</script>"#,
+            ctx.request_host
+        )]
+    }
+}
+```
+
+`html_processor.rs` calls `head_inserts` once per HTML response when the `<head>` element is first encountered. The returned snippets are concatenated after the unified script tag and prepended together, so ordering between integrations is not guaranteed — keep snippets self-contained.
+
+::: tip When to Use Head Injection
+Use `IntegrationHeadInjector` when you need to emit configuration, inline scripts, or `<meta>` tags that must appear early in `<head>`. For attribute or script content changes on existing elements, prefer `IntegrationAttributeRewriter` or `IntegrationScriptRewriter` instead.
+:::
+
 ### 6. Register the Module
 
 Add the module to `crates/common/src/integrations/mod.rs`'s builder list. The registry will call its `register` function automatically. Once registered:
diff --git a/docs/guide/integrations-overview.md b/docs/guide/integrations-overview.md
@@ -174,11 +174,12 @@ All integrations support:
 
 ### Rewriting System
 
-Integrations can implement three types of rewriting:
+Integrations can implement four types of rewriting:
 
 1. **HTTP Proxying** - Route requests through first-party domain
 2. **HTML Attribute Rewriting** - Modify element attributes during streaming
 3. **Script Content Rewriting** - Transform inline script content
+4. **Head Injection** - Insert HTML snippets at the start of `<head>`
 
 ## Choosing an Integration
 
@@ -242,6 +243,7 @@ You can create your own integrations by implementing the integration traits:
 - `IntegrationProxy` - For HTTP endpoint proxying
 - `IntegrationAttributeRewriter` - For HTML attribute rewriting
 - `IntegrationScriptRewriter` - For script content transformation
+- `IntegrationHeadInjector` - For injecting HTML snippets into `<head>`
 
 See the [Integration Guide](./integration-guide.md) for details on building custom integrations.
 

Original file line number	Diff line number	Diff line change
`@@ -503,6 +503,7 @@ pub struct IntegrationMetadata {`
`503`	`503`	`pub routes: Vec<IntegrationEndpoint>,`
`504`	`504`	`pub attribute_rewriters: usize,`
`505`	`505`	`pub script_selectors: Vec<&'static str>,`
	`506`	`+ pub head_injectors: usize,`
`506`	`507`	`}`
`507`	`508`
`508`	`509`	`impl IntegrationMetadata {`
`@@ -512,6 +513,7 @@ impl IntegrationMetadata {`
`512`	`513`	`routes: Vec::new(),`
`513`	`514`	`attribute_rewriters: 0,`
`514`	`515`	`script_selectors: Vec::new(),`
	`516`	`+ head_injectors: 0,`
`515`	`517`	`}`
`516`	`518`	`}`
`517`	`519`	`}`
`@@ -725,6 +727,13 @@ impl IntegrationRegistry {`
`725`	`727`	`entry.script_selectors.push(rewriter.selector());`
`726`	`728`	`}`
`727`	`729`
	`730`	`+ for injector in &self.inner.head_injectors {`
	`731`	`+ let entry = map`
	`732`	`+ .entry(injector.integration_id())`
	`733`	`+ .or_insert_with(\|\| IntegrationMetadata::new(injector.integration_id()));`
	`734`	`+ entry.head_injectors += 1;`
	`735`	`+ }`
	`736`	`+`
`728`	`737`	`map.into_values().collect()`
`729`	`738`	`}`
`730`	`739`
`@@ -751,7 +760,7 @@ impl IntegrationRegistry {`
`751`	`760`	`}`
`752`	`761`
`753`	`762`	`#[cfg(test)]`
`754`		`- #[must_use]`
	`763`	`+ #[must_use]`
`755`	`764`	`pub fn from_rewriters_with_head_injectors(`
`756`	`765`	`attribute_rewriters: Vec<Arc<dyn IntegrationAttributeRewriter>>,`
`757`	`766`	`script_rewriters: Vec<Arc<dyn IntegrationScriptRewriter>>,`