diff --git a/build/build_number.txt b/build/build_number.txt
index cdffbbc..f2c1eee 100644
--- a/build/build_number.txt
+++ b/build/build_number.txt
@@ -1 +1 @@
-166
+167
diff --git a/src/services/assistant.ts b/src/services/assistant.ts
index e3a7999..fd5e0b0 100644
--- a/src/services/assistant.ts
+++ b/src/services/assistant.ts
@@ -180,13 +180,10 @@ export default class {
       if (error.name !== 'AbortError') {
         if (error.status === 401 || error.message.includes('401') || error.message.toLowerCase().includes('apikey')) {
           message.setText('You need to enter your API key in the Models tab of <a href="#settings">Settings</a> in order to chat.')
-          callback?.call(null, { done: true })
         } else if (error.status === 400 && (error.message.includes('credit') || error.message.includes('balance'))) {
           message.setText('Sorry, it seems you have run out of credits. Check the balance of your LLM provider account.')
-          callback?.call(null, { done: true })
         } else if (message.content === '') {
           message.setText('Sorry, I could not generate text for that prompt.')
-          callback?.call(null, { done: true })
         } else {
           message.appendText({ text: '\n\nSorry, I am not able to continue here.', done: true })
         }
diff --git a/tests/unit/assistant.test.ts b/tests/unit/assistant.test.ts
index 4f8ffb6..6d9cf46 100644
--- a/tests/unit/assistant.test.ts
+++ b/tests/unit/assistant.test.ts
@@ -27,16 +27,14 @@ let assistant: Assistant = null
 const prompt = async (prompt: string, opts: LlmCompletionOpts = {}): Promise<string> => {
 
   // callback
-  let done = false
   let content = ''
   const callback = (chunk: LlmStream) => {
-    done = chunk?.done
     content += chunk?.text || ''
   }
   
   // call and wait
   await assistant.prompt(prompt, { ...opts, save: false }, callback)
-  await vi.waitUntil(async () => done)
+  await vi.waitUntil(async () => !assistant.chat.lastMessage().transient)
 
   // return
   return content