diff --git a/site/src/components/ai-elements/response.stories.tsx b/site/src/components/ai-elements/response.stories.tsx
index 50ec73650e..f5be2b4bcc 100644
--- a/site/src/components/ai-elements/response.stories.tsx
+++ b/site/src/components/ai-elements/response.stories.tsx
@@ -1,5 +1,5 @@
 import type { Meta, StoryObj } from "@storybook/react-vite";
-import { expect, within } from "storybook/test";
+import { expect, waitFor, within } from "storybook/test";
 import { Response } from "./response";
 
 const sampleMarkdown = `
@@ -125,15 +125,21 @@ export const StreamingInlineMarkdown: Story = {
 
 // Verifies that an incomplete fenced code block in streaming mode
 // renders inside a code element rather than showing raw backticks.
+// The FileViewer renders a <diffs-container> web component whose
+// content lives in Shadow DOM, so we assert on DOM structure rather
+// than text content inside the web component.
 export const StreamingCodeFence: Story = {
 	args: {
 		children: "```ts\nconst x = 1",
 		streaming: true,
 	},
 	play: async ({ canvasElement }) => {
-		const canvas = within(canvasElement);
-		const codeContent = await canvas.findByText(/const x = 1/);
-		expect(codeContent).toBeInTheDocument();
+		// The code fence should be parsed into a FileViewer (web component),
+		// not rendered as raw backtick text.
+		await waitFor(() => {
+			const viewer = canvasElement.querySelector("diffs-container");
+			expect(viewer).toBeInTheDocument();
+		});
 		// The raw triple-backtick should not appear as visible text.
 		const bodyText = canvasElement.textContent ?? "";
 		expect(bodyText).not.toContain("```");