chore(amazonq): refactor codeWhispereServer.ts

Author: atonaamz

server/aws-lsp-codewhisperer/src/language-server/inline-completion/codeWhispererServer.ts

@@ -9,35 +9,35 @@ import {
     ResponseError,
     TextDocument,
 } from '@aws/language-server-runtimes/protocol'
-import { RecentEditTracker } from './tracker/codeEditTracker'
+import { RecentEditTracker } from '../tracker/codeEditTracker'
 import { CredentialsProvider, Logging, Telemetry, Workspace } from '@aws/language-server-runtimes/server-interface'
 import {
     CodeWhispererServiceToken,
     GenerateSuggestionsRequest,
     GenerateSuggestionsResponse,
     getFileContext,
     SuggestionType,
-} from '../../shared/codeWhispererService'
-import { CodeWhispererSession, SessionManager } from './session/sessionManager'
-import { CursorTracker } from './tracker/cursorTracker'
-import { CodewhispererLanguage, getSupportedLanguageId } from '../../shared/languageDetection'
-import { WorkspaceFolderManager } from '../workspaceContext/workspaceFolderManager'
-import { shouldTriggerEdits } from './trigger'
+} from '../../../shared/codeWhispererService'
+import { CodeWhispererSession, SessionManager } from '../session/sessionManager'
+import { CursorTracker } from '../tracker/cursorTracker'
+import { CodewhispererLanguage, getSupportedLanguageId } from '../../../shared/languageDetection'
+import { WorkspaceFolderManager } from '../../workspaceContext/workspaceFolderManager'
+import { shouldTriggerEdits } from '../trigger'
 import {
     emitEmptyUserTriggerDecisionTelemetry,
     emitServiceInvocationFailure,
     emitServiceInvocationTelemetry,
     emitUserTriggerDecisionTelemetry,
-} from './telemetry/telemetry'
-import { TelemetryService } from '../../shared/telemetry/telemetryService'
+} from '../telemetry/telemetry'
+import { TelemetryService } from '../../../shared/telemetry/telemetryService'
 import { textUtils } from '@aws/lsp-core'
-import { AmazonQBaseServiceManager } from '../../shared/amazonQServiceManager/BaseAmazonQServiceManager'
-import { RejectedEditTracker } from './tracker/rejectedEditTracker'
-import { getErrorMessage, hasConnectionExpired } from '../../shared/utils'
-import { AmazonQError, AmazonQServiceConnectionExpiredError } from '../../shared/amazonQServiceManager/errors'
-import { DocumentChangedListener } from './documentChangedListener'
-import { EMPTY_RESULT, EDIT_DEBOUNCE_INTERVAL_MS } from './contants/constants'
-import { StreakTracker } from './tracker/streakTracker'
+import { AmazonQBaseServiceManager } from '../../../shared/amazonQServiceManager/BaseAmazonQServiceManager'
+import { RejectedEditTracker } from '../tracker/rejectedEditTracker'
+import { getErrorMessage, hasConnectionExpired } from '../../../shared/utils'
+import { AmazonQError, AmazonQServiceConnectionExpiredError } from '../../../shared/amazonQServiceManager/errors'
+import { DocumentChangedListener } from '../documentChangedListener'
+import { EMPTY_RESULT, EDIT_DEBOUNCE_INTERVAL_MS } from '../contants/constants'
+import { StreakTracker } from '../tracker/streakTracker'
 
 export class EditCompletionHandler {
     private readonly editsEnabled: boolean

server/aws-lsp-codewhisperer/src/language-server/inline-completion/editCompletionHandler.ts renamed to server/aws-lsp-codewhisperer/src/language-server/inline-completion/handler/editCompletionHandler.ts

@@ -0,0 +1,146 @@
+import * as assert from 'assert'
+import * as sinon from 'sinon'
+import { TextDocument } from 'vscode-languageserver-textdocument'
+import { InlineCompletionHandler } from './inlineCompletionHandler'
+import { SessionManager } from '../session/sessionManager'
+import { CodePercentageTracker } from '../tracker/codePercentageTracker'
+import { RecentEditTracker } from '../tracker/codeEditTracker'
+import { CursorTracker } from '../tracker/cursorTracker'
+import { StreakTracker } from '../tracker/streakTracker'
+import { TelemetryService } from '../../../shared/telemetry/telemetryService'
+import { UserWrittenCodeTracker } from '../../../shared/userWrittenCodeTracker'
+import { InlineCompletionTriggerKind, CancellationToken } from '@aws/language-server-runtimes/server-interface'
+import { EMPTY_RESULT } from '../contants/constants'
+import * as IdleWorkspaceManagerModule from '../../workspaceContext/IdleWorkspaceManager'
+import * as telemetryModule from '../telemetry/telemetry'
+
+describe('InlineCompletionHandler', () => {
+    const testDocument = TextDocument.create('file:///test.cs', 'csharp', 1, 'test content')
+
+    const completionParams = {
+        textDocument: { uri: testDocument.uri },
+        position: { line: 0, character: 0 },
+        context: { triggerKind: InlineCompletionTriggerKind.Invoked },
+    }
+
+    let handler: InlineCompletionHandler
+    let completionSessionManager: SessionManager
+    let amazonQServiceManager: any
+    let codePercentageTracker: sinon.SinonStubbedInstance<CodePercentageTracker>
+    let userWrittenCodeTracker: sinon.SinonStubbedInstance<UserWrittenCodeTracker>
+    let recentEditTracker: sinon.SinonStubbedInstance<RecentEditTracker>
+    let cursorTracker: sinon.SinonStubbedInstance<CursorTracker>
+    let streakTracker: sinon.SinonStubbedInstance<StreakTracker>
+    let telemetryService: sinon.SinonStubbedInstance<TelemetryService>
+    let lsp: any
+    let telemetry: any
+    let credentialsProvider: any
+    let workspace: any
+    let logging: any
+    let getTextDocument: sinon.SinonStub
+
+    beforeEach(() => {
+        SessionManager.reset()
+        completionSessionManager = SessionManager.getInstance('COMPLETIONS')
+
+        amazonQServiceManager = {
+            getCodewhispererService: sinon.stub(),
+            getConfiguration: sinon.stub().returns({ inlineSuggestions: {} }),
+        }
+        codePercentageTracker = sinon.createStubInstance(CodePercentageTracker)
+        userWrittenCodeTracker = sinon.createStubInstance(UserWrittenCodeTracker)
+        recentEditTracker = sinon.createStubInstance(RecentEditTracker)
+        cursorTracker = sinon.createStubInstance(CursorTracker)
+        streakTracker = sinon.createStubInstance(StreakTracker)
+        telemetryService = sinon.createStubInstance(TelemetryService)
+
+        workspace = { getWorkspaceFolder: sinon.stub() }
+        logging = { log: sinon.stub(), debug: sinon.stub() }
+        getTextDocument = sinon.stub().resolves(testDocument)
+        lsp = { getClientInitializeParams: sinon.stub() }
+        telemetry = { emitMetric: sinon.stub() } as any
+        credentialsProvider = { getConnectionMetadata: sinon.stub() } as any
+
+        // Stub IdleWorkspaceManager and telemetry functions
+        sinon.stub(IdleWorkspaceManagerModule.IdleWorkspaceManager, 'recordActivityTimestamp')
+        sinon.stub(telemetryModule, 'emitServiceInvocationTelemetry')
+        sinon.stub(telemetryModule, 'emitServiceInvocationFailure')
+        sinon.stub(telemetryModule, 'emitUserTriggerDecisionTelemetry')
+
+        handler = new InlineCompletionHandler(
+            logging,
+            workspace,
+            amazonQServiceManager,
+            completionSessionManager,
+            codePercentageTracker,
+            userWrittenCodeTracker,
+            recentEditTracker,
+            cursorTracker,
+            streakTracker,
+            telemetry,
+            telemetryService,
+            credentialsProvider,
+            false,
+            1000,
+            lsp,
+            getTextDocument
+        )
+    })
+
+    afterEach(() => {
+        sinon.restore()
+    })
+
+    it('should return empty result when concurrent request is in progress', async () => {
+        // Make handler busy
+        handler['isOnInlineCompletionHandlerInProgress'] = true
+
+        const result = await handler.onInlineCompletion(completionParams, CancellationToken.None)
+
+        assert.deepEqual(result, EMPTY_RESULT)
+        sinon.assert.calledWith(logging.log, 'Skip concurrent inline completion')
+    })
+
+    it('should return empty result when service manager not initialized', async () => {
+        handler = new InlineCompletionHandler(
+            logging,
+            workspace,
+            null as any,
+            completionSessionManager,
+            codePercentageTracker,
+            userWrittenCodeTracker,
+            recentEditTracker,
+            cursorTracker,
+            streakTracker,
+            { emitMetric: sinon.stub() } as any,
+            telemetryService,
+            { getConnectionMetadata: sinon.stub() } as any,
+            false,
+            1000,
+            { getClientInitializeParams: sinon.stub() },
+            getTextDocument
+        )
+
+        const result = await handler.onInlineCompletion(completionParams, CancellationToken.None)
+
+        assert.deepEqual(result, EMPTY_RESULT)
+        sinon.assert.calledWith(logging.log, 'Amazon Q Service Manager not initialized yet')
+    })
+
+    it('should return empty result when text document not found', async () => {
+        getTextDocument.resolves(null)
+
+        const result = await handler.onInlineCompletion(completionParams, CancellationToken.None)
+
+        assert.deepEqual(result, EMPTY_RESULT)
+        sinon.assert.calledWith(logging.log, `textDocument [${testDocument.uri}] not found`)
+    })
+
+    it('should track cursor position when cursor tracker available', async () => {
+        getTextDocument.resolves(null) // Will return early
+
+        await handler.onInlineCompletion(completionParams, CancellationToken.None)
+
+        sinon.assert.calledWith(cursorTracker.trackPosition, testDocument.uri, completionParams.position)
+    })
+})