diff --git a/packages/jotai-ai/package.json b/packages/jotai-ai/package.json
index b365d61..0def70b 100644
--- a/packages/jotai-ai/package.json
+++ b/packages/jotai-ai/package.json
@@ -23,8 +23,7 @@
     "build": "bunchee",
     "dev": "bunchee --watch",
     "lint": "prettier --c .",
-    "test": "echo \"Error: no test specified\" && exit 1",
-    "ui-test": "vitest --config vitest.ui.config.ts"
+    "test": "vitest"
   },
   "keywords": [
     "ai",
diff --git a/packages/jotai-ai/src/utils/index.ts b/packages/jotai-ai/src/utils/index.ts
index 7f941a5..5ca5307 100644
--- a/packages/jotai-ai/src/utils/index.ts
+++ b/packages/jotai-ai/src/utils/index.ts
@@ -1,9 +1,6 @@
 import type { Message } from '@ai-sdk/ui-utils';
-import type { CoreToolMessage } from 'ai';
 
-import { generateId as defaultGenerateId } from '@ai-sdk/ui-utils';
-
-export { defaultGenerateId };
+export { generateId } from '@ai-sdk/ui-utils';
 
 export const isPromiseLike = (
   value: unknown,
@@ -16,51 +13,20 @@ export const isPromiseLike = (
   );
 };
 
-// Modified from https://github.com/vercel/ai-chatbot/blob/e6806aaa542c831e874681a50df39fe26c59d25e/lib/utils.ts
-export function addToolMessageToChat({
-  toolMessage,
-  messages,
-}: {
-  toolMessage: CoreToolMessage;
-  messages: Message[];
-}): Message[] {
-  return messages.map(message => {
-    if (message.toolInvocations) {
-      return {
-        ...message,
-        toolInvocations: message.toolInvocations.map(toolInvocation => {
-          const toolResult = toolMessage.content.find(
-            tool => tool.toolCallId === toolInvocation.toolCallId,
-          );
-
-          if (toolResult) {
-            return {
-              ...toolInvocation,
-              state: 'result',
-              result: toolResult.result,
-            };
-          }
-
-          return toolInvocation;
-        }),
-      };
-    }
-
-    return message;
-  });
-}
-
 /**
  Check if the message is an assistant message with completed tool calls.
  The message must have at least one tool invocation and all tool invocations
  must have a result.
  */
 export function isAssistantMessageWithCompletedToolCalls(message: Message) {
+  const toolInvocationParts = message.parts?.filter(
+    part => part.type === 'tool-invocation',
+  );
   return (
     message.role === 'assistant' &&
-    message.toolInvocations &&
-    message.toolInvocations.length > 0 &&
-    message.toolInvocations.every(toolInvocation => 'result' in toolInvocation)
+    toolInvocationParts &&
+    toolInvocationParts.length > 0 &&
+    toolInvocationParts.every(part => part.toolInvocation.state === 'result')
   );
 }
 
diff --git a/packages/jotai-ai/vitest.ui.config.ts b/packages/jotai-ai/vitest.config.ts
similarity index 67%
rename from packages/jotai-ai/vitest.ui.config.ts
rename to packages/jotai-ai/vitest.config.ts
index 3e9aaae..d358967 100644
--- a/packages/jotai-ai/vitest.ui.config.ts
+++ b/packages/jotai-ai/vitest.config.ts
@@ -6,9 +6,6 @@ export default defineConfig({
   test: {
     environment: 'jsdom',
     globals: true,
-    includeSource: [
-      'src/**/*.ui.test.tsx',
-      // 'tests/**/*.?(c|m)[jt]s?(x)',
-    ],
+    includeSource: ['src/**/*.ui.test.tsx'],
   },
 });