fix: chat push url and references

html-router/src/middlewares/response_middleware.rs

@@ -1,6 +1,6 @@
 use axum::{
     extract::State,
-    http::StatusCode,
+    http::{HeaderName, StatusCode},
     response::{Html, IntoResponse, Response},
     Extension,
 };
@@ -106,13 +106,31 @@ pub async fn with_template_response<S>(
 where
     S: ProvidesTemplateEngine + Clone + Send + Sync + 'static,
 {
+    // Headers to forward from the original response
+    const HTMX_HEADERS_TO_FORWARD: &[&str] = &["HX-Push", "HX-Trigger", "HX-Redirect"];
+
     if let Some(template_response) = response.extensions().get::<TemplateResponse>().cloned() {
         let template_engine = state.template_engine();
 
+        // Helper to forward relevant headers
+        fn forward_headers(from: &axum::http::HeaderMap, to: &mut axum::http::HeaderMap) {
+            for &header_name in HTMX_HEADERS_TO_FORWARD {
+                if let Ok(name) = HeaderName::from_bytes(header_name.as_bytes()) {
+                    if let Some(value) = from.get(&name) {
+                        to.insert(name.clone(), value.clone());
+                    }
+                }
+            }
+        }
+
         match &template_response.template_kind {
             TemplateKind::Full(name) => {
                 match template_engine.render(name, &template_response.context) {
-                    Ok(html) => Html(html).into_response(),
+                    Ok(html) => {
+                        let mut final_response = Html(html).into_response();
+                        forward_headers(response.headers(), final_response.headers_mut());
+                        final_response
+                    }
                     Err(e) => {
                         error!("Failed to render template '{}': {:?}", name, e);
                         (StatusCode::INTERNAL_SERVER_ERROR, Html(fallback_error())).into_response()
@@ -121,7 +139,11 @@ where
             }
             TemplateKind::Partial(template, block) => {
                 match template_engine.render_block(template, block, &template_response.context) {
-                    Ok(html) => Html(html).into_response(),
+                    Ok(html) => {
+                        let mut final_response = Html(html).into_response();
+                        forward_headers(response.headers(), final_response.headers_mut());
+                        final_response
+                    }
                     Err(e) => {
                         error!("Failed to render block '{}/{}': {:?}", template, block, e);
                         (StatusCode::INTERNAL_SERVER_ERROR, Html(fallback_error())).into_response()

html-router/src/routes/chat/chat_handlers.rs

@@ -222,5 +222,5 @@ pub async fn new_chat_user_message(
         HeaderValue::from_str(&format!("/chat/{}", conversation.id)).unwrap(),
     );
 
-    Ok(response)
+    Ok(response.into_response())
 }

html-router/src/routes/knowledge/handlers.rs

@@ -9,7 +9,7 @@ use axum_htmx::{HxBoosted, HxRequest};
 use plotly::{
     common::{Line, Marker, Mode},
     layout::{Axis, LayoutScene},
-    Layout, Plot, Scatter, Scatter3D,
+    Layout, Plot, Scatter3D,
 };
 use serde::{Deserialize, Serialize};
 

html-router/templates/chat/streaming_response.html

@@ -10,6 +10,7 @@
       <span class="loading loading-dots loading-sm loading-id-{{user_message.id}}"></span>
       <div class="markdown-content" id="ai-message-content-{{user_message.id}}" sse-swap="chat_message"></div>
     </div>
+    <div sse-swap="references"></div>
   </div>
 </div>
 <script>
@@ -22,17 +23,19 @@
     const spinner = document.querySelector('.loading-id-{{user_message.id}}');
     if (spinner) spinner.style.display = 'none';
 
-    const el = document.getElementById('ai-message-content-{{user_message.id}}');
-    if (e.detail.elt !== el) return;
-    e.preventDefault(); // Prevent htmx from swapping
-    // Use message id as buffer key
-    const msgId = '{{user_message.id}}';
-    window.markdownBuffer[msgId] = (window.markdownBuffer[msgId] || '') + (e.detail.data || '');
-    // Render buffer (with newline fix) on *every* chunk
-    el.innerHTML = marked.parse(
-      window.markdownBuffer[msgId].replace(/\\n/g, '\n')
-    );
-    scrollChatToBottom();
+    if (e.detail.event === 'chat_message') {
+      const el = document.getElementById('ai-message-content-{{user_message.id}}');
+      if (e.detail.elt !== el) return;
+      e.preventDefault(); // Prevent htmx from swapping
+      // Use message id as buffer key
+      const msgId = '{{user_message.id}}';
+      window.markdownBuffer[msgId] = (window.markdownBuffer[msgId] || '') + (e.detail.data || '');
+      // Render buffer (with newline fix) on *every* chunk
+      el.innerHTML = marked.parse(
+        window.markdownBuffer[msgId].replace(/\\n/g, '\n')
+      );
+      scrollChatToBottom();
+    }
   });
   document.body.addEventListener('htmx:sseClose', function () {
     const el = document.getElementById('ai-message-content-{{user_message.id}}');