feat(ai-assistant): add support for codeblocks

hlx_statics/blocks/ai-assistant/ai-assistant.css

@@ -422,6 +422,7 @@
     justify-content: center;
     padding: 24px;
     box-sizing: border-box;
+    z-index: 11;
 }
 
 .ai-assistant-panel .chat-window .chat-window-dialog-card {

hlx_statics/blocks/ai-assistant/ai-assistant_api-client.js

@@ -396,7 +396,7 @@ export class AiApiClient {
     callbacks = {},
   }) {
     const defaultSystemPrompt = `
-      Use markdown formatting for the response.
+      Use markdown formatting and codeblocks for the response.
     `;
 
     /** @type {RequestBody} */

hlx_statics/blocks/ai-assistant/ai-assistant_chat-bubble.js

@@ -1,5 +1,7 @@
 // @ts-check
 import { createTag } from "../../scripts/lib-adobeio.js";
+import decoratePreformattedCode from "../../components/code.js";
+import { ensurePrismLoaded } from "../../scripts/prism-loader.js";
 import { aiApiClient } from "./ai-assistant_api-client.js";
 import { chatHistory } from "./ai-assistant_chat-history.js";
 import { createAiAvatar } from "./ai-assistant_chat-ui.js";
@@ -74,6 +76,8 @@ export class ChatBubble {
         bubble.dataset.messageId = this.id;
         // If we have an ID already (e.g., restored from history), append the actions
         contentElement.appendChild(this._actionsContainer);
+        // Restored-history bubbles never call completeBubble(), so decorate here.
+        this.#_decorateCodeBlocks(contentElement);
       }
       // Otherwise, the actions will be appended when completeBubble is called
     }
