diff --git a/apps/server/src/services/__test__/providers.test.ts b/apps/server/src/services/__test__/providers.test.ts
index 073bd3ab4..a3c0ad8e3 100644
--- a/apps/server/src/services/__test__/providers.test.ts
+++ b/apps/server/src/services/__test__/providers.test.ts
@@ -39,9 +39,8 @@ describe('providerService', () => {
     expect(result.name).toBe('My OpenAI')
 
     const found = await service.findUserConfigById('prov-1')
-    expect(found).toBeDefined()
-    expect(found!.definitionId).toBe('openai')
-    expect((found!.config as Record<string, string>).apiKey).toBe('sk-123')
+    expect(found?.definitionId).toBe('openai')
+    expect((found?.config as any).apiKey).toBe('sk-123')
   })
 
   it('findUserConfigsByOwnerId should return providers for the user', async () => {
diff --git a/packages/stage-ui/src/components/scenes/Stage.vue b/packages/stage-ui/src/components/scenes/Stage.vue
index 23c3fc07e..b23a8c262 100644
--- a/packages/stage-ui/src/components/scenes/Stage.vue
+++ b/packages/stage-ui/src/components/scenes/Stage.vue
@@ -163,56 +163,67 @@ function playSpecialToken(special: string) {
 const lipSyncNode = ref<AudioNode>()
 
 async function playFunction(item: Parameters<Parameters<typeof createPlaybackManager<AudioBuffer>>[0]['play']>[0], signal: AbortSignal): Promise<void> {
-  if (!audioContext || !item.audio)
-    return
-
-  // Ensure audio context is resumed (browsers suspend it by default until user interaction)
-  if (audioContext.state === 'suspended') {
-    try {
-      await audioContext.resume()
-    }
-    catch {
+  return new Promise<void>(async (resolve) => {
+    if (!audioContext) {
+      resolve()
       return
     }
-  }
 
-  const source = audioContext.createBufferSource()
-  currentAudioSource.value = source
-  source.buffer = item.audio
-
-  source.connect(audioContext.destination)
-  if (audioAnalyser.value)
-    source.connect(audioAnalyser.value)
-  if (lipSyncNode.value)
-    source.connect(lipSyncNode.value)
-
-  const stopPlayback = () => {
-    try {
-      source.stop()
-      source.disconnect()
+    if (!item.audio) {
+      resolve()
+      return
     }
-    catch {}
-    if (currentAudioSource.value === source)
-      currentAudioSource.value = undefined
-  }
 
-  if (signal.aborted) {
-    stopPlayback()
-    return
-  }
+    // Ensure audio context is resumed (browsers suspend it by default until user interaction)
+    if (audioContext.state === 'suspended') {
+      try {
+        await audioContext.resume()
+      }
+      catch {
+        resolve()
+        return
+      }
+    }
 
-  signal.addEventListener('abort', stopPlayback, { once: true })
-  source.onended = () => {
-    signal.removeEventListener('abort', stopPlayback)
-    stopPlayback()
-  }
+    const source = audioContext.createBufferSource()
+    currentAudioSource.value = source
+    source.buffer = item.audio
 
-  try {
-    source.start(0)
-  }
-  catch {
-    stopPlayback()
-  }
+    source.connect(audioContext.destination)
+    if (audioAnalyser.value)
+      source.connect(audioAnalyser.value)
+    if (lipSyncNode.value)
+      source.connect(lipSyncNode.value)
+
+    const stopPlayback = () => {
+      try {
+        source.stop()
+        source.disconnect()
+      }
+      catch {}
+      if (currentAudioSource.value === source)
+        currentAudioSource.value = undefined
+      resolve()
+    }
+
+    if (signal.aborted) {
+      stopPlayback()
+      return
+    }
+
+    signal.addEventListener('abort', stopPlayback, { once: true })
+    source.onended = () => {
+      signal.removeEventListener('abort', stopPlayback)
+      stopPlayback()
+    }
+
+    try {
+      source.start(0)
+    }
+    catch {
+      stopPlayback()
+    }
+  })
 }
 
 const playbackManager = createPlaybackManager<AudioBuffer>({
diff --git a/packages/stage-ui/src/stores/llm.ts b/packages/stage-ui/src/stores/llm.ts
index 532019560..a2da200f6 100644
--- a/packages/stage-ui/src/stores/llm.ts
+++ b/packages/stage-ui/src/stores/llm.ts
@@ -53,29 +53,40 @@ async function streamFrom(model: string, chatProvider: ChatProvider, messages: M
     return tools ?? []
   }
 
-  const supportedTools = streamOptionsToolsCompatibilityOk(model, chatProvider, messages, options)
+  return new Promise<void>(async (resolve, reject) => {
+    try {
+      const supportedTools = streamOptionsToolsCompatibilityOk(model, chatProvider, messages, options)
 
-  // TODO: we need Automatic tools discovery
-  const tools = supportedTools
-    ? [
-        ...await mcp(),
-        ...await debug(),
-        ...await resolveTools(),
-      ]
-    : undefined
-
-  streamText({
-    ...chatProvider.chat(model),
-    maxSteps: 10,
-    messages: sanitized,
-    headers,
-    tools,
-    onEvent: async (event) => {
-      await options?.onStreamEvent?.(event as StreamEvent)
-
-      if (event.type === 'error')
-        throw event.error ?? new Error('Stream error')
-    },
+      await streamText({
+        ...chatProvider.chat(model),
+        maxSteps: 10,
+        messages: sanitized,
+        headers,
+        // TODO: we need Automatic tools discovery
+        tools: supportedTools
+          ? [
+              ...await mcp(),
+              ...await debug(),
+              ...await resolveTools(),
+            ]
+          : undefined,
+        async onEvent(event) {
+          try {
+            await options?.onStreamEvent?.(event as StreamEvent)
+            if (event.type === 'finish' && (event.finishReason !== 'tool_calls' || !options?.waitForTools))
+              resolve()
+            else if (event.type === 'error')
+              reject(event.error ?? new Error('Stream error'))
+          }
+          catch (err) {
+            reject(err)
+          }
+        },
+      })
+    }
+    catch (err) {
+      reject(err)
+    }
   })
 }
 
diff --git a/packages/stage-ui/src/tools/mcp.test.ts b/packages/stage-ui/src/tools/mcp.test.ts
index 385ccde9e..e1731c498 100644
--- a/packages/stage-ui/src/tools/mcp.test.ts
+++ b/packages/stage-ui/src/tools/mcp.test.ts
@@ -26,8 +26,8 @@ describe('tools mcp schema', () => {
     const callTool = tools.find(entry => entry.function.name === 'mcp_call_tool')
 
     expect(callTool).toBeDefined()
-    const items = ((callTool!.function.parameters as JsonSchema).properties?.parameters as JsonSchema)?.items as JsonSchema
+    const items = ((callTool?.function.parameters as JsonSchema).properties?.parameters as any)?.items
     expect(items).toBeDefined()
-    expect(items.additionalProperties).toBe(false)
+    expect(items?.additionalProperties).toBe(false)
   })
 })