diff --git a/webui/src/web-components/demo/mocks/handlers.ts b/webui/src/web-components/demo/mocks/handlers.ts
index ad3426e..29f0710 100644
--- a/webui/src/web-components/demo/mocks/handlers.ts
+++ b/webui/src/web-components/demo/mocks/handlers.ts
@@ -3,7 +3,6 @@
 import { AgentMessage, State } from "../../../types";
 
 // Mock state updates for SSE simulation
-let currentState = { ...initialState };
 const EMPTY_CONVERSATION =
   new URL(window.location.href).searchParams.get("emptyConversation") === "1";
 const ADD_NEW_MESSAGES =
@@ -11,6 +10,12 @@
 
 const messages = EMPTY_CONVERSATION ? [] : [...initialMessages];
 
+// Initialize state with correct message_count
+let currentState = {
+  ...initialState,
+  message_count: messages.length,
+};
+
 // Text encoder for SSE messages
 const encoder = new TextEncoder();