@@ -301,7 +305,7 @@ export class ChatBubble {
   }
 
   /**
-   * Adds the missing actions to an AI bubble
+   * Final processing that needs to be done after streaming finishes.
    */
   completeBubble() {
     if (this.source !== "ai" || !this._actionsContainer) return;
@@ -310,6 +314,52 @@ export class ChatBubble {
 
     const contentElement = this.element.querySelector(".chat-bubble-content");
     contentElement?.appendChild(this._actionsContainer);
+    if (contentElement) {
+      this.#_decorateCodeBlocks(contentElement);
+    }
+  }
+
+  /**
+   * @param {Element} contentElement - The `.chat-bubble-content` element
+   */
+  #_decorateCodeBlocks(contentElement) {
+    const preBlocks = contentElement.querySelectorAll("pre");
+    if (!preBlocks.length) return;
+
+    let decoratedAny = false;
+    preBlocks.forEach((pre) => {
+      // decoratePreformattedCode dereferences a <code> child unconditionally.
+      if (!pre.querySelector("code")) return;
+      decoratePreformattedCode(pre);
+      decoratedAny = true;
+    });
+
+    if (!decoratedAny) return;
+
+    ensurePrismLoaded().then(() => {
+      // @ts-expect-error - Prism is not on the Window type
+      window.Prism?.highlightAllUnder?.(contentElement);
+    });
+  }
+
+  /**
+   * Recomputes Prism's line-number row heights for every decorated code block
+   * inside a container.
+   * This is required to correctly align line numbers after restoring history.
+   * @param {Element | null | undefined} container
+   */
+  static resizeCodeBlockLineNumbers(container) {
+    const preBlocks = container?.querySelectorAll("pre.line-numbers");
+    if (!preBlocks?.length) return;
+
+    ensurePrismLoaded().then(() => {
+      // @ts-expect-error - Prism is not on the Window type
+      const resize = window.Prism?.plugins?.lineNumbers?.resize;
+      if (typeof resize !== "function") return;
+      preBlocks.forEach((pre) => {
+        resize(pre);
+      });
+    });
   }
 
   /**

hlx_statics/blocks/ai-assistant/ai-assistant_chat-controller.js

@@ -82,6 +82,25 @@ export const openChatWindow = () => {
   ELEMENTS.CHAT_WINDOW.classList.add("show");
   ELEMENTS.CHAT_BUTTON.classList.add("hidden");
 
+  // Code blocks in restored-history bubbles are decorated while the window is
+  // hidden (scaled to 0), so Prism's line-number rows collapse onto line 1.
+  // Recompute them once the open transition settles and the window has real
+  // layout. transitionend is the precise signal; the timeout is a fallback in
+  // case the transition is skipped (e.g. reduced motion / 0s duration).
+  const chatWindow = ELEMENTS.CHAT_WINDOW;
+  let fixedLineNumbers = false;
+  const fixLineNumbers = (event) => {
+    // transitionend bubbles, so ignore descendant transitions and only react to
+    // the window's own transform settling (or the timeout, which has no event).
+    if (event && (event.target !== chatWindow || event.propertyName !== "transform")) return;
+    if (fixedLineNumbers) return;
+    fixedLineNumbers = true;
+    chatWindow.removeEventListener("transitionend", fixLineNumbers);
+    ChatBubble.resizeCodeBlockLineNumbers(ELEMENTS.CHAT_WINDOW_CONTENT);
+  };
+  chatWindow.addEventListener("transitionend", fixLineNumbers);
+  window.setTimeout(fixLineNumbers, 400);
+
   // Initial messages
   if (chatHistory.isEmpty()) {
     sendInitialMessages();

hlx_statics/scripts/prism-loader.js

@@ -0,0 +1,35 @@
+import { loadCSS } from "./lib-helix.js";
+
+/**
+ * Cached promise so the Prism import/config runs at most once per page,
+ * regardless of how many callers (docs code blocks, AI assistant, …) need it.
+ * @type {Promise<void> | null}
+ */
+let prismPromise = null;
+
+/**
+ * Load Prism. Call this before using `window.Prism`.
+ * @returns {Promise<void>} resolves once `window.Prism` is ready to use.
+ */
+export function ensurePrismLoaded() {
+  if (prismPromise) return prismPromise;
+
+  prismPromise = (async () => {
+    window.Prism = { manual: true };
+    loadCSS(`${window.hlx.codeBasePath}/styles/prism.css`);
+    await import("./prism.js");
+
+    // Ensure Prism autoloader knows where to fetch language components
+    if (
+      window.Prism &&
+      window.Prism.plugins &&
+      window.Prism.plugins.autoloader
+    ) {
+      window.Prism.plugins.autoloader.languages_path =
+        "/hlx_statics/scripts/prism-grammars/";
+      window.Prism.plugins.autoloader.use_minified = true;
+    }
+  })();
+
+  return prismPromise;
+}

hlx_statics/scripts/scripts.js

@@ -24,6 +24,8 @@ import {
   getCodePlaygroundJsonPath
 } from './lib-helix.js';
 
+import { ensurePrismLoaded } from './prism-loader.js';
+
 import {
   buildAiAssistant,
   buildBreadcrumbs,
@@ -898,15 +900,7 @@ async function loadPrism(document) {
 
       if (!prismLoaded) {
         prismLoaded = true;
-        window.Prism = { manual: true };
-        loadCSS(`${window.hlx.codeBasePath}/styles/prism.css`);
-        import('./prism.js').then(() => {
-          // Ensure Prism autoloader knows where to fetch language components
-          if (window.Prism && window.Prism.plugins && window.Prism.plugins.autoloader) {
-            window.Prism.plugins.autoloader.languages_path = '/hlx_statics/scripts/prism-grammars/';
-            window.Prism.plugins.autoloader.use_minified = true;
-          }
-
+        ensurePrismLoaded().then(() => {
           // Register "Try it" button for code blocks with "try" class
           if (window.Prism?.plugins?.toolbar) {
             window.Prism.plugins.toolbar.registerButton('try-code', (env) => {