chore: address comments

Author: KevinDing1

packages/core/src/amazonq/util/files.ts

@@ -13,7 +13,6 @@ import {
 } from '../../shared/utilities/workspaceUtils'
 
 import { ContentLengthError, PrepareRepoFailedError } from '../../amazonqFeatureDev/errors'
-import { ContentLengthError as DocContentLengthError } from '../../amazonqDoc/errors'
 import { getLogger } from '../../shared/logger/logger'
 import { maxFileSizeBytes } from '../../amazonqFeatureDev/limits'
 import { CurrentWsFolders, DeletedFileInfo, NewFileInfo, NewFileZipContents } from '../../amazonqDoc/types'
@@ -49,7 +48,6 @@ export type PrepareRepoDataOptions = {
     telemetry?: TelemetryHelper
     zip?: ZipStream
     isIncludeInfraDiagram?: boolean
-    featureName?: 'featureDev' | 'docGeneration'
 }
 
 /**
@@ -188,9 +186,6 @@ export async function prepareRepoData(
     } catch (error) {
         getLogger().debug(`featureDev: Failed to prepare repo: ${error}`)
         if (error instanceof ToolkitError && error.code === 'ContentLengthError') {
-            if (options?.featureName === 'docGeneration') {
-                throw new DocContentLengthError()
-            }
             throw new ContentLengthError()
         }
         throw new PrepareRepoFailedError()

packages/core/src/amazonqDoc/session/session.ts

@@ -30,6 +30,8 @@ import fs from '../../shared/fs/fs'
 import globals from '../../shared/extensionGlobals'
 import { extensionVersion } from '../../shared/vscode/env'
 import { getLogger } from '../../shared/logger/logger'
+import { ContentLengthError } from '../errors'
+import { ContentLengthError as FeatureDevContentLengthError } from '../../amazonqFeatureDev/errors'
 
 export class Session {
     private _state?: SessionState | Omit<SessionState, 'uploadId'>
@@ -126,28 +128,35 @@ export class Session {
         return this.nextInteraction(msg, mode, folderPath)
     }
     private async nextInteraction(msg: string, mode: Mode, folderPath?: string) {
-        const resp = await this.state.interact({
-            task: this.task,
-            msg,
-            fs: this.config.fs,
-            mode: mode,
-            folderPath: folderPath,
-            messenger: this.messenger,
-            telemetry: this.telemetry,
-            tokenSource: this.state.tokenSource,
-            uploadHistory: this.state.uploadHistory,
-        })
+        try {
+            const resp = await this.state.interact({
+                task: this.task,
+                msg,
+                fs: this.config.fs,
+                mode: mode,
+                folderPath: folderPath,
+                messenger: this.messenger,
+                telemetry: this.telemetry,
+                tokenSource: this.state.tokenSource,
+                uploadHistory: this.state.uploadHistory,
+            })
+
+            if (resp.nextState) {
+                if (!this.state?.tokenSource?.token.isCancellationRequested) {
+                    this.state?.tokenSource?.cancel()
+                }
 
-        if (resp.nextState) {
-            if (!this.state?.tokenSource?.token.isCancellationRequested) {
-                this.state?.tokenSource?.cancel()
+                // Move to the next state
+                this._state = resp.nextState
             }
 
-            // Move to the next state
-            this._state = resp.nextState
+            return resp.interaction
+        } catch (e) {
+            if (e instanceof FeatureDevContentLengthError) {
+                throw new ContentLengthError()
+            }
+            throw e
         }
-
-        return resp.interaction
     }
 
     public async updateFilesPaths(

packages/core/src/amazonqDoc/session/sessionState.ts

@@ -144,7 +144,6 @@ export class DocPrepareCodeGenState extends BasePrepareCodeGenState {
         return await prepareRepoData(workspaceRoots, workspaceFolders, span, {
             ...options,
             isIncludeInfraDiagram: true,
-            featureName: 'docGeneration',
         })
     }
 }