feat: memory bank support (aws#2314)

* feat: support memory bank (aws#2299)

* feat: poc for adding memory bank as rules

* feat: changed to use simlar prompt from kiro

* feat: implement iteration based guidelines generation from science doc

* fix: fixing issues from poc code

* fix: refine UX messages

* fix: improving memory bank pre-processing logic (aws#2301)

* fix: limit file size for during memory bank analysis (aws#2304)

* fix: ehance memory bank prompt for more determinastic file name and path (aws#2305)

* fix: fix windows path issue for memory bank (aws#2307)

* fix: small path fix for button text, remove unnecessary logs (aws#2310)

---------

Co-authored-by: aws-toolkit-automation <[email protected]>

Author: bywang56

chat-client/src/client/features/rules.test.ts

@@ -22,6 +22,7 @@ describe('rules', () => {
 
         messager = {
             onRuleClick: sinon.stub(),
+            onChatPrompt: sinon.stub(),
         } as unknown as Messager
 
         rulesList = new RulesList(mynahUi, messager)
@@ -142,6 +143,22 @@ describe('rules', () => {
             assert.equal(formArgs[2][1].id, ContextRule.SubmitButtonId)
         })
 
+        it('calls messager when create memory bank is clicked', () => {
+            const createMemoryBankItem: DetailedListItem = {
+                id: ContextRule.CreateMemoryBankId,
+                description: 'Generate Memory Bank',
+            }
+
+            onItemClick(createMemoryBankItem)
+
+            // Should send a chat prompt
+            sinon.assert.calledOnce(messager.onChatPrompt as sinon.SinonStub)
+
+            const chatPromptArgs = (messager.onChatPrompt as sinon.SinonStub).getCall(0).args[0]
+            assert.equal(chatPromptArgs.prompt.prompt, 'Generate a Memory Bank for this project')
+            assert.equal(chatPromptArgs.prompt.escapedPrompt, 'Generate a Memory Bank for this project')
+        })
+
         it('calls messager when regular rule is clicked', () => {
             const ruleItem: DetailedListItem = {
                 id: 'test-rule-id',
@@ -267,21 +284,27 @@ describe('rules', () => {
 
             const result = convertRulesListToDetailedListGroup(rulesFolder)
 
-            assert.equal(result.length, 3) // 2 folders + create rule group
+            assert.equal(result.length, 3) // 2 folders + actions group
             assert.equal(result[0].groupName, 'test-folder')
             assert.equal(result[0].children?.length, 2)
             assert.equal(result[0].children?.[0].id, 'rule-1')
             assert.equal(result[0].children?.[0].description, 'Test Rule 1')
             assert.equal(result[1].groupName, 'inactive-folder')
             assert.equal(result[1].children?.length, 0)
-            assert.equal(result[2].children?.[0].id, ContextRule.CreateRuleId)
+            assert.equal(result[2].groupName, 'Actions')
+            assert.equal(result[2].children?.length, 2) // Memory Bank + Create Rule
+            assert.equal(result[2].children?.[0].id, ContextRule.CreateMemoryBankId)
+            assert.equal(result[2].children?.[1].id, ContextRule.CreateRuleId)
         })
 
         it('handles empty rules array', () => {
             const result = convertRulesListToDetailedListGroup([])
 
-            assert.equal(result.length, 1) // Only create rule group
-            assert.equal(result[0].children?.[0].id, ContextRule.CreateRuleId)
+            assert.equal(result.length, 1) // Only actions group
+            assert.equal(result[0].groupName, 'Actions')
+            assert.equal(result[0].children?.length, 2) // Memory Bank + Create Rule
+            assert.equal(result[0].children?.[0].id, ContextRule.CreateMemoryBankId)
+            assert.equal(result[0].children?.[1].id, ContextRule.CreateRuleId)
         })
     })
 })

chat-client/src/client/features/rules.ts

@@ -6,6 +6,7 @@ import { MynahDetailedList } from './history'
 
 export const ContextRule = {
     CreateRuleId: 'create-rule',
+    CreateMemoryBankId: 'create-memory-bank',
     CancelButtonId: 'cancel-create-rule',
     SubmitButtonId: 'submit-create-rule',
     RuleNameFieldId: 'rule-name',
@@ -68,12 +69,29 @@ export class RulesList {
                     ],
                     `Create a rule`
                 )
+            } else if (item.id === ContextRule.CreateMemoryBankId) {
+                this.rulesList?.close()
+                this.handleMemoryBankCreation()
             } else {
                 this.messager.onRuleClick({ tabId: this.tabId, type: 'rule', id: item.id })
             }
         }
     }
 
+    private handleMemoryBankCreation = () => {
+        // Close the rules list first
+        this.rulesList?.close()
+
+        // Use the current tab, the tabId should be the same as the one used for the rules list
+        this.messager.onChatPrompt({
+            prompt: {
+                prompt: 'Generate a Memory Bank for this project',
+                escapedPrompt: 'Generate a Memory Bank for this project',
+            },
+            tabId: this.tabId,
+        })
+    }
+
     showLoading(tabId: string) {
         this.tabId = tabId
         const rulesList = this.mynahUi.openTopBarButtonOverlay({
@@ -156,6 +174,24 @@ const createRuleListItem: DetailedListItem = {
     id: ContextRule.CreateRuleId,
 }
 
+function createMemoryBankListItem(rules: RulesFolder[]): DetailedListItem {
+    // Handles button text changes between "Generation" and "Regenerate"
+    const memoryBankFiles = ['product', 'structure', 'tech', 'guidelines']
+
+    const memoryBankFolder = rules.find(folder => folder.folderName === 'memory-bank')
+
+    const hasMemoryBankFiles =
+        memoryBankFolder && memoryBankFolder.rules.some(rule => memoryBankFiles.includes(rule.name))
+
+    const buttonText = hasMemoryBankFiles ? 'Regenerate Memory Bank' : 'Generate Memory Bank'
+
+    return {
+        description: buttonText,
+        icon: MynahIcons.FOLDER,
+        id: ContextRule.CreateMemoryBankId,
+    }
+}
+
 export function convertRulesListToDetailedListGroup(rules: RulesFolder[]): DetailedListItemGroup[] {
     return rules
         .map(
@@ -179,7 +215,10 @@ export function convertRulesListToDetailedListGroup(rules: RulesFolder[]): Detai
                     })),
                 }) as DetailedListItemGroup
         )
-        .concat({ children: [createRuleListItem] })
+        .concat({
+            groupName: 'Actions',
+            children: [createMemoryBankListItem(rules), createRuleListItem],
+        })
 }
 
 function convertRuleStatusToIcon(status: boolean | 'indeterminate'): MynahIcons | undefined {

server/aws-lsp-codewhisperer/src/language-server/agenticChat/agenticChatController.ts

@@ -231,6 +231,7 @@ import { IDE } from '../../shared/constants'
 import { IdleWorkspaceManager } from '../workspaceContext/IdleWorkspaceManager'
 import escapeHTML = require('escape-html')
 import { SemanticSearch } from './tools/workspaceContext/semanticSearch'
+import { MemoryBankController } from './context/memorybank/memoryBankController'
 
 type ChatHandlers = Omit<
     LspHandlers<Chat>,
@@ -267,6 +268,7 @@ export class AgenticChatController implements ChatHandlers {
     #chatHistoryDb: ChatDatabase
     #additionalContextProvider: AdditionalContextProvider
     #contextCommandsProvider: ContextCommandsProvider
+    #memoryBankController: MemoryBankController
     #stoppedToolUses = new Set<string>()
     #userWrittenCodeTracker: UserWrittenCodeTracker | undefined
     #toolUseStartTimes: Record<string, number> = {}
@@ -370,6 +372,7 @@ export class AgenticChatController implements ChatHandlers {
         this.#mcpEventHandler = new McpEventHandler(features, telemetryService)
         this.#origin = getOriginFromClientInfo(getClientName(this.#features.lsp.getClientInitializeParams()))
         this.#activeUserTracker = ActiveUserTracker.getInstance(this.#features)
+        this.#memoryBankController = MemoryBankController.getInstance(features)
     }
 
     async onExecuteCommand(params: ExecuteCommandParams, _token: CancellationToken): Promise<any> {
@@ -833,6 +836,93 @@ export class AgenticChatController implements ChatHandlers {
 
         IdleWorkspaceManager.recordActivityTimestamp()
 
+        // Memory Bank Creation Flow - Delegate to MemoryBankController
+        if (this.#memoryBankController.isMemoryBankCreationRequest(params.prompt.prompt)) {
+            this.#features.logging.info(`Memory Bank creation request detected for tabId: ${params.tabId}`)
+
+            // Store original prompt to prevent data loss on failure
+            const originalPrompt = params.prompt.prompt
+
+            try {
+                const workspaceFolders = workspaceUtils.getWorkspaceFolderPaths(this.#features.workspace)
+                const workspaceUri = workspaceFolders.length > 0 ? workspaceFolders[0] : ''
+
+                if (!workspaceUri) {
+                    throw new Error('No workspace folder found for Memory Bank creation')
+                }
+
+                // Check if memory bank already exists to provide appropriate user feedback
+                const memoryBankExists = await this.#memoryBankController.memoryBankExists(workspaceUri)
+                const actionType = memoryBankExists ? 'Regenerating' : 'Generating'
+                this.#features.logging.info(`${actionType} Memory Bank for workspace: ${workspaceUri}`)
+
+                const resultStream = this.#getChatResultStream(params.partialResultToken)
+                await resultStream.writeResultBlock({
+                    body: `Preparing to analyze your project...`,
+                    type: 'answer',
+                    messageId: crypto.randomUUID(),
+                })
+
+                const comprehensivePrompt = await this.#memoryBankController.prepareComprehensiveMemoryBankPrompt(
+                    workspaceUri,
+                    async (prompt: string) => {
+                        // Direct LLM call for ranking - no agentic loop
+                        try {
+                            if (!this.#serviceManager) {
+                                throw new Error('amazonQServiceManager is not initialized')
+                            }
+
+                            const client = this.#serviceManager.getStreamingClient()
+                            const requestInput: SendMessageCommandInput = {
+                                conversationState: {
+                                    chatTriggerType: ChatTriggerType.MANUAL,
+                                    currentMessage: {
+                                        userInputMessage: {
+                                            content: prompt,
+                                        },
+                                    },
+                                },
+                            }
+
+                            const response = await client.sendMessage(requestInput)
+
+                            let responseContent = ''
+                            const maxResponseSize = 50000 // 50KB limit
+
+                            if (response.sendMessageResponse) {
+                                for await (const chatEvent of response.sendMessageResponse) {
+                                    if (chatEvent.assistantResponseEvent?.content) {
+                                        responseContent += chatEvent.assistantResponseEvent.content
+                                        if (responseContent.length > maxResponseSize) {
+                                            this.#features.logging.warn('LLM response exceeded size limit, truncating')
+                                            break
+                                        }
+                                    }
+                                }
+                            }
+
+                            return responseContent.trim()
+                        } catch (error) {
+                            this.#features.logging.error(`Memory Bank LLM ranking failed: ${error}`)
+                            return '' // Empty string triggers TF-IDF fallback
+                        }
+                    }
+                )
+
+                // Only update prompt if we got a valid comprehensive prompt
+                if (comprehensivePrompt && comprehensivePrompt.trim().length > 0) {
+                    params.prompt.prompt = comprehensivePrompt
+                } else {
+                    this.#features.logging.warn('Empty comprehensive prompt received, using original prompt')
+                    params.prompt.prompt = originalPrompt
+                }
+            } catch (error) {
+                this.#features.logging.error(`Memory Bank preparation failed: ${error}`)
+                // Restore original prompt to ensure no data loss
+                params.prompt.prompt = originalPrompt
+            }
+        }
+
         const maybeDefaultResponse = !params.prompt.command && getDefaultChatResponse(params.prompt.prompt)
         if (maybeDefaultResponse) {
             return maybeDefaultResponse
@@ -909,7 +999,8 @@ export class AgenticChatController implements ChatHandlers {
             const additionalContext = await this.#additionalContextProvider.getAdditionalContext(
                 triggerContext,
                 params.tabId,
-                params.context
+                params.context,
+                params.prompt.prompt
             )
             // Add active file to context list if it's not already there
             const activeFile =

server/aws-lsp-codewhisperer/src/language-server/agenticChat/context/additionalContextProvider.ts

@@ -33,7 +33,7 @@ import { ChatDatabase } from '../tools/chatDb/chatDb'
 import { ChatMessage, ImageBlock, ImageFormat } from '@amzn/codewhisperer-streaming'
 import { getRelativePathWithUri, getRelativePathWithWorkspaceFolder } from '../../workspaceContext/util'
 import { isSupportedImageExtension, MAX_IMAGE_CONTEXT_COUNT } from '../../../shared/imageVerification'
-import { mergeFileLists } from './contextUtils'
+import { MemoryBankController } from './memorybank/memoryBankController'
 
 export const ACTIVE_EDITOR_CONTEXT_ID = 'active-editor'
 
@@ -151,6 +151,30 @@ export class AdditionalContextProvider {
 
         // Filter rules based on user's rules preferences for current tab
         let rulesState = this.chatDb.getRules(tabId) || { folders: {}, rules: {} }
+
+        // Ensure memory bank files are active by default when first discovered
+        const memoryBankFiles = rulesFiles.filter(rule => rule.id?.includes('memory-bank'))
+        if (memoryBankFiles.length > 0) {
+            let needsUpdate = false
+
+            const memoryBankFolderName = 'memory-bank'
+            if (rulesState.folders[memoryBankFolderName] === undefined) {
+                rulesState.folders[memoryBankFolderName] = true
+                needsUpdate = true
+            }
+
+            memoryBankFiles.forEach(file => {
+                if (rulesState.rules[file.id] === undefined) {
+                    rulesState.rules[file.id] = true
+                    needsUpdate = true
+                }
+            })
+
+            if (needsUpdate) {
+                this.chatDb.setRules(tabId, rulesState)
+            }
+        }
+
         return rulesFiles.filter(rule => {
             // If the rule has an explicit state in rulesState, use that value
             if (rulesState.rules[rule.id] !== undefined) {
@@ -199,7 +223,8 @@ export class AdditionalContextProvider {
     async getAdditionalContext(
         triggerContext: TriggerContext,
         tabId: string,
-        context?: ContextCommand[]
+        context?: ContextCommand[],
+        prompt?: string
     ): Promise<AdditionalContentEntryAddition[]> {
         triggerContext.contextInfo = getInitialContextInfo()
 
@@ -220,7 +245,33 @@ export class AdditionalContextProvider {
             : workspaceUtils.getWorkspaceFolderPaths(this.features.workspace)[0]
 
         if (workspaceRules.length > 0) {
-            pinnedContextCommands.push(...workspaceRules)
+            // Check if this is a memory bank generation request
+            const isMemoryBankRequest = prompt
+                ? new MemoryBankController(this.features).isMemoryBankCreationRequest(prompt)
+                : false
+
+            let rulesToInclude = workspaceRules
+
+            if (isMemoryBankRequest) {
+                // Exclude memory bank files from context when regenerating memory bank
+                const memoryBankFiles = workspaceRules.filter(rule => rule.id?.includes('memory-bank'))
+                rulesToInclude = workspaceRules.filter(rule => !rule.id?.includes('memory-bank'))
+
+                if (memoryBankFiles.length > 0) {
+                    this.features.logging.info(
+                        `Memory Bank: excluding ${memoryBankFiles.length} existing memory bank files from context`
+                    )
+                }
+            } else {
+                // Normal behavior: include all workspace rules (including memory bank files)
+                const memoryBankFiles = workspaceRules.filter(rule => rule.id?.includes('memory-bank'))
+                if (memoryBankFiles.length > 0) {
+                    this.features.logging.info(`Including ${memoryBankFiles.length} memory bank files in chat context`)
+                }
+            }
+
+            // Add the filtered rules to pinned context
+            pinnedContextCommands.push(...rulesToInclude)
         }
 
         // Merge pinned context with context added to prompt, avoiding duplicates
@@ -675,7 +726,12 @@ export class AdditionalContextProvider {
                 if (dirPath === '.') {
                     folderName = undefined
                 } else {
-                    folderName = dirPath
+                    // Special handling for memory bank files
+                    if (dirPath === '.amazonq/rules/memory-bank') {
+                        folderName = 'memory-bank'
+                    } else {
+                        folderName = dirPath
+                    }
                 }
             } else {
                 // In multi-workspace: include workspace folder name for all files