Merge branch 'feature/stepfunctions-execution' into feature/stepfunctions-execution

Author: laileni-aws

packages/amazonq/src/app/inline/EditRendering/displayImage.ts

@@ -17,6 +17,7 @@ import type { AmazonQInlineCompletionItemProvider } from '../completion'
 import { vsCodeState } from 'aws-core-vscode/codewhisperer'
 
 const autoRejectEditCursorDistance = 25
+const autoDiscardEditCursorDistance = 10
 
 export class EditDecorationManager {
     private imageDecorationType: vscode.TextEditorDecorationType
@@ -312,6 +313,18 @@ export async function displaySvgDecoration(
     item: InlineCompletionItemWithReferences,
     inlineCompletionProvider?: AmazonQInlineCompletionItemProvider
 ) {
+    // Check if edit is too far from current cursor position
+    const currentCursorLine = editor.selection.active.line
+    if (Math.abs(startLine - currentCursorLine) >= autoDiscardEditCursorDistance) {
+        // Emit DISCARD telemetry for edit suggestion that can't be shown because the suggestion is too far away
+        const params = createDiscardTelemetryParams(session, item)
+        languageClient.sendNotification('aws/logInlineCompletionSessionResults', params)
+        getLogger('nextEditPrediction').debug(
+            `Auto discarded edit suggestion for suggestion that is too far away: ${item.insertText as string}`
+        )
+        return
+    }
+
     const originalCode = editor.document.getText()
 
     // Set edit state immediately to prevent race condition with completion requests
@@ -399,9 +412,6 @@ export async function displaySvgDecoration(
             const endPosition = getEndOfEditPosition(originalCode, newCode)
             editor.selection = new vscode.Selection(endPosition, endPosition)
 
-            // Move cursor to end of the actual changed content
-            editor.selection = new vscode.Selection(endPosition, endPosition)
-
             await decorationManager.clearDecorations(editor)
             documentChangeListener.dispose()
             cursorChangeListener.dispose()
@@ -420,19 +430,6 @@ export async function displaySvgDecoration(
             }
             languageClient.sendNotification('aws/logInlineCompletionSessionResults', params)
             session.triggerOnAcceptance = true
-            // VS Code triggers suggestion on every keystroke, temporarily disable trigger on acceptance
-            // if (inlineCompletionProvider && session.editsStreakPartialResultToken) {
-            //     await inlineCompletionProvider.provideInlineCompletionItems(
-            //         editor.document,
-            //         endPosition,
-            //         {
-            //             triggerKind: vscode.InlineCompletionTriggerKind.Automatic,
-            //             selectedCompletionInfo: undefined,
-            //         },
-            //         new vscode.CancellationTokenSource().token,
-            //         { emitTelemetry: false, showUi: false, editsStreakToken: session.editsStreakPartialResultToken }
-            //     )
-            // }
         },
         async (isDiscard: boolean) => {
             // Handle reject

packages/amazonq/src/app/inline/completion.ts

@@ -213,6 +213,8 @@ export class InlineCompletionManager implements Disposable {
 
 export class AmazonQInlineCompletionItemProvider implements InlineCompletionItemProvider {
     private logger = getLogger()
+    private pendingRequest: Promise<InlineCompletionItem[]> | undefined
+
     constructor(
         private readonly languageClient: LanguageClient,
         private readonly recommendationService: RecommendationService,
@@ -300,6 +302,48 @@ export class AmazonQInlineCompletionItemProvider implements InlineCompletionItem
             options: JSON.stringify(getAllRecommendationsOptions),
         })
 
+        // If there's already a pending request, wait for it to complete instead of starting a new one
+        // This prevents race conditions where multiple concurrent calls cause the later (empty) response
+        // to override the earlier (valid) response
+        if (this.pendingRequest) {
+            getLogger().info('Reusing pending inline completion request to avoid race condition')
+            try {
+                const result = await this.pendingRequest
+                // Check if THIS call's token was cancelled (not the original call's token)
+                if (token.isCancellationRequested) {
+                    getLogger().info('Reused request completed but this call was cancelled')
+                    return []
+                }
+                return result
+            } catch (e) {
+                // If the pending request failed, continue with a new request
+                getLogger().info('Pending request failed, starting new request: %O', e)
+            }
+        }
+
+        // Start a new request and track it
+        this.pendingRequest = this._provideInlineCompletionItemsImpl(
+            document,
+            position,
+            context,
+            token,
+            getAllRecommendationsOptions
+        )
+
+        try {
+            return await this.pendingRequest
+        } finally {
+            this.pendingRequest = undefined
+        }
+    }
+
+    private async _provideInlineCompletionItemsImpl(
+        document: TextDocument,
+        position: Position,
+        context: InlineCompletionContext,
+        token: CancellationToken,
+        getAllRecommendationsOptions?: GetAllRecommendationsOptions
+    ): Promise<InlineCompletionItem[]> {
         if (vsCodeState.isCodeWhispererEditing) {
             getLogger().info('Q is editing, returning empty')
             return []

packages/amazonq/test/unit/app/inline/EditRendering/displayImage.test.ts

@@ -188,6 +188,89 @@ describe('EditDecorationManager', function () {
     })
 })
 
+describe('displaySvgDecoration cursor distance auto-discard', function () {
+    let sandbox: sinon.SinonSandbox
+    let editorStub: sinon.SinonStubbedInstance<vscode.TextEditor>
+    let languageClientStub: any
+    let sessionStub: any
+    let itemStub: any
+
+    beforeEach(function () {
+        sandbox = sinon.createSandbox()
+        const commonStubs = createCommonStubs(sandbox)
+        editorStub = commonStubs.editorStub
+
+        languageClientStub = {
+            sendNotification: sandbox.stub(),
+        }
+
+        sessionStub = {
+            sessionId: 'test-session',
+            requestStartTime: Date.now(),
+            firstCompletionDisplayLatency: 100,
+        }
+
+        itemStub = {
+            itemId: 'test-item',
+            insertText: 'test content',
+        }
+    })
+
+    afterEach(function () {
+        sandbox.restore()
+    })
+
+    it('should send discard telemetry and return early when edit is 10+ lines away from cursor', async function () {
+        // Set cursor at line 5
+        editorStub.selection = {
+            active: new vscode.Position(5, 0),
+        } as any
+        // Try to display edit at line 20 (15 lines away)
+        await displaySvgDecoration(
+            editorStub as unknown as vscode.TextEditor,
+            vscode.Uri.parse('data:image/svg+xml;base64,test'),
+            20,
+            'new code',
+            [],
+            sessionStub,
+            languageClientStub,
+            itemStub
+        )
+
+        // Verify discard telemetry was sent
+        sinon.assert.calledOnce(languageClientStub.sendNotification)
+        const call = languageClientStub.sendNotification.getCall(0)
+        assert.strictEqual(call.args[0], 'aws/logInlineCompletionSessionResults')
+        assert.strictEqual(call.args[1].sessionId, 'test-session')
+        assert.strictEqual(call.args[1].completionSessionResult['test-item'].discarded, true)
+    })
+
+    it('should proceed normally when edit is within 10 lines of cursor', async function () {
+        // Set cursor at line 5
+        editorStub.selection = {
+            active: new vscode.Position(5, 0),
+        } as any
+        // Mock required dependencies for normal flow
+        sandbox.stub(vscode.workspace, 'onDidChangeTextDocument').returns({ dispose: sandbox.stub() })
+        sandbox.stub(vscode.window, 'onDidChangeTextEditorSelection').returns({ dispose: sandbox.stub() })
+
+        // Try to display edit at line 10 (5 lines away)
+        await displaySvgDecoration(
+            editorStub as unknown as vscode.TextEditor,
+            vscode.Uri.parse('data:image/svg+xml;base64,test'),
+            10,
+            'new code',
+            [],
+            sessionStub,
+            languageClientStub,
+            itemStub
+        )
+
+        // Verify no discard telemetry was sent (function should proceed normally)
+        sinon.assert.notCalled(languageClientStub.sendNotification)
+    })
+})
+
 describe('displaySvgDecoration cursor distance auto-reject', function () {
     let sandbox: sinon.SinonSandbox
     let editorStub: sinon.SinonStubbedInstance<vscode.TextEditor>
@@ -253,6 +336,10 @@ describe('displaySvgDecoration cursor distance auto-reject', function () {
     })
 
     it('should not reject when cursor moves less than 25 lines away', async function () {
+        // Set cursor at line 50
+        editorStub.selection = {
+            active: new vscode.Position(50, 0),
+        } as any
         const startLine = 50
         await setupDisplaySvgDecoration(startLine)
 
@@ -262,6 +349,10 @@ describe('displaySvgDecoration cursor distance auto-reject', function () {
     })
 
     it('should not reject when cursor moves exactly 25 lines away', async function () {
+        // Set cursor at line 50
+        editorStub.selection = {
+            active: new vscode.Position(50, 0),
+        } as any
         const startLine = 50
         await setupDisplaySvgDecoration(startLine)
 
@@ -271,6 +362,10 @@ describe('displaySvgDecoration cursor distance auto-reject', function () {
     })
 
     it('should reject when cursor moves more than 25 lines away', async function () {
+        // Set cursor at line 50
+        editorStub.selection = {
+            active: new vscode.Position(50, 0),
+        } as any
         const startLine = 50
         await setupDisplaySvgDecoration(startLine)
 
@@ -280,6 +375,10 @@ describe('displaySvgDecoration cursor distance auto-reject', function () {
     })
 
     it('should reject when cursor moves more than 25 lines before the edit', async function () {
+        // Set cursor at line 50
+        editorStub.selection = {
+            active: new vscode.Position(50, 0),
+        } as any
         const startLine = 50
         await setupDisplaySvgDecoration(startLine)
 
@@ -289,6 +388,10 @@ describe('displaySvgDecoration cursor distance auto-reject', function () {
     })
 
     it('should not reject when edit is near beginning of file and cursor cannot move far enough', async function () {
+        // Set cursor at line 10
+        editorStub.selection = {
+            active: new vscode.Position(10, 0),
+        } as any
         const startLine = 10
         await setupDisplaySvgDecoration(startLine)
 
@@ -298,6 +401,10 @@ describe('displaySvgDecoration cursor distance auto-reject', function () {
     })
 
     it('should not reject when edit suggestion is not active', async function () {
+        // Set cursor at line 50
+        editorStub.selection = {
+            active: new vscode.Position(50, 0),
+        } as any
         editSuggestionStateStub.returns(false)
 
         const startLine = 50

packages/core/src/shared/clients/sagemaker.ts

@@ -52,6 +52,11 @@ import { yes, no, continueText, cancel } from '../localizedText'
 import { AwsCredentialIdentity } from '@aws-sdk/types'
 import globals from '../extensionGlobals'
 
+const appTypeSettingsMap: Record<string, string> = {
+    [AppType.JupyterLab as string]: 'JupyterLabAppSettings',
+    [AppType.CodeEditor as string]: 'CodeEditorAppSettings',
+} as const
+
 export interface SagemakerSpaceApp extends SpaceDetails {
     App?: AppDetails
     DomainSpaceKey: string
@@ -136,13 +141,13 @@ export class SagemakerClient extends ClientWrapper<SageMakerClient> {
 
         // Get app type
         const appType = spaceDetails.SpaceSettings?.AppType
-        if (appType !== 'JupyterLab' && appType !== 'CodeEditor') {
+        if (!appType || !(appType in appTypeSettingsMap)) {
             throw new ToolkitError(`Unsupported AppType "${appType}" for space "${spaceName}"`)
         }
 
         // Get app resource spec
         const requestedResourceSpec =
-            appType === 'JupyterLab'
+            appType === AppType.JupyterLab
                 ? spaceDetails.SpaceSettings?.JupyterLabAppSettings?.DefaultResourceSpec
                 : spaceDetails.SpaceSettings?.CodeEditorAppSettings?.DefaultResourceSpec
 
@@ -181,16 +186,30 @@ export class SagemakerClient extends ClientWrapper<SageMakerClient> {
             instanceType = InstanceTypeMinimum
         }
 
-        // Get remote access flag
-        if (!spaceDetails.SpaceSettings?.RemoteAccess || spaceDetails.SpaceSettings?.RemoteAccess === 'DISABLED') {
+        // First, update the space if needed
+        const needsRemoteAccess =
+            !spaceDetails.SpaceSettings?.RemoteAccess || spaceDetails.SpaceSettings?.RemoteAccess === 'DISABLED'
+        const instanceTypeChanged = requestedResourceSpec?.InstanceType !== instanceType
+
+        if (needsRemoteAccess || instanceTypeChanged) {
+            const updateSpaceRequest: UpdateSpaceCommandInput = {
+                DomainId: domainId,
+                SpaceName: spaceName,
+                SpaceSettings: {
+                    ...(needsRemoteAccess && { RemoteAccess: 'ENABLED' }),
+                    ...(instanceTypeChanged && {
+                        [appTypeSettingsMap[appType]]: {
+                            DefaultResourceSpec: {
+                                InstanceType: instanceType,
+                            },
+                        },
+                    }),
+                },
+            }
+
             try {
-                await this.updateSpace({
-                    DomainId: domainId,
-                    SpaceName: spaceName,
-                    SpaceSettings: {
-                        RemoteAccess: 'ENABLED',
-                    },
-                })
+                getLogger().debug('SagemakerClient: Updating space: domainId=%s, spaceName=%s', domainId, spaceName)
+                await this.updateSpace(updateSpaceRequest)
                 await this.waitForSpaceInService(spaceName, domainId)
             } catch (err) {
                 throw this.handleStartSpaceError(err)
@@ -214,6 +233,7 @@ export class SagemakerClient extends ClientWrapper<SageMakerClient> {
                 ? { ...resourceSpec, EnvironmentArn: undefined, EnvironmentVersionArn: undefined }
                 : resourceSpec
 
+        // Second, create the App
         const createAppRequest: CreateAppCommandInput = {
             DomainId: domainId,
             SpaceName: spaceName,
@@ -223,6 +243,7 @@ export class SagemakerClient extends ClientWrapper<SageMakerClient> {
         }
 
         try {
+            getLogger().debug('SagemakerClient: Creating app: domainId=%s, spaceName=%s', domainId, spaceName)
             await this.createApp(createAppRequest)
         } catch (err) {
             throw this.handleStartSpaceError(err)

packages/core/src/test/shared/clients/sagemakerClient.test.ts

@@ -360,6 +360,7 @@ describe('SagemakerClient.startSpace', function () {
         getTestWindow().getFirstMessage().selectItem('Yes')
 
         await promise
+        sinon.assert.calledOnce(updateSpaceStub)
         sinon.assert.calledOnce(createAppStub)
     })
 

packages/toolkit/.changes/next-release/Bug Fix-f521ea3f-0ac5-4fc0-b939-9776b7a7fd88.json

@@ -0,0 +1,4 @@
+{
+	"type": "Bug Fix",
+	"description": "The space is updated upon creation of a new app with the requested settings"
+}