From 7738cfc1210cffbb5a06b69b29bd29cc883f26e1 Mon Sep 17 00:00:00 2001
From: Matt Rossman <22670878+mattrossman@users.noreply.github.com>
Date: Thu, 11 Dec 2025 11:49:15 -0500
Subject: [PATCH] fix: aggregate tools across all steps

---
 apps/studio/evals/assistant.eval.ts    | 15 +++++++----
 apps/studio/lib/ai/tools/index.ts      | 32 -----------------------
 apps/studio/lib/ai/tools/mock-tools.ts | 35 ++++++++++++++++++++++++++
 3 files changed, 45 insertions(+), 37 deletions(-)
 create mode 100644 apps/studio/lib/ai/tools/mock-tools.ts

diff --git a/apps/studio/evals/assistant.eval.ts b/apps/studio/evals/assistant.eval.ts
index 3fc773dd5e7..ef8f7e69381 100644
--- a/apps/studio/evals/assistant.eval.ts
+++ b/apps/studio/evals/assistant.eval.ts
@@ -2,7 +2,7 @@ import { openai } from '@ai-sdk/openai'
 import { Eval } from 'braintrust'
 import { AssertionScorer } from 'lib/ai/evals/scorer'
 import { generateAssistantResponse } from 'lib/ai/generate-assistant-response'
-import { getMockTools } from 'lib/ai/tools'
+import { getMockTools } from 'lib/ai/tools/mock-tools'
 import assert from 'node:assert'
 
 assert(process.env.BRAINTRUST_PROJECT_ID, 'BRAINTRUST_PROJECT_ID is not set')
@@ -29,12 +29,17 @@ Eval('Assistant', {
       tools: await getMockTools(),
     })
 
-    const text = await result.text
-    const toolCalls = await result.toolCalls
+    const textLastStep = await result.text
+
+    // `result.toolCalls` only shows the last step, instead aggregate tools across all steps
+    const steps = await result.steps
+    const toolNamesAllStepsFlattened = steps.flatMap((step) =>
+      step.toolCalls.map((call) => call.toolName)
+    )
 
     return {
-      text: text,
-      tools: toolCalls.map((call) => call.toolName),
+      text: textLastStep,
+      tools: toolNamesAllStepsFlattened,
     }
   },
   scores: [AssertionScorer],
diff --git a/apps/studio/lib/ai/tools/index.ts b/apps/studio/lib/ai/tools/index.ts
index ee2b4b51a95..9082cfdb81f 100644
--- a/apps/studio/lib/ai/tools/index.ts
+++ b/apps/studio/lib/ai/tools/index.ts
@@ -58,35 +58,3 @@ export const getTools = async ({
 
   return filteredTools
 }
-
-/**
- * For evals
- */
-export async function getMockTools() {
-  // Always include rendering tools
-  let tools: ToolSet = getRenderingTools()
-
-  const mcpTools = await getMcpTools({
-    accessToken: 'mock-access-token',
-    projectRef: 'mock-project-ref',
-    aiOptInLevel: 'schema_and_log_and_data',
-  })
-
-  // Override execute handlers for all MCP tools to return mock responses
-  const mockedMcpTools = Object.fromEntries(
-    Object.entries(mcpTools).map(([name, baseTool]) => [
-      name,
-      {
-        ...baseTool,
-        execute: async () => ({ status: 'successful' }),
-      },
-    ])
-  ) as ToolSet
-
-  tools = {
-    ...tools,
-    ...mockedMcpTools,
-  }
-
-  return tools
-}
diff --git a/apps/studio/lib/ai/tools/mock-tools.ts b/apps/studio/lib/ai/tools/mock-tools.ts
new file mode 100644
index 00000000000..98991d550ac
--- /dev/null
+++ b/apps/studio/lib/ai/tools/mock-tools.ts
@@ -0,0 +1,35 @@
+import { ToolSet } from 'ai'
+import { getMcpTools } from './mcp-tools'
+import { getRenderingTools } from './rendering-tools'
+
+/**
+ * For evals. See `index.ts` for the real tools
+ */
+export async function getMockTools() {
+  // Always include rendering tools
+  let tools: ToolSet = getRenderingTools()
+
+  const mcpTools = await getMcpTools({
+    accessToken: 'mock-access-token',
+    projectRef: 'mock-project-ref',
+    aiOptInLevel: 'schema_and_log_and_data',
+  })
+
+  // Override execute handlers for all MCP tools to return mock responses
+  const mockedMcpTools = Object.fromEntries(
+    Object.entries(mcpTools).map(([name, baseTool]) => [
+      name,
+      {
+        ...baseTool,
+        execute: async () => ({ status: 'successful' }),
+      },
+    ])
+  ) as ToolSet
+
+  tools = {
+    ...tools,
+    ...mockedMcpTools,
+  }
+
+  return tools
+}