refactor(amazonq): refactor doc and featureDev SessionState

packages/core/src/amazonqFeatureCommon/client/client.ts → packages/core/src/amazonq/client/client.ts

@@ -5,7 +5,7 @@
 
 import { ServiceOptions } from '../../shared'
 
-export abstract class Client {
+export abstract class FeatureClient {
     public abstract getClient(options?: Partial<ServiceOptions>): Promise<any>
 
     public abstract createConversation(): Promise<string>

packages/core/src/amazonq/commons/connector/baseMessenger.ts

@@ -27,7 +27,7 @@ import {
 } from './connectorMessages'
 import { FollowUpTypes } from '../types'
 import { messageWithConversationId } from '../../../amazonqFeatureDev/userFacingText'
-import { DeletedFileInfo, NewFileInfo } from '../../../amazonqFeatureCommon/types'
+import { DeletedFileInfo, NewFileInfo } from '../../types'
 
 export class Messenger {
     public constructor(

packages/core/src/amazonq/commons/connector/connectorMessages.ts

@@ -8,7 +8,7 @@ import { MessagePublisher } from '../../messages/messagePublisher'
 import { CodeReference } from '../../webview/ui/connector'
 import { ChatItemAction, ProgressField, SourceLink } from '@aws/mynah-ui'
 import { ChatItemType } from '../model'
-import { DeletedFileInfo, NewFileInfo } from '../../../amazonqFeatureCommon/types'
+import { DeletedFileInfo, NewFileInfo } from '../../types'
 import { licenseText } from '../../../amazonqFeatureDev/constants'
 
 class UiMessage {

packages/core/src/amazonq/commons/session/sessionConfigFactory.ts

@@ -6,7 +6,7 @@
 import * as vscode from 'vscode'
 import { WorkspaceFolderNotFoundError } from '../../../amazonqFeatureDev/errors'
 import { VirtualFileSystem, VirtualMemoryFile } from '../../../shared'
-import { CurrentWsFolders } from '../../../amazonqFeatureCommon/types'
+import { CurrentWsFolders } from '../../types'
 
 export interface SessionConfig {
     // The paths on disk to where the source code lives

packages/core/src/amazonq/index.ts

@@ -66,3 +66,5 @@ export function createMynahUI(
 }
 
 export * from './commons/types'
+export * from './types'
+export * from './constants'

packages/core/src/amazonqFeatureCommon/session/sessionState.ts → packages/core/src/amazonq/session/sessionState.ts

@@ -9,7 +9,7 @@ import globals from '../../shared/extensionGlobals'
 import { getLogger } from '../../shared/logger'
 import { telemetry } from '../../shared/telemetry/telemetry'
 import { VirtualFileSystem } from '../../shared/virtualFilesystem'
-import { CodeReference, UploadHistory } from '../../amazonq/webview/ui/connector'
+import { CodeReference, UploadHistory } from '../webview/ui/connector'
 import { AuthUtil } from '../../codewhisperer/util/authUtil'
 import { randomUUID } from '../../shared/crypto'
 import { i18n } from '../../shared/i18n-helper'
@@ -185,7 +185,30 @@ export abstract class BasePrepareCodeGenState implements SessionState {
         this.config.workspaceRoots = [workspaceRoot]
     }
 
-    protected abstract createNextState(config: SessionStateConfig): SessionState
+    protected createNextState(
+        config: SessionStateConfig,
+        StateClass?: new (
+            config: SessionStateConfig,
+            filePaths: NewFileInfo[],
+            deletedFiles: DeletedFileInfo[],
+            references: CodeReference[],
+            tabID: string,
+            currentIteration: number,
+            uploadHistory: UploadHistory,
+            codeGenerationRemainingIterationCount?: number,
+            codeGenerationTotalIterationCount?: number
+        ) => SessionState
+    ): SessionState {
+        return new StateClass!(
+            config,
+            this.filePaths,
+            this.deletedFiles,
+            this.references,
+            this.tabID,
+            this.currentIteration,
+            this.uploadHistory
+        )
+    }
 
     async interact(action: SessionStateAction): Promise<SessionStateInteraction> {
         const uploadId = await telemetry.amazonq_createUpload.run(async (span) => {
@@ -226,6 +249,19 @@ export interface CodeGenerationParams {
     workspaceFolders: CurrentWsFolders
 }
 
+export interface CreateNextStateParams {
+    filePaths: NewFileInfo[]
+    deletedFiles: DeletedFileInfo[]
+    references: CodeReference[]
+    currentIteration: number
+    remainingIterations?: number
+    totalIterations?: number
+    uploadHistory: UploadHistory
+    tokenSource: vscode.CancellationTokenSource
+    currentCodeGenerationId?: string
+    codeGenerationId?: string
+}
+
 export abstract class BaseCodeGenState extends CodeGenBase implements SessionState {
     constructor(
         config: SessionStateConfig,
@@ -241,21 +277,39 @@ export abstract class BaseCodeGenState extends CodeGenBase implements SessionSta
         super(config, tabID)
     }
 
-    protected abstract createNextState(
+    protected createNextState(
         config: SessionStateConfig,
-        params: {
-            filePaths: NewFileInfo[]
-            deletedFiles: DeletedFileInfo[]
-            references: CodeReference[]
-            currentIteration: number
-            remainingIterations?: number
-            totalIterations?: number
-            uploadHistory: UploadHistory
-            tokenSource: vscode.CancellationTokenSource
-            currentCodeGenerationId?: string
+        params: CreateNextStateParams,
+        StateClass?: new (
+            config: SessionStateConfig,
+            filePaths: NewFileInfo[],
+            deletedFiles: DeletedFileInfo[],
+            references: CodeReference[],
+            tabID: string,
+            currentIteration: number,
+            remainingIterations?: number,
+            totalIterations?: number,
+            uploadHistory?: UploadHistory,
+            tokenSource?: vscode.CancellationTokenSource,
+            currentCodeGenerationId?: string,
             codeGenerationId?: string
-        }
-    ): SessionState
+        ) => SessionState
+    ): SessionState {
+        return new StateClass!(
+            config,
+            params.filePaths,
+            params.deletedFiles,
+            params.references,
+            this.tabID,
+            params.currentIteration,
+            params.remainingIterations,
+            params.totalIterations,
+            params.uploadHistory,
+            params.tokenSource,
+            params.currentCodeGenerationId,
+            params.codeGenerationId
+        )
+    }
 
     async interact(action: SessionStateAction): Promise<SessionStateInteraction> {
         return telemetry.amazonq_codeGenerationInvoke.run(async (span) => {

packages/core/src/amazonqFeatureCommon/types.ts → packages/core/src/amazonq/types.ts

@@ -6,10 +6,10 @@
 import * as vscode from 'vscode'
 import { VirtualFileSystem } from '../shared/virtualFilesystem'
 import type { CancellationTokenSource } from 'vscode'
-import { CodeReference, UploadHistory } from '../amazonq/webview/ui/connector'
-import { DiffTreeFileInfo } from '../amazonq/webview/ui/diffTree/types'
-import { Messenger } from '../amazonq/commons/connector/baseMessenger'
-import { Client } from './client/client'
+import { CodeReference, UploadHistory } from './webview/ui/connector'
+import { DiffTreeFileInfo } from './webview/ui/diffTree/types'
+import { Messenger } from './commons/connector/baseMessenger'
+import { FeatureClient } from './client/client'
 import { TelemetryHelper } from './util/telemetryHelper'
 
 export type Interaction = {
@@ -70,7 +70,7 @@ export interface SessionStateConfig {
     workspaceRoots: string[]
     workspaceFolders: CurrentWsFolders
     conversationId: string
-    proxyClient: Client
+    proxyClient: FeatureClient
     uploadId: string
     currentCodeGenerationId?: string
 }

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

@@ -12,7 +12,7 @@ import { ContentLengthError, PrepareRepoFailedError } from '../../amazonqFeature
 import { getLogger } from '../../shared/logger/logger'
 import { maxFileSizeBytes } from '../../amazonqFeatureDev/limits'
 import { createHash } from 'crypto'
-import { CurrentWsFolders, DeletedFileInfo, NewFileInfo, NewFileZipContents } from '../../amazonqFeatureCommon/types'
+import { CurrentWsFolders, DeletedFileInfo, NewFileInfo, NewFileZipContents } from '../types'
 import { hasCode, ToolkitError } from '../../shared/errors'
 import { AmazonqCreateUpload, Span, telemetry as amznTelemetry, telemetry } from '../../shared/telemetry/telemetry'
 import { TelemetryHelper } from './telemetryHelper'

packages/core/src/amazonqDoc/controllers/chat/controller.ts

@@ -40,7 +40,7 @@ import {
     getWorkspaceRelativePath,
     isMultiRootWorkspace,
 } from '../../../shared/utilities/workspaceUtils'
-import { getPathsFromZipFilePath } from '../../../amazonqFeatureCommon/util/files'
+import { getPathsFromZipFilePath } from '../../../amazonq/util/files'
 import { FollowUpTypes } from '../../../amazonq/commons/types'
 import { DocGenerationTask } from '../docGenerationTask'
 

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

@@ -13,7 +13,7 @@ import { SessionConfig } from '../../amazonq/commons/session/sessionConfigFactor
 import { ReferenceLogViewProvider } from '../../codewhisperer'
 import path from 'path'
 import { FeatureDevClient } from '../../amazonqFeatureDev/client/featureDev'
-import { TelemetryHelper } from '../../amazonqFeatureCommon/util/telemetryHelper'
+import { TelemetryHelper } from '../../amazonq/util/telemetryHelper'
 import { ConversationNotStartedState } from '../../amazonqFeatureDev/session/sessionState'
 import { logWithConversationId } from '../../amazonqFeatureDev/userFacingText'
 import { ConversationIdNotFoundError } from '../../amazonqFeatureDev/errors'

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

@@ -3,14 +3,12 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import * as vscode from 'vscode'
 import { ToolkitError } from '../../shared/errors'
 import { DocGenerationStep, docScheme, getFileSummaryPercentage, Mode } from '../constants'
 
-import { CodeReference, UploadHistory } from '../../amazonq/webview/ui/connector'
 import { i18n } from '../../shared/i18n-helper'
 
-import { DeletedFileInfo, NewFileInfo, SessionState, SessionStateAction, SessionStateConfig } from '../types'
+import { NewFileInfo, SessionState, SessionStateAction, SessionStateConfig } from '../types'
 import { Intent } from '../../amazonqFeatureDev'
 import {
     ContentLengthError,
@@ -24,7 +22,7 @@ import {
     WorkspaceEmptyError,
 } from '../errors'
 import { DocMessenger } from '../messenger'
-import { BaseCodeGenState, BasePrepareCodeGenState } from '../../amazonqFeatureCommon/session/sessionState'
+import { BaseCodeGenState, BasePrepareCodeGenState, CreateNextStateParams } from '../../amazonq/session/sessionState'
 
 export class DocCodeGenState extends BaseCodeGenState {
     protected handleProgress(messenger: DocMessenger, detail?: string): void {
@@ -112,48 +110,13 @@ export class DocCodeGenState extends BaseCodeGenState {
         )
     }
 
-    protected createNextState(
-        config: SessionStateConfig,
-        params: {
-            filePaths: NewFileInfo[]
-            deletedFiles: DeletedFileInfo[]
-            references: CodeReference[]
-            currentIteration: number
-            remainingIterations?: number
-            totalIterations?: number
-            uploadHistory: UploadHistory
-            tokenSource: vscode.CancellationTokenSource
-            currentCodeGenerationId?: string
-            codeGenerationId?: string
-        }
-    ): SessionState {
-        return new DocPrepareCodeGenState(
-            config,
-            params.filePaths,
-            params.deletedFiles,
-            params.references,
-            this.tabID,
-            params.currentIteration,
-            params.remainingIterations,
-            params.totalIterations,
-            params.uploadHistory,
-            params.tokenSource,
-            params.currentCodeGenerationId,
-            params.codeGenerationId
-        )
+    protected override createNextState(config: SessionStateConfig, params: CreateNextStateParams): SessionState {
+        return super.createNextState(config, params, DocPrepareCodeGenState)
     }
 }
 
 export class DocPrepareCodeGenState extends BasePrepareCodeGenState {
-    protected createNextState(config: SessionStateConfig): SessionState {
-        return new DocCodeGenState(
-            config,
-            this.filePaths,
-            this.deletedFiles,
-            this.references,
-            this.tabID,
-            this.currentIteration,
-            this.uploadHistory
-        )
+    protected override createNextState(config: SessionStateConfig): SessionState {
+        return super.createNextState(config, DocCodeGenState)
     }
 }

packages/core/src/amazonqDoc/types.ts

@@ -20,7 +20,7 @@ import {
     SessionState as FeatureDevSessionState,
     SessionStateAction as FeatureDevSessionStateAction,
     SessionStateInteraction as FeatureDevSessionStateInteraction,
-} from '../amazonqFeatureCommon/types'
+} from '../amazonq/types'
 
 import { Mode } from './constants'
 import { DocMessenger } from './messenger'

packages/core/src/amazonqFeatureDev/client/featureDev.ts

@@ -32,7 +32,7 @@ import {
     MetricData,
     TelemetryEvent,
 } from './featuredevproxyclient'
-import { Client } from '../../amazonqFeatureCommon/client/client'
+import { FeatureClient } from '../../amazonq/client/client'
 
 // Re-enable once BE is able to handle retries.
 const writeAPIRetryOptions = {
@@ -63,7 +63,7 @@ export async function createFeatureDevProxyClient(options?: Partial<ServiceOptio
     )) as FeatureDevProxyClient
 }
 
-export class FeatureDevClient implements Client {
+export class FeatureDevClient implements FeatureClient {
     public async getClient(options?: Partial<ServiceOptions>) {
         // Should not be stored for the whole session.
         // Client has to be reinitialized for each request so we always have a fresh bearerToken

packages/core/src/amazonqFeatureDev/controllers/chat/controller.ts

@@ -36,15 +36,15 @@ import {
     MetricDataOperationName,
     MetricDataResult,
     type NewFileInfo,
-} from '../../../amazonqFeatureCommon/types'
+} from '../../../amazonq'
 import { AuthUtil } from '../../../codewhisperer/util/authUtil'
 import { AuthController } from '../../../amazonq/auth/controller'
 import { getLogger } from '../../../shared/logger'
 import { submitFeedback } from '../../../feedback/vue/submitFeedback'
 import { placeholder } from '../../../shared/vscode/commands2'
 import { EditorContentController } from '../../../amazonq/commons/controllers/contentController'
 import { openUrl } from '../../../shared/utilities/vsCodeUtils'
-import { getPathsFromZipFilePath } from '../../../amazonqFeatureCommon/util/files'
+import { getPathsFromZipFilePath } from '../../../amazonq/util/files'
 import { examples, messageWithConversationId } from '../../userFacingText'
 import { getWorkspaceFoldersByPrefixes } from '../../../shared/utilities/workspaceUtils'
 import { openDeletedDiff, openDiff } from '../../../amazonq/commons/diff'

packages/core/src/amazonqFeatureDev/index.ts

@@ -3,14 +3,14 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-export * from '../amazonqFeatureCommon/types'
+export * from '../amazonq/types'
 export * from './userFacingText'
 export * from './errors'
 export * from './session/sessionState'
 export * from './constants'
 export { Session } from './session/session'
 export { FeatureDevClient } from './client/featureDev'
 export { FeatureDevChatSessionStorage } from './storages/chatSession'
-export { TelemetryHelper } from '../amazonqFeatureCommon/util/telemetryHelper'
-export { prepareRepoData } from '../amazonqFeatureCommon/util/files'
+export { TelemetryHelper } from '../amazonq/util/telemetryHelper'
+export { prepareRepoData } from '../amazonq/util/files'
 export { ChatControllerEventEmitters, FeatureDevController } from './controllers/chat/controller'

packages/core/src/amazonqFeatureDev/session/session.ts

@@ -13,14 +13,14 @@ import {
     type SessionState,
     type SessionStateConfig,
     UpdateFilesPathsParams,
-} from '../../amazonqFeatureCommon/types'
+} from '../../amazonq/types'
 import { ConversationIdNotFoundError } from '../errors'
 import { featureDevChat, referenceLogText, featureDevScheme } from '../constants'
 import fs from '../../shared/fs/fs'
 import { FeatureDevClient } from '../client/featureDev'
 import { codeGenRetryLimit } from '../limits'
 import { telemetry } from '../../shared/telemetry/telemetry'
-import { TelemetryHelper } from '../../amazonqFeatureCommon/util/telemetryHelper'
+import { TelemetryHelper } from '../../amazonq/util/telemetryHelper'
 import { ReferenceLogViewProvider } from '../../codewhisperer/service/referenceLogViewProvider'
 import { AuthUtil } from '../../codewhisperer/util/authUtil'
 import { getLogger } from '../../shared'