diff --git a/src/store/prompt.ts b/src/store/prompt.ts
index 13cf9a4..cf74621 100644
--- a/src/store/prompt.ts
+++ b/src/store/prompt.ts
@@ -14,6 +14,7 @@ export const conversationId = ref<string>("")
 export const connected = ref(false)
 export const streaming = ref(false)
 export const streamingContent = ref("")
+export let onCodeComplete: ((code: { html: string | null; css: string | null; js: string | null }) => void) | null = null
 
 let ws: WebSocket | null = null
 
@@ -54,6 +55,9 @@ export function connectPrompt(taskId: number) {
       // Apply code to editors
       if (data.code) {
         applyCode(data.code)
+        if (onCodeComplete) {
+          onCodeComplete(data.code)
+        }
       }
     } else if (data.type === "error") {
       streaming.value = false
@@ -80,6 +84,7 @@ export function disconnectPrompt() {
   connected.value = false
   streaming.value = false
   streamingContent.value = ""
+  onCodeComplete = null
 }
 
 export function sendPrompt(content: string) {