diff --git a/packages/ai/src/tools/communication-tools/done-tool/done-tool-delta.ts b/packages/ai/src/tools/communication-tools/done-tool/done-tool-delta.ts
index 0f00ed407..4678d6052 100644
--- a/packages/ai/src/tools/communication-tools/done-tool/done-tool-delta.ts
+++ b/packages/ai/src/tools/communication-tools/done-tool/done-tool-delta.ts
@@ -1,6 +1,5 @@
 import {
   type UpdateMessageEntriesParams,
-  getAssetLatestVersion,
   updateChat,
   updateMessage,
   updateMessageEntries,
@@ -52,6 +51,7 @@ export function createDoneToolDelta(context: DoneToolContext, doneToolState: Don
       assetId: string;
       assetName: string;
       assetType: ResponseMessageFileType;
+      versionNumber: number;
     };
 
     function isAssetToReturn(value: unknown): value is AssetToReturn {
@@ -63,7 +63,8 @@ export function createDoneToolDelta(context: DoneToolContext, doneToolState: Don
       const typeOk =
         typeof typeVal === 'string' &&
         ResponseMessageFileTypeSchema.options.includes(typeVal as ResponseMessageFileType);
-      return idOk && nameOk && typeOk;
+      const versionOk = typeof obj.versionNumber === 'number' && obj.versionNumber > 0;
+      return idOk && nameOk && typeOk && versionOk;
     }
 
     let assetsToInsert: AssetToReturn[] = [];
@@ -95,7 +96,7 @@ export function createDoneToolDelta(context: DoneToolContext, doneToolState: Don
           type: 'file' as const,
           file_type: a.assetType,
           file_name: a.assetName,
-          version_number: 1,
+          version_number: a.versionNumber,
           filter_version_id: null,
           metadata: [
             {
@@ -128,7 +129,7 @@ export function createDoneToolDelta(context: DoneToolContext, doneToolState: Don
       if (newAssets.length > 0) {
         doneToolState.addedAssets = [
           ...(doneToolState.addedAssets || []),
-          ...newAssets.map((a) => ({ assetId: a.assetId, assetType: a.assetType })),
+          ...newAssets.map((a) => ({ assetId: a.assetId, assetType: a.assetType, versionNumber: a.versionNumber })),
         ];
       }
     }
@@ -159,16 +160,10 @@ export function createDoneToolDelta(context: DoneToolContext, doneToolState: Don
               const firstAsset = doneToolState.addedAssets[0];
 
               if (firstAsset) {
-                // Get the actual version number from the database
-                const versionNumber = await getAssetLatestVersion({
-                  assetId: firstAsset.assetId,
-                  assetType: firstAsset.assetType,
-                });
-
                 await updateChat(context.chatId, {
                   mostRecentFileId: firstAsset.assetId,
                   mostRecentFileType: firstAsset.assetType,
-                  mostRecentVersionNumber: versionNumber,
+                  mostRecentVersionNumber: firstAsset.versionNumber,
                 });
               }
             } catch (error) {
diff --git a/packages/ai/src/tools/communication-tools/done-tool/done-tool.ts b/packages/ai/src/tools/communication-tools/done-tool/done-tool.ts
index 3cbb2263d..6b0dab26c 100644
--- a/packages/ai/src/tools/communication-tools/done-tool/done-tool.ts
+++ b/packages/ai/src/tools/communication-tools/done-tool/done-tool.ts
@@ -15,6 +15,7 @@ export const DoneToolInputSchema = z.object({
         assetId: z.string().uuid(),
         assetName: z.string(),
         assetType: AssetTypeSchema,
+        versionNumber: z.number().int().positive().describe('The version number of the asset to return'),
       })
     )
     .describe(
@@ -60,17 +61,12 @@ const DoneToolStateSchema = z.object({
     .array(
       z.object({
         assetId: z.string(),
-        assetType: z.enum([
-          'metric_file',
-          'dashboard_file',
-          'report_file',
-          'analyst_chat',
-          'collection',
-        ]),
+        assetType: AssetTypeSchema,
+        versionNumber: z.number(),
       })
     )
     .optional()
-    .describe('Assets that have been added with their types for chat update'),
+    .describe('Assets that have been added with their types and version numbers for chat update'),
 });
 
 export type DoneToolInput = z.infer<typeof DoneToolInputSchema>;