wip

hayemaxi · Feb 11, 2025 · ca955b6 · ca955b6
1 parent b1a396c
commit ca955b6
Show file tree

Hide file tree

Showing 21 changed files with 578 additions and 195 deletions.
diff --git a/packages/amazonq/.vscode/launch.json b/packages/amazonq/.vscode/launch.json
@@ -142,10 +142,14 @@
             "type": "node",
             "request": "attach",
             "port": 6080, // Hard defined in core/src/shared/lsp/platform.ts
-            "outFiles": ["${workspaceFolder}/../../../language-servers/**/out/**/*.js"],
+            "outFiles": [
+                "${workspaceFolder}/../../../language-servers/**/out/**/*.js",
+                "${workspaceFolder}/../../../language-server-runtimes/**/out/**/*.js"
+            ],
             "skipFiles": [
                 "<node_internals>/**",
-                "${workspaceFolder}/../../../language-servers/**/node_modules/**/*.js"
+                "${workspaceFolder}/../../../language-servers/**/node_modules/**/*.js",
+                "${workspaceFolder}/../../../language-server-runtimes/**/node_modules/**/*.js"
             ],
             "restart": {
                 "maxAttempts": 10,

diff --git a/packages/amazonq/src/extension.ts b/packages/amazonq/src/extension.ts
@@ -4,7 +4,12 @@
  */
 
 import { AuthUtils, CredentialsStore, LoginManager, initializeAuth } from 'aws-core-vscode/auth'
-import { activate as activateCodeWhisperer, shutdown as shutdownCodeWhisperer } from 'aws-core-vscode/codewhisperer'
+import {
+    AuthUtil,
+    activate as activateCodeWhisperer,
+    shutdown as shutdownCodeWhisperer,
+    postInit,
+} from 'aws-core-vscode/codewhisperer'
 import { makeEndpointsProvider, registerGenericCommands } from 'aws-core-vscode'
 import { CommonAuthWebview } from 'aws-core-vscode/login'
 import {
@@ -117,10 +122,16 @@ export async function activateAmazonQCommon(context: vscode.ExtensionContext, is
         extensionContext: context,
     }
     // This contains every lsp agnostic things (auth, security scan, code scan)
-    await activateCodeWhisperer(extContext as ExtContext)
+    // await activateCodeWhisperer(extContext as ExtContext)
     if (Experiments.instance.get('amazonqLSP', false)) {
+        await activateCodeWhisperer(extContext as ExtContext)
         await activateAmazonqLsp(context)
+        AuthUtil.instance.initCodeWhispererHooks()
+        await postInit(extContext as ExtContext)
     } else {
+        await activateCodeWhisperer(extContext as ExtContext)
+        AuthUtil.instance.initCodeWhispererHooks()
+        await postInit(extContext as ExtContext)
         await activateInlineCompletion()
     }
 

diff --git a/packages/amazonq/src/lsp/auth.ts b/packages/amazonq/src/lsp/auth.ts
@@ -10,12 +10,10 @@ import {
     ResponseMessage,
 } from '@aws/language-server-runtimes/protocol'
 import * as jose from 'jose'
-import * as crypto from 'crypto'
 import { LanguageClient } from 'vscode-languageclient'
 import { AuthUtil } from 'aws-core-vscode/codewhisperer'
 import { Writable } from 'stream'
-
-export const encryptionKey = crypto.randomBytes(32)
+import { Auth2 } from 'aws-core-vscode/auth'
 
 /**
  * Sends a json payload to the language server, who is waiting to know what the encryption key is.
@@ -25,7 +23,7 @@ export function writeEncryptionInit(stream: Writable): void {
     const request = {
         version: '1.0',
         mode: 'JWT',
-        key: encryptionKey.toString('base64'),
+        key: Auth2.encryptionKey.toString('base64'),
     }
     stream.write(JSON.stringify(request))
     stream.write('\n')
@@ -88,7 +86,7 @@ export class AmazonQLspAuth {
 
         const jwt = await new jose.CompactEncrypt(payload)
             .setProtectedHeader({ alg: 'dir', enc: 'A256GCM' })
-            .encrypt(encryptionKey)
+            .encrypt(Auth2.encryptionKey)
 
         return {
             data: jwt,

diff --git a/packages/amazonq/src/lsp/client.ts b/packages/amazonq/src/lsp/client.ts
@@ -8,20 +8,26 @@ import * as nls from 'vscode-nls'
 import * as crypto from 'crypto'
 import { LanguageClient, LanguageClientOptions } from 'vscode-languageclient'
 import { registerInlineCompletion } from '../app/inline/completion'
-import { AmazonQLspAuth, encryptionKey, notificationTypes } from './auth'
+import { notificationTypes } from './auth'
 import { AuthUtil } from 'aws-core-vscode/codewhisperer'
-import { ConnectionMetadata } from '@aws/language-server-runtimes/protocol'
-import { ResourcePaths, Settings, createServerOptions, globals } from 'aws-core-vscode/shared'
+import { ResourcePaths, Settings, createServerOptions, globals, getLogger, openUrl } from 'aws-core-vscode/shared'
+import {
+    ConnectionMetadata,
+    ShowDocumentParams,
+    ShowDocumentRequest,
+    ShowDocumentResult,
+} from '@aws/language-server-runtimes/protocol'
+import { Auth2, SsoConnection, getRegistrationCacheFile, getTokenCacheFile, getCacheDir } from 'aws-core-vscode/auth'
 
 const localize = nls.loadMessageBundle()
 
-export async function startLanguageServer(extensionContext: vscode.ExtensionContext, resourcePaths: ResourcePaths) {
+export function startLanguageServer(extensionContext: vscode.ExtensionContext, resourcePaths: ResourcePaths) {
     const toDispose = extensionContext.subscriptions
 
     const serverModule = resourcePaths.lsp
 
     const serverOptions = createServerOptions({
-        encryptionKey,
+        encryptionKey: Auth2.encryptionKey,
         executable: resourcePaths.node,
         serverModule,
         execArgv: [
@@ -75,20 +81,13 @@ export async function startLanguageServer(extensionContext: vscode.ExtensionCont
               }),
     }
 
-    const client = new LanguageClient(
-        clientId,
-        localize('amazonq.server.name', 'Amazon Q Language Server'),
-        serverOptions,
-        clientOptions
-    )
+    const lspName = localize('amazonq.server.name', 'Amazon Q Language Server')
+    const client = new LanguageClient(clientId, lspName, serverOptions, clientOptions)
 
     const disposable = client.start()
     toDispose.push(disposable)
 
-    const auth = new AmazonQLspAuth(client)
-
     return client.onReady().then(async () => {
-        await auth.init()
         registerInlineCompletion(client)
 
         // Request handler for when the server wants to know about the clients auth connnection
@@ -100,13 +99,37 @@ export async function startLanguageServer(extensionContext: vscode.ExtensionCont
             }
         })
 
-        toDispose.push(
-            AuthUtil.instance.auth.onDidChangeActiveConnection(async () => {
-                await auth.init()
-            }),
-            AuthUtil.instance.auth.onDidDeleteConnection(async () => {
-                client.sendNotification(notificationTypes.deleteBearerToken.method)
-            })
-        )
+        client.onRequest<ShowDocumentResult, Error>(ShowDocumentRequest.method, async (params: ShowDocumentParams) => {
+            try {
+                return { success: await openUrl(vscode.Uri.parse(params.uri), lspName) }
+            } catch (err: any) {
+                getLogger().error(`Failed to open document for LSP: ${lspName}, error: %s`, err)
+                return { success: false }
+            }
+        })
+
+        // toDispose.push(
+        //     AuthUtil.instance.auth.onDidChangeActiveConnection(async () => {
+        //         await auth.init()
+        //     }),
+        //     AuthUtil.instance.auth.onDidDeleteConnection(async () => {
+        //         client.sendNotification(notificationTypes.deleteBearerToken.method)
+        //     })
+        // )
+        Auth2.create(client)
+        const conn = (AuthUtil.instance.conn as SsoConnection) ?? AuthUtil.instance.auth.activeConnection
+        if (conn) {
+            await Auth2.instance.importOldSsoSession(
+                conn.startUrl,
+                conn.ssoRegion,
+                getRegistrationCacheFile(getCacheDir(), {
+                    startUrl: conn.startUrl,
+                    region: conn.ssoRegion,
+                    scopes: conn.scopes,
+                }),
+                getTokenCacheFile(getCacheDir(), (conn.id ?? conn.startUrl) as any)
+            )
+            await AuthUtil.instance.secondaryAuth.deleteConnection()
+        }
     })
 }
diff --git a/packages/core/src/amazonq/lsp/workspaceInstaller.ts b/packages/core/src/amazonq/lsp/workspaceInstaller.ts
@@ -35,7 +35,7 @@ export class WorkspaceLSPResolver implements LspResolver {
 
         const deletedVersions = await cleanLspDownloads(
             manifest.versions,
-            path.basename(installationResult.assetDirectory)
+            path.dirname(installationResult.assetDirectory)
         )
         logger.debug(`cleaning old LSP versions deleted ${deletedVersions.length} versions`)
         return {