Skip to content

Commit 3057d56

Browse files
tsmithsztsmithsz
authored
fix: optimize memory bank token usage and add new tab support (#2366)
1 parent 1f6b7f7 commit 3057d56

File tree

7 files changed

+113
-29
lines changed

7 files changed

+113
-29
lines changed

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

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -16,13 +16,17 @@ describe('rules', () => {
1616
mynahUi = {
1717
openTopBarButtonOverlay: sinon.stub(),
1818
showCustomForm: sinon.stub(),
19+
getAllTabs: sinon.stub().returns({}),
20+
updateStore: sinon.stub().returns('new-tab-id'),
21+
notify: sinon.stub(),
1922
} as unknown as MynahUI
2023
openTopBarButtonOverlayStub = mynahUi.openTopBarButtonOverlay as sinon.SinonStub
2124
showCustomFormStub = mynahUi.showCustomForm as sinon.SinonStub
2225

2326
messager = {
2427
onRuleClick: sinon.stub(),
2528
onChatPrompt: sinon.stub(),
29+
onTabAdd: sinon.stub(),
2630
} as unknown as Messager
2731

2832
rulesList = new RulesList(mynahUi, messager)
@@ -151,12 +155,17 @@ describe('rules', () => {
151155

152156
onItemClick(createMemoryBankItem)
153157

154-
// Should send a chat prompt
158+
// Should create new tab and send chat prompt
159+
sinon.assert.calledOnce(messager.onTabAdd as sinon.SinonStub)
155160
sinon.assert.calledOnce(messager.onChatPrompt as sinon.SinonStub)
156161

162+
const tabAddArgs = (messager.onTabAdd as sinon.SinonStub).getCall(0).args[0]
163+
assert.equal(tabAddArgs, 'new-tab-id')
164+
157165
const chatPromptArgs = (messager.onChatPrompt as sinon.SinonStub).getCall(0).args[0]
158166
assert.equal(chatPromptArgs.prompt.prompt, 'Generate a Memory Bank for this project')
159167
assert.equal(chatPromptArgs.prompt.escapedPrompt, 'Generate a Memory Bank for this project')
168+
assert.equal(chatPromptArgs.tabId, 'new-tab-id')
160169
})
161170

162171
it('calls messager when regular rule is clicked', () => {

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

Lines changed: 40 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,11 @@
1-
import { MynahIconsType, MynahUI, DetailedListItem, DetailedListItemGroup, MynahIcons } from '@aws/mynah-ui'
1+
import {
2+
MynahIconsType,
3+
MynahUI,
4+
DetailedListItem,
5+
DetailedListItemGroup,
6+
MynahIcons,
7+
NotificationType,
8+
} from '@aws/mynah-ui'
29
import { Messager } from '../messager'
310
import { ListRulesResult } from '@aws/language-server-runtimes-types'
411
import { RulesFolder } from '@aws/language-server-runtimes-types'
@@ -82,14 +89,38 @@ export class RulesList {
8289
// Close the rules list first
8390
this.rulesList?.close()
8491

85-
// Use the current tab, the tabId should be the same as the one used for the rules list
86-
this.messager.onChatPrompt({
87-
prompt: {
88-
prompt: 'Generate a Memory Bank for this project',
89-
escapedPrompt: 'Generate a Memory Bank for this project',
90-
},
91-
tabId: this.tabId,
92-
})
92+
// Check if we're at the tab limit (10 tabs max)
93+
const currentTabCount = Object.keys(this.mynahUi.getAllTabs()).length
94+
if (currentTabCount >= 10) {
95+
// Show notification that max tabs reached
96+
this.mynahUi.notify({
97+
content: 'You can only open ten conversation tabs at a time.',
98+
type: NotificationType.WARNING,
99+
})
100+
return
101+
}
102+
103+
// Create a new tab for the memory bank generation
104+
const newTabId = this.mynahUi.updateStore('', { tabTitle: 'Memory Bank' })
105+
if (newTabId) {
106+
// Add the new tab and switch to it
107+
this.messager.onTabAdd(newTabId)
108+
109+
// Send the chat prompt to the new tab
110+
this.messager.onChatPrompt({
111+
prompt: {
112+
prompt: 'Generate a Memory Bank for this project',
113+
escapedPrompt: 'Generate a Memory Bank for this project',
114+
},
115+
tabId: newTabId,
116+
})
117+
} else {
118+
// Show error notification if tab creation failed
119+
this.mynahUi.notify({
120+
content: 'Failed to create new tab for Memory Bank generation.',
121+
type: NotificationType.ERROR,
122+
})
123+
}
93124
}
94125

95126
showLoading(tabId: string) {

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

Lines changed: 29 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -189,6 +189,8 @@ import {
189189
DEFAULT_WINDOW_STOP_SHORTCUT,
190190
COMPACTION_CHARACTER_THRESHOLD,
191191
MAX_OVERALL_CHARACTERS,
192+
FSREAD_MEMORY_BANK_MAX_PER_FILE,
193+
FSREAD_MEMORY_BANK_MAX_TOTAL,
192194
} from './constants/constants'
193195
import {
194196
AgenticChatError,
@@ -837,9 +839,17 @@ export class AgenticChatController implements ChatHandlers {
837839

838840
IdleWorkspaceManager.recordActivityTimestamp()
839841

842+
const sessionResult = this.#chatSessionManagementService.getSession(params.tabId)
843+
const { data: session, success } = sessionResult
844+
845+
if (!success) {
846+
return new ResponseError<ChatResult>(ErrorCodes.InternalError, sessionResult.error)
847+
}
848+
840849
// Memory Bank Creation Flow - Delegate to MemoryBankController
841850
if (this.#memoryBankController.isMemoryBankCreationRequest(params.prompt.prompt)) {
842851
this.#features.logging.info(`Memory Bank creation request detected for tabId: ${params.tabId}`)
852+
session.isMemoryBankGeneration = true
843853

844854
// Store original prompt to prevent data loss on failure
845855
const originalPrompt = params.prompt.prompt
@@ -921,6 +931,8 @@ export class AgenticChatController implements ChatHandlers {
921931
this.#features.logging.error(`Memory Bank preparation failed: ${error}`)
922932
// Restore original prompt to ensure no data loss
923933
params.prompt.prompt = originalPrompt
934+
// Reset memory bank flag since preparation failed
935+
session.isMemoryBankGeneration = false
924936
}
925937
}
926938

@@ -929,14 +941,6 @@ export class AgenticChatController implements ChatHandlers {
929941
return maybeDefaultResponse
930942
}
931943

932-
const sessionResult = this.#chatSessionManagementService.getSession(params.tabId)
933-
934-
const { data: session, success } = sessionResult
935-
936-
if (!success) {
937-
return new ResponseError<ChatResult>(ErrorCodes.InternalError, sessionResult.error)
938-
}
939-
940944
const compactIds = session.getAllDeferredCompactMessageIds()
941945
await this.#invalidateCompactCommand(params.tabId, compactIds)
942946
session.rejectAllDeferredToolExecutions(new ToolApprovalException('Command ignored: new prompt', false))
@@ -1930,7 +1934,14 @@ export class AgenticChatController implements ChatHandlers {
19301934
}
19311935

19321936
const { Tool } = toolMap[toolUse.name as keyof typeof toolMap]
1933-
const tool = new Tool(this.#features)
1937+
const tool =
1938+
toolUse.name === FS_READ && session.isMemoryBankGeneration
1939+
? new Tool(
1940+
this.#features,
1941+
FSREAD_MEMORY_BANK_MAX_PER_FILE,
1942+
FSREAD_MEMORY_BANK_MAX_TOTAL
1943+
)
1944+
: new Tool(this.#features)
19341945

19351946
// For MCP tools, get the permission from McpManager
19361947
// const permission = McpManager.instance.getToolPerm('Built-in', toolUse.name)
@@ -3458,6 +3469,9 @@ export class AgenticChatController implements ChatHandlers {
34583469
},
34593470
})
34603471

3472+
// Reset memory bank flag after completion
3473+
session.isMemoryBankGeneration = false
3474+
34613475
return chatResultStream.getResult()
34623476
}
34633477

@@ -3485,6 +3499,12 @@ export class AgenticChatController implements ChatHandlers {
34853499
const errorCode = err.code ?? ''
34863500
await this.#telemetryController.emitAddMessageMetric(tabId, metric.metric, 'Failed', errorMessage, errorCode)
34873501

3502+
// Reset memory bank flag on request error
3503+
const sessionResult = this.#chatSessionManagementService.getSession(tabId)
3504+
if (sessionResult.success) {
3505+
sessionResult.data.isMemoryBankGeneration = false
3506+
}
3507+
34883508
if (isUsageLimitError(err)) {
34893509
if (this.#paidTierMode !== 'paidtier') {
34903510
this.setPaidTierMode(tabId, 'freetier-limit')

‎server/aws-lsp-codewhisperer/src/language-server/agenticChat/constants/constants.ts‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -81,6 +81,16 @@ The summary should have following main sections:
8181
</example_output>
8282
`
8383

84+
// FsRead limits
85+
export const FSREAD_MAX_PER_FILE = 200_000
86+
export const FSREAD_MAX_TOTAL = 400_000
87+
export const FSREAD_MEMORY_BANK_MAX_PER_FILE = 20_000
88+
export const FSREAD_MEMORY_BANK_MAX_TOTAL = 100_000
89+
90+
// Memory Bank constants
91+
// Temporarily reduced from recommended 20 to 5 for token optimization
92+
export const MAX_NUMBER_OF_FILES_FOR_MEMORY_BANK_RANKING = 5
93+
8494
// shortcut constant
8595
export const DEFAULT_MACOS_RUN_SHORTCUT = '&#8679; &#8984; &#8629;'
8696
export const DEFAULT_WINDOW_RUN_SHORTCUT = 'Ctrl + &#8679; + &#8629;'

‎server/aws-lsp-codewhisperer/src/language-server/agenticChat/context/memorybank/memoryBankController.ts‎

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@
66
import { Features } from '@aws/language-server-runtimes/server-interface/server'
77
import { MemoryBankPrompts } from './memoryBankPrompts'
88
import { normalizePathFromUri } from '../../tools/mcp/mcpUtils'
9+
import { MAX_NUMBER_OF_FILES_FOR_MEMORY_BANK_RANKING } from '../../constants/constants'
910

1011
const MEMORY_BANK_DIRECTORY = '.amazonq/rules/memory-bank'
1112
const MEMORY_BANK_FILES = {
@@ -71,7 +72,10 @@ export class MemoryBankController {
7172
const analysisResults = await this.executeGuidelinesGenerationPipeline(workspaceFolderUri)
7273

7374
// Step 3: Make LLM call for file ranking
74-
const rankingPrompt = MemoryBankPrompts.getFileRankingPrompt(analysisResults.formattedFilesString, 10)
75+
const rankingPrompt = MemoryBankPrompts.getFileRankingPrompt(
76+
analysisResults.formattedFilesString,
77+
MAX_NUMBER_OF_FILES_FOR_MEMORY_BANK_RANKING
78+
)
7579
const rankedFilesResponse = await llmCallFunction(rankingPrompt)
7680

7781
// Step 4: Parse ranked files
@@ -111,7 +115,7 @@ export class MemoryBankController {
111115
this.features.logging.warn(
112116
`Memory Bank: failed to parse LLM ranking response, using TF-IDF fallback: ${error}`
113117
)
114-
rankedFilesList = analysisResults.rankedFilesList.slice(0, 10)
118+
rankedFilesList = analysisResults.rankedFilesList.slice(0, MAX_NUMBER_OF_FILES_FOR_MEMORY_BANK_RANKING)
115119
}
116120

117121
this.features.logging.info(
@@ -477,7 +481,7 @@ export class MemoryBankController {
477481
// Step 5: Create fallback ranking (deterministic, for when LLM fails)
478482
const rankedFilesList = filesWithDissimilarity
479483
.sort((a, b) => b.dissimilarity - a.dissimilarity)
480-
.slice(0, 10)
484+
.slice(0, MAX_NUMBER_OF_FILES_FOR_MEMORY_BANK_RANKING)
481485
.map(f => f.path)
482486

483487
return {

‎server/aws-lsp-codewhisperer/src/language-server/agenticChat/tools/fsRead.ts‎

Lines changed: 16 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
import { sanitize } from '@aws/lsp-core/out/util/path'
22
import { CommandValidation, InvokeOutput, requiresPathAcceptance, validatePath } from './toolShared'
33
import { Features } from '@aws/language-server-runtimes/server-interface/server'
4+
import { FSREAD_MAX_PER_FILE, FSREAD_MAX_TOTAL } from '../constants/constants'
45

56
export interface FsReadParams {
67
paths: string[]
@@ -13,16 +14,24 @@ export interface FileReadResult {
1314
}
1415

1516
export class FsRead {
16-
static maxResponseSize = 200_000
17-
static maxResponseSizeTotal = 400_000
17+
static maxResponseSize = FSREAD_MAX_PER_FILE
18+
static maxResponseSizeTotal = FSREAD_MAX_TOTAL
1819
private readonly logging: Features['logging']
1920
private readonly workspace: Features['workspace']
2021
private readonly lsp: Features['lsp']
22+
private readonly maxPerFile: number
23+
private readonly maxTotal: number
2124

22-
constructor(features: Pick<Features, 'lsp' | 'workspace' | 'logging'> & Partial<Features>) {
25+
constructor(
26+
features: Pick<Features, 'lsp' | 'workspace' | 'logging'> & Partial<Features>,
27+
maxPerFile?: number,
28+
maxTotal?: number
29+
) {
2330
this.logging = features.logging
2431
this.workspace = features.workspace
2532
this.lsp = features.lsp
33+
this.maxPerFile = maxPerFile ?? FsRead.maxResponseSize
34+
this.maxTotal = maxTotal ?? FsRead.maxResponseSizeTotal
2635
}
2736

2837
public async validate(params: FsReadParams): Promise<void> {
@@ -62,16 +71,16 @@ export class FsRead {
6271
private createOutput(fileResult: FileReadResult[]): InvokeOutput {
6372
let totalSize = 0
6473
for (const result of fileResult) {
65-
const exceedsMaxSize = result.content.length > FsRead.maxResponseSize
74+
const exceedsMaxSize = result.content.length > this.maxPerFile
6675
if (exceedsMaxSize) {
67-
this.logging.info(`FsRead: truncating ${result.path} to first ${FsRead.maxResponseSize} characters`)
68-
result.content = result.content.substring(0, FsRead.maxResponseSize - 3) + '...'
76+
this.logging.info(`FsRead: truncating ${result.path} to first ${this.maxPerFile} characters`)
77+
result.content = result.content.substring(0, this.maxPerFile - 3) + '...'
6978
result.truncated = true
7079
}
7180
totalSize += result.content.length
7281
}
7382

74-
if (totalSize > FsRead.maxResponseSizeTotal) {
83+
if (totalSize > this.maxTotal) {
7584
throw Error('Files are too large, please break the file read into smaller chunks')
7685
}
7786

‎server/aws-lsp-codewhisperer/src/language-server/chat/chatSessionService.ts‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -41,6 +41,7 @@ export class ChatSessionService {
4141
public pairProgrammingMode: boolean = true
4242
public contextListSent: boolean = false
4343
public modelId: string | undefined
44+
public isMemoryBankGeneration: boolean = false
4445
#lsp?: Features['lsp']
4546
#abortController?: AbortController
4647
#currentPromptId?: string

0 commit comments

Comments
 (0)