diff --git a/package.json b/package.json
index 7f6465377..d25f6d7c8 100644
--- a/package.json
+++ b/package.json
@@ -92,7 +92,7 @@
     "@cucumber/messages": "31.2.0",
     "@xmldom/xmldom": "0.9.8",
     "acorn": "8.15.0",
-    "ai": "^5.0.60",
+    "ai": "^6.0.43",
     "arrify": "3.0.0",
     "axios": "1.13.2",
     "chalk": "4.1.2",
diff --git a/test/support/mock-ai-provider.js b/test/support/mock-ai-provider.js
index cc7f97f97..a92f0b671 100644
--- a/test/support/mock-ai-provider.js
+++ b/test/support/mock-ai-provider.js
@@ -4,12 +4,10 @@ export function createMockModel(config = {}) {
   let callCount = 0
 
   const mockModel = {
-    specificationVersion: 'v2',
-    provider: 'mock',
+    specificationVersion: 'v3',
     modelId: 'mock-model',
-    config: {
-      provider: 'mock',
-    },
+    provider: 'mock',
+    defaultObjectGenerationMode: 'json',
 
     async doGenerate(options) {
       if (delay > 0) {
@@ -25,6 +23,11 @@ export function createMockModel(config = {}) {
 
       const textContent = response.text || 'Mock response'
 
+      // Calculate token values
+      const promptTokens = response.promptTokens || 50
+      const totalTokens = response.totalTokens || 100
+      const completionTokens = response.completionTokens || (totalTokens - promptTokens)
+
       return {
         text: textContent,
         content: [
@@ -34,12 +37,20 @@ export function createMockModel(config = {}) {
           },
         ],
         usage: {
-          promptTokens: response.promptTokens || 50,
-          completionTokens: response.completionTokens || 50,
-          totalTokens: response.totalTokens || 100,
+          promptTokens,
+          completionTokens,
+          totalTokens,
+          inputTokens: {
+            total: promptTokens,
+          },
+          outputTokens: {
+            total: completionTokens,
+          },
         },
         finishReason: response.finishReason || 'stop',
-        rawResponse: { headers: {} },
+        rawResponse: {
+          headers: {},
+        },
         warnings: [],
         logprobs: undefined,
         response: {
@@ -54,6 +65,7 @@ export function createMockModel(config = {}) {
         request: {
           body: JSON.stringify({ messages: options.prompt }),
         },
+        experimental_providerMetadata: undefined,
       }
     },
 
diff --git a/test/unit/ai_test.js b/test/unit/ai_test.js
index 47024ad5c..e94a8b942 100644
--- a/test/unit/ai_test.js
+++ b/test/unit/ai_test.js
@@ -27,7 +27,7 @@ describe('AI module', () => {
       },
     }
 
-    AiAssistant.enable(config)
+    await AiAssistant.enable(config)
     await AiAssistant.setHtmlContext(html)
     expect(AiAssistant.minifiedHtml).to.include('<a data-qa="ok">Hey</a>')
   })
@@ -64,7 +64,7 @@ describe('AI module', () => {
 describe('AI module with mock provider', () => {
   const tempDir = path.join(__dirname, '../data/sandbox/ai-prompts-test')
 
-  beforeEach(async () => {
+  beforeEach(() => {
     AiAssistant.reset()
     config.reset()