diff --git a/CHANGELOG.md b/CHANGELOG.md
index fdb77c1ebe9..e296c49c750 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,5 +1,11 @@
 ## vNEXT (not yet published)
 
+### `@liveblocks/client`
+
+- Fix race condition where AI tools were not always executing. This could happen
+  when using `useSendAiMessage` first and then immediately opening the
+  `<AiChat />` afterwards.
+
 ## v3.3.1
 
 ### `@liveblocks/react-ui`
diff --git a/packages/liveblocks-core/src/ai.ts b/packages/liveblocks-core/src/ai.ts
index 43205ef11d4..a19a6fbe81f 100644
--- a/packages/liveblocks-core/src/ai.ts
+++ b/packages/liveblocks-core/src/ai.ts
@@ -564,7 +564,12 @@ function createStore_forChatMessages(
           }
         }
       } else {
-        myMessages.delete(message.id);
+        // Clean up the ownership administration
+        if (message.role === "assistant" && message.status === "generating") {
+          // ...unless it's still generating
+        } else {
+          myMessages.delete(message.id);
+        }
       }
     });
   }
@@ -1293,6 +1298,7 @@ export function createAi(config: AiConfig): Ai {
         const combinedKnowledge = [...globalKnowledge, ...requestKnowledge];
         const tools = context.toolsStore.getToolDescriptions(chatId);
 
+        messagesStore.markMine(targetMessageId);
         const resp: AskInChatResponse = await sendClientMsgWithResponse({
           cmd: "ask-in-chat",
           chatId,
@@ -1309,7 +1315,6 @@ export function createAi(config: AiConfig): Ai {
             tools: tools.length > 0 ? tools : undefined,
           },
         });
-        messagesStore.markMine(resp.targetMessage.id);
         return resp;
       },