fix: add in-loop compaction (#2387)

Author: tsmithsz

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

@@ -1185,15 +1185,21 @@ describe('AgenticChatController', () => {
 
         it('truncate input to 500k character ', async function () {
             const input = 'X'.repeat(GENERATE_ASSISTANT_RESPONSE_INPUT_LIMIT + 10)
-            generateAssistantResponseStub.restore()
-            generateAssistantResponseStub = sinon.stub(CodeWhispererStreaming.prototype, 'generateAssistantResponse')
-            generateAssistantResponseStub.callsFake(() => {})
-            await chatController.onChatPrompt({ tabId: mockTabId, prompt: { prompt: input } }, mockCancellationToken)
-            assert.ok(generateAssistantResponseStub.called)
-            const calledRequestInput: GenerateAssistantResponseCommandInput =
-                generateAssistantResponseStub.firstCall.firstArg
+            const request: GenerateAssistantResponseCommandInput = {
+                conversationState: {
+                    currentMessage: {
+                        userInputMessage: {
+                            content: input,
+                        },
+                    },
+                    chatTriggerType: undefined,
+                },
+            }
+
+            chatController.truncateRequest(request)
+
             assert.deepStrictEqual(
-                calledRequestInput.conversationState?.currentMessage?.userInputMessage?.content?.length,
+                request.conversationState?.currentMessage?.userInputMessage?.content?.length,
                 GENERATE_ASSISTANT_RESPONSE_INPUT_LIMIT
             )
         })

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

@@ -191,6 +191,8 @@ import {
     MAX_OVERALL_CHARACTERS,
     FSREAD_MEMORY_BANK_MAX_PER_FILE,
     FSREAD_MEMORY_BANK_MAX_TOTAL,
+    MID_LOOP_COMPACTION_HANDOFF_PROMPT,
+    COMPACTION_PROMPT,
 } from './constants/constants'
 import {
     AgenticChatError,
@@ -943,6 +945,11 @@ export class AgenticChatController implements ChatHandlers {
 
         const compactIds = session.getAllDeferredCompactMessageIds()
         await this.#invalidateCompactCommand(params.tabId, compactIds)
+        // Set compactionDeclined flag if there were pending compaction requests
+        // This prevents endless compaction warning loops when user declines compaction once
+        if (compactIds.length > 0) {
+            session.compactionDeclined = true
+        }
         session.rejectAllDeferredToolExecutions(new ToolApprovalException('Command ignored: new prompt', false))
         await this.#invalidateAllShellCommands(params.tabId, session)
 
@@ -978,6 +985,10 @@ export class AgenticChatController implements ChatHandlers {
                 session.abortRequest()
                 const compactIds = session.getAllDeferredCompactMessageIds()
                 await this.#invalidateCompactCommand(params.tabId, compactIds)
+                // Set compactionDeclined flag if there were pending compaction requests
+                if (compactIds.length > 0) {
+                    session.compactionDeclined = true
+                }
                 void this.#invalidateAllShellCommands(params.tabId, session)
                 session.rejectAllDeferredToolExecutions(new CancellationError('user'))
 
@@ -1087,7 +1098,7 @@ export class AgenticChatController implements ChatHandlers {
             }
 
             // Result Handling - This happens only once
-            return await this.#handleFinalResult(
+            const result = await this.#handleFinalResult(
                 finalResult,
                 session,
                 params.tabId,
@@ -1096,6 +1107,13 @@ export class AgenticChatController implements ChatHandlers {
                 isNewConversation,
                 chatResultStream
             )
+
+            // Reset compactionDeclined flag after successful completion
+            if (session.compactionDeclined) {
+                session.compactionDeclined = false
+            }
+
+            return result
         } catch (err) {
             // HACK: the chat-client needs to have a partial event with the associated messageId sent before it can accept the final result.
             // Without this, the `working` indicator never goes away.
@@ -1167,25 +1185,29 @@ export class AgenticChatController implements ChatHandlers {
     /**
      * Prepares the initial request input for the chat prompt
      */
-    #getCompactionRequestInput(session: ChatSessionService): ChatCommandInput {
+    #getCompactionRequestInput(session: ChatSessionService, toolResults?: any[]): ChatCommandInput {
         this.#debug('Preparing compaction request input')
         // Get profileArn from the service manager if available
         const profileArn = this.#serviceManager?.getActiveProfileArn()
         const requestInput = this.#triggerContext.getCompactionChatCommandInput(
             profileArn,
             this.#getTools(session),
             session.modelId,
-            this.#origin
+            this.#origin,
+            toolResults
         )
         return requestInput
     }
 
     /**
      * Runs the compaction, making requests and processing tool uses until completion
      */
-    #shouldCompact(currentRequestCount: number): boolean {
-        if (currentRequestCount > COMPACTION_CHARACTER_THRESHOLD) {
-            this.#debug(`Current request total character count is: ${currentRequestCount}, prompting user to compact`)
+    #shouldCompact(currentRequestCount: number, session: ChatSessionService): boolean {
+        const EFFECTIVE_COMPACTION_THRESHOLD = COMPACTION_CHARACTER_THRESHOLD - COMPACTION_PROMPT.length
+        if (currentRequestCount > EFFECTIVE_COMPACTION_THRESHOLD && !session.compactionDeclined) {
+            this.#debug(
+                `Current request total character count is: ${currentRequestCount}, prompting user to compact (threshold: ${EFFECTIVE_COMPACTION_THRESHOLD})`
+            )
             return true
         } else {
             return false
@@ -1368,6 +1390,10 @@ export class AgenticChatController implements ChatHandlers {
         let currentRequestCount = 0
         const pinnedContext = additionalContext?.filter(item => item.pinned)
 
+        // Store initial non-empty prompt for compaction handoff
+        const initialPrompt =
+            initialRequestInput.conversationState?.currentMessage?.userInputMessage?.content?.trim() || ''
+
         metric.recordStart()
         this.logSystemInformation()
         while (true) {
@@ -1432,6 +1458,63 @@ export class AgenticChatController implements ChatHandlers {
             this.#llmRequestStartTime = Date.now()
             // Phase 3: Request Execution
             currentRequestInput = sanitizeRequestInput(currentRequestInput)
+
+            if (this.#shouldCompact(currentRequestCount, session)) {
+                this.#features.logging.info(
+                    `Entering mid-loop compaction at iteration ${iterationCount} with ${currentRequestCount} characters`
+                )
+                this.#telemetryController.emitMidLoopCompaction(
+                    currentRequestCount,
+                    iterationCount,
+                    this.#features.runtime.serverInfo.version ?? ''
+                )
+                const messageId = this.#getMessageIdForCompact(uuid())
+                const confirmationResult = this.#processCompactConfirmation(messageId, currentRequestCount)
+                const cachedButtonBlockId = await chatResultStream.writeResultBlock(confirmationResult)
+                await this.waitForCompactApproval(messageId, chatResultStream, cachedButtonBlockId, session)
+
+                // Run compaction
+                const toolResults =
+                    currentRequestInput.conversationState?.currentMessage?.userInputMessage?.userInputMessageContext
+                        ?.toolResults || []
+                const compactionRequestInput = this.#getCompactionRequestInput(session, toolResults)
+                const compactionResult = await this.#runCompaction(
+                    compactionRequestInput,
+                    session,
+                    metric,
+                    chatResultStream,
+                    tabId,
+                    promptId,
+                    CompactHistoryActionType.Nudge,
+                    session.conversationId,
+                    token,
+                    documentReference
+                )
+
+                if (!compactionResult.success) {
+                    this.#features.logging.error(`Compaction failed: ${compactionResult.error}`)
+                    return compactionResult
+                }
+
+                // Show compaction summary to user before continuing
+                await chatResultStream.writeResultBlock({
+                    type: 'answer',
+                    body:
+                        (compactionResult.data?.chatResult.body || '') +
+                        '\n\nConversation history has been compacted successfully!',
+                    messageId: uuid(),
+                })
+
+                currentRequestInput = this.#updateRequestInputWithToolResults(
+                    currentRequestInput,
+                    [],
+                    MID_LOOP_COMPACTION_HANDOFF_PROMPT + initialPrompt
+                )
+                shouldDisplayMessage = false
+                this.#features.logging.info(`Completed mid-loop compaction, restarting loop with handoff prompt`)
+                continue
+            }
+
             // Note: these logs are very noisy, but contain information redacted on the backend.
             this.#debug(
                 `generateAssistantResponse/SendMessage Request: ${JSON.stringify(currentRequestInput, this.#imageReplacer, 2)}`
@@ -1660,7 +1743,7 @@ export class AgenticChatController implements ChatHandlers {
             currentRequestInput = this.#updateRequestInputWithToolResults(currentRequestInput, toolResults, content)
         }
 
-        if (this.#shouldCompact(currentRequestCount)) {
+        if (this.#shouldCompact(currentRequestCount, session)) {
             this.#telemetryController.emitCompactNudge(
                 currentRequestCount,
                 this.#features.runtime.serverInfo.version ?? ''

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

@@ -18,6 +18,8 @@ export const GENERATE_ASSISTANT_RESPONSE_INPUT_LIMIT = 500_000
 // 200K tokens * 3.5 = 700K characters, intentionally overestimating with 3.5:1 ratio
 export const MAX_OVERALL_CHARACTERS = 700_000
 export const COMPACTION_CHARACTER_THRESHOLD = 0.7 * MAX_OVERALL_CHARACTERS
+// TODO: We need to carefully craft a prompt for this supported by the science team
+export const MID_LOOP_COMPACTION_HANDOFF_PROMPT = `CONTEXT HANDOFF: Previous conversation was compacted. Continue with user's request: `
 export const COMPACTION_BODY = (threshold: number) =>
     `The context window is almost full (${threshold}%) and exceeding it will clear your history. Amazon Q can compact your history instead.`
 export const COMPACTION_HEADER_BODY = 'Compact chat history?'

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

@@ -113,13 +113,15 @@ export class AgenticChatTriggerContext {
      * @param tools Optional Bedrock tools
      * @param modelId Optional model ID
      * @param origin Optional origin
+     * @param toolResults Optional tool results to include in compaction context
      * @returns ChatCommandInput - which is either SendMessageInput or GenerateAssistantResponseInput
      */
     getCompactionChatCommandInput(
         profileArn?: string,
         tools: BedrockTools = [],
         modelId?: string,
-        origin?: Origin
+        origin?: Origin,
+        toolResults?: any[]
     ): ChatCommandInput {
         const data: ChatCommandInput = {
             conversationState: {
@@ -130,6 +132,7 @@ export class AgenticChatTriggerContext {
                         userInputMessageContext: {
                             tools,
                             envState: this.#mapPlatformToEnvState(process.platform),
+                            toolResults: toolResults,
                         },
                         userIntent: undefined,
                         origin: origin ? origin : 'IDE',

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

@@ -42,6 +42,7 @@ export class ChatSessionService {
     public contextListSent: boolean = false
     public modelId: string | undefined
     public isMemoryBankGeneration: boolean = false
+    public compactionDeclined: boolean = false
     #lsp?: Features['lsp']
     #abortController?: AbortController
     #currentPromptId?: string

server/aws-lsp-codewhisperer/src/language-server/chat/telemetry/chatTelemetryController.ts

@@ -243,6 +243,18 @@ export class ChatTelemetryController {
         })
     }
 
+    public emitMidLoopCompaction(characters: number, iterationCount: number, languageServerVersion: string) {
+        this.#telemetry.emitMetric({
+            name: ChatTelemetryEventName.MidLoopCompaction,
+            data: {
+                characters,
+                iterationCount,
+                credentialStartUrl: this.#credentialsProvider.getConnectionMetadata()?.sso?.startUrl,
+                languageServerVersion: languageServerVersion,
+            },
+        })
+    }
+
     public emitToolUseSuggested(
         toolUse: ToolUse,
         conversationId: string,

server/aws-lsp-codewhisperer/src/shared/telemetry/types.ts

@@ -207,6 +207,7 @@ export enum ChatTelemetryEventName {
     LoadHistory = 'amazonq_loadHistory',
     CompactHistory = 'amazonq_compactHistory',
     CompactNudge = 'amazonq_compactNudge',
+    MidLoopCompaction = 'amazonq_midLoopCompaction',
     ChatHistoryAction = 'amazonq_performChatHistoryAction',
     ExportTab = 'amazonq_exportTab',
     UiClick = 'ui_click',
@@ -233,6 +234,7 @@ export interface ChatTelemetryEventMap {
     [ChatTelemetryEventName.LoadHistory]: LoadHistoryEvent
     [ChatTelemetryEventName.CompactHistory]: CompactHistoryEvent
     [ChatTelemetryEventName.CompactNudge]: CompactNudgeEvent
+    [ChatTelemetryEventName.MidLoopCompaction]: MidLoopCompactionEvent
     [ChatTelemetryEventName.ChatHistoryAction]: ChatHistoryActionEvent
     [ChatTelemetryEventName.ExportTab]: ExportTabEvent
     [ChatTelemetryEventName.UiClick]: UiClickEvent
@@ -404,6 +406,13 @@ export type CompactNudgeEvent = {
     languageServerVersion?: string
 }
 
+export type MidLoopCompactionEvent = {
+    characters: number
+    iterationCount: number
+    credentialStartUrl?: string
+    languageServerVersion?: string
+}
+
 export type ChatHistoryActionEvent = {
     action: ChatHistoryActionType
     result: Result