fix(amazonq): not set IAM auth as default for SMAI remote ssh

Author: yueny2020

packages/core/src/shared/lsp/utils/platform.ts

@@ -8,7 +8,7 @@ import { ToolkitError } from '../../errors'
 import { Logger } from '../../logger/logger'
 import { ChildProcess } from '../../utilities/processUtils'
 import { waitUntil } from '../../utilities/timeoutUtils'
-import { isDebugInstance } from '../../vscode/env'
+import { isDebugInstance, isRemoteWorkspace } from '../../vscode/env'
 import { isSageMaker } from '../../extensionUtilities'
 import { getLogger } from '../../logger/logger'
 
@@ -124,8 +124,16 @@ export function createServerOptions({
                     getLogger().info(`[SageMaker Debug] Using SSO auth mode, not setting USE_IAM_AUTH`)
                 }
             } catch (err) {
-                getLogger().warn(`[SageMaker Debug] Failed to parse SageMaker cookies, defaulting to IAM auth: ${err}`)
-                env.USE_IAM_AUTH = 'true'
+                if (isRemoteWorkspace() && env.SERVICE_NAME !== 'SageMakerUnifiedStudio') {
+                    getLogger().warn(
+                        `[SageMaker Debug] Failed to parse SageMaker cookies in remote space, not SMUS env, not defaulting to IAM auth: ${err}`
+                    )
+                } else {
+                    getLogger().warn(
+                        `[SageMaker Debug] Failed to parse SageMaker cookies, defaulting to IAM auth: ${err}`
+                    )
+                    env.USE_IAM_AUTH = 'true'
+                }
             }
 
             // Log important environment variables for debugging

packages/core/src/test/shared/lsp/utils/platform.test.ts

@@ -0,0 +1,125 @@
+/*!
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import assert from 'assert'
+import * as sinon from 'sinon'
+import * as vscode from 'vscode'
+import { createServerOptions } from '../../../../shared/lsp/utils/platform'
+import * as extensionUtilities from '../../../../shared/extensionUtilities'
+import * as env from '../../../../shared/vscode/env'
+import { getLogger } from '../../../../shared/logger/logger'
+import { ChildProcess } from '../../../../shared/utilities/processUtils'
+
+describe('createServerOptions - SageMaker Authentication', function () {
+    let sandbox: sinon.SinonSandbox
+    let isSageMakerStub: sinon.SinonStub
+    let isRemoteWorkspaceStub: sinon.SinonStub
+    let executeCommandStub: sinon.SinonStub
+    let loggerInfoStub: sinon.SinonStub
+    let loggerWarnStub: sinon.SinonStub
+    let originalEnv: NodeJS.ProcessEnv
+
+    beforeEach(function () {
+        sandbox = sinon.createSandbox()
+        originalEnv = { ...process.env }
+
+        isSageMakerStub = sandbox.stub(extensionUtilities, 'isSageMaker')
+        isRemoteWorkspaceStub = sandbox.stub(env, 'isRemoteWorkspace')
+        sandbox.stub(env, 'isDebugInstance').returns(false)
+        executeCommandStub = sandbox.stub(vscode.commands, 'executeCommand')
+
+        const logger = getLogger()
+        loggerInfoStub = sandbox.stub(logger, 'info')
+        loggerWarnStub = sandbox.stub(logger, 'warn')
+
+        sandbox.stub(ChildProcess.prototype, 'run').resolves()
+        sandbox.stub(ChildProcess.prototype, 'send').resolves()
+        sandbox.stub(ChildProcess.prototype, 'proc').returns({} as any)
+    })
+
+    afterEach(function () {
+        sandbox.restore()
+        process.env = originalEnv
+    })
+
+    it('sets USE_IAM_AUTH=true when authMode is Iam', async function () {
+        isSageMakerStub.returns(true)
+        executeCommandStub.withArgs('sagemaker.parseCookies').resolves({ authMode: 'Iam' })
+
+        const serverOptions = createServerOptions({
+            encryptionKey: Buffer.from('test-key'),
+            executable: ['node'],
+            serverModule: 'test-module.js',
+            execArgv: ['--stdio'],
+        })
+
+        await serverOptions()
+
+        assert.ok(
+            loggerInfoStub.calledWith(
+                '[SageMaker Debug] Setting USE_IAM_AUTH=true for language server process (authMode: Iam)'
+            )
+        )
+    })
+
+    it('does not set USE_IAM_AUTH when authMode is Sso', async function () {
+        isSageMakerStub.returns(true)
+        executeCommandStub.withArgs('sagemaker.parseCookies').resolves({ authMode: 'Sso' })
+
+        const serverOptions = createServerOptions({
+            encryptionKey: Buffer.from('test-key'),
+            executable: ['node'],
+            serverModule: 'test-module.js',
+            execArgv: ['--stdio'],
+        })
+
+        await serverOptions()
+
+        assert.ok(loggerInfoStub.calledWith('[SageMaker Debug] Using SSO auth mode, not setting USE_IAM_AUTH'))
+    })
+
+    it('defaults to IAM auth when parseCookies fails', async function () {
+        isSageMakerStub.returns(true)
+        isRemoteWorkspaceStub.returns(false)
+        executeCommandStub.withArgs('sagemaker.parseCookies').rejects(new Error('Command failed'))
+
+        const serverOptions = createServerOptions({
+            encryptionKey: Buffer.from('test-key'),
+            executable: ['node'],
+            serverModule: 'test-module.js',
+            execArgv: ['--stdio'],
+        })
+
+        await serverOptions()
+
+        assert.ok(
+            loggerWarnStub.calledWith(
+                '[SageMaker Debug] Failed to parse SageMaker cookies, defaulting to IAM auth: Error: Command failed'
+            )
+        )
+    })
+
+    it('does not default to IAM in remote workspace without SMUS', async function () {
+        isSageMakerStub.returns(true)
+        isRemoteWorkspaceStub.returns(true)
+        process.env.SERVICE_NAME = 'OtherService'
+        executeCommandStub.withArgs('sagemaker.parseCookies').rejects(new Error('Command failed'))
+
+        const serverOptions = createServerOptions({
+            encryptionKey: Buffer.from('test-key'),
+            executable: ['node'],
+            serverModule: 'test-module.js',
+            execArgv: ['--stdio'],
+        })
+
+        await serverOptions()
+
+        assert.ok(
+            loggerWarnStub.calledWith(
+                '[SageMaker Debug] Failed to parse SageMaker cookies in remote space, not SMUS env, not defaulting to IAM auth: Error: Command failed'
+            )
+        )
+    })
+})