@@ -13,81 +13,175 @@ import {
13
13
TextDocument ,
14
14
commands ,
15
15
languages ,
16
+ Disposable ,
16
17
} from 'vscode'
17
18
import { LanguageClient } from 'vscode-languageclient'
18
- import {
19
- InlineCompletionListWithReferences ,
20
- InlineCompletionWithReferencesParams ,
21
- inlineCompletionWithReferencesRequestType ,
22
- logInlineCompletionSessionResultsNotificationType ,
23
- LogInlineCompletionSessionResultsParams ,
24
- } from '@aws/language-server-runtimes/protocol'
19
+ import { LogInlineCompletionSessionResultsParams } from '@aws/language-server-runtimes/protocol'
20
+ import { SessionManager } from './sessionManager'
21
+ import { RecommendationService } from './recommendationService'
25
22
import { CodeWhispererConstants } from 'aws-core-vscode/codewhisperer'
26
23
27
- export function registerInlineCompletion ( languageClient : LanguageClient ) {
28
- const inlineCompletionProvider = new AmazonQInlineCompletionItemProvider ( languageClient )
29
- languages . registerInlineCompletionItemProvider ( CodeWhispererConstants . platformLanguageIds , inlineCompletionProvider )
24
+ export class InlineCompletionManager implements Disposable {
25
+ private disposable : Disposable
26
+ private inlineCompletionProvider : AmazonQInlineCompletionItemProvider
27
+ private languageClient : LanguageClient
28
+ private sessionManager : SessionManager
29
+ private recommendationService : RecommendationService
30
+ private readonly logSessionResultMessageName = 'aws/logInlineCompletionSessionResults'
31
+
32
+ constructor ( languageClient : LanguageClient ) {
33
+ this . languageClient = languageClient
34
+ this . sessionManager = new SessionManager ( )
35
+ this . recommendationService = new RecommendationService ( this . sessionManager )
36
+ this . inlineCompletionProvider = new AmazonQInlineCompletionItemProvider (
37
+ languageClient ,
38
+ this . recommendationService ,
39
+ this . sessionManager
40
+ )
41
+ this . disposable = languages . registerInlineCompletionItemProvider (
42
+ CodeWhispererConstants . platformLanguageIds ,
43
+ this . inlineCompletionProvider
44
+ )
45
+ }
46
+
47
+ public dispose ( ) : void {
48
+ if ( this . disposable ) {
49
+ this . disposable . dispose ( )
50
+ }
51
+ }
52
+
53
+ public registerInlineCompletion ( ) {
54
+ const onInlineAcceptance = async (
55
+ sessionId : string ,
56
+ itemId : string ,
57
+ requestStartTime : number ,
58
+ firstCompletionDisplayLatency ?: number
59
+ ) => {
60
+ // TODO: also log the seen state for other suggestions in session
61
+ const params : LogInlineCompletionSessionResultsParams = {
62
+ sessionId : sessionId ,
63
+ completionSessionResult : {
64
+ [ itemId ] : {
65
+ seen : true ,
66
+ accepted : true ,
67
+ discarded : false ,
68
+ } ,
69
+ } ,
70
+ totalSessionDisplayTime : Date . now ( ) - requestStartTime ,
71
+ firstCompletionDisplayLatency : firstCompletionDisplayLatency ,
72
+ }
73
+ this . languageClient . sendNotification ( this . logSessionResultMessageName , params )
74
+ this . disposable . dispose ( )
75
+ this . disposable = languages . registerInlineCompletionItemProvider (
76
+ CodeWhispererConstants . platformLanguageIds ,
77
+ this . inlineCompletionProvider
78
+ )
79
+ }
80
+ commands . registerCommand ( 'aws.amazonq.acceptInline' , onInlineAcceptance )
30
81
31
- const onInlineAcceptance = async (
32
- sessionId : string ,
33
- itemId : string ,
34
- requestStartTime : number ,
35
- firstCompletionDisplayLatency ?: number
36
- ) => {
37
- const params : LogInlineCompletionSessionResultsParams = {
38
- sessionId : sessionId ,
39
- completionSessionResult : {
40
- [ itemId ] : {
41
- seen : true ,
42
- accepted : true ,
43
- discarded : false ,
82
+ const onInlineRejection = async ( ) => {
83
+ await commands . executeCommand ( 'editor.action.inlineSuggest.hide' )
84
+ // TODO: also log the seen state for other suggestions in session
85
+ this . disposable . dispose ( )
86
+ this . disposable = languages . registerInlineCompletionItemProvider (
87
+ CodeWhispererConstants . platformLanguageIds ,
88
+ this . inlineCompletionProvider
89
+ )
90
+ const sessionId = this . sessionManager . getActiveSession ( ) ?. sessionId
91
+ const itemId = this . sessionManager . getActiveRecommendation ( ) [ 0 ] ?. itemId
92
+ if ( ! sessionId || ! itemId ) {
93
+ return
94
+ }
95
+ const params : LogInlineCompletionSessionResultsParams = {
96
+ sessionId : sessionId ,
97
+ completionSessionResult : {
98
+ [ itemId ] : {
99
+ seen : true ,
100
+ accepted : false ,
101
+ discarded : false ,
102
+ } ,
44
103
} ,
45
- } ,
46
- totalSessionDisplayTime : Date . now ( ) - requestStartTime ,
47
- firstCompletionDisplayLatency : firstCompletionDisplayLatency ,
104
+ }
105
+ this . languageClient . sendNotification ( this . logSessionResultMessageName , params )
48
106
}
49
- languageClient . sendNotification ( logInlineCompletionSessionResultsNotificationType as any , params )
107
+ commands . registerCommand ( 'aws.amazonq.rejectCodeSuggestion' , onInlineRejection )
108
+
109
+ /*
110
+ We have to overwrite the prev. and next. commands because the inlineCompletionProvider only contained the current item
111
+ To show prev. and next. recommendation we need to re-register a new provider with the previous or next item
112
+ */
113
+
114
+ const swapProviderAndShow = async ( ) => {
115
+ await commands . executeCommand ( 'editor.action.inlineSuggest.hide' )
116
+ this . disposable . dispose ( )
117
+ this . disposable = languages . registerInlineCompletionItemProvider (
118
+ CodeWhispererConstants . platformLanguageIds ,
119
+ new AmazonQInlineCompletionItemProvider (
120
+ this . languageClient ,
121
+ this . recommendationService ,
122
+ this . sessionManager ,
123
+ false
124
+ )
125
+ )
126
+ await commands . executeCommand ( 'editor.action.inlineSuggest.trigger' )
127
+ }
128
+
129
+ const prevCommandHandler = async ( ) => {
130
+ this . sessionManager . decrementActiveIndex ( )
131
+ await swapProviderAndShow ( )
132
+ }
133
+ commands . registerCommand ( 'editor.action.inlineSuggest.showPrevious' , prevCommandHandler )
134
+
135
+ const nextCommandHandler = async ( ) => {
136
+ this . sessionManager . incrementActiveIndex ( )
137
+ await swapProviderAndShow ( )
138
+ }
139
+ commands . registerCommand ( 'editor.action.inlineSuggest.showNext' , nextCommandHandler )
50
140
}
51
- commands . registerCommand ( 'aws.sample-vscode-ext-amazonq.accept' , onInlineAcceptance )
52
141
}
53
142
54
143
export class AmazonQInlineCompletionItemProvider implements InlineCompletionItemProvider {
55
- constructor ( private readonly languageClient : LanguageClient ) { }
144
+ constructor (
145
+ private readonly languageClient : LanguageClient ,
146
+ private readonly recommendationService : RecommendationService ,
147
+ private readonly sessionManager : SessionManager ,
148
+ private readonly isNewSession : boolean = true
149
+ ) { }
56
150
57
151
async provideInlineCompletionItems (
58
152
document : TextDocument ,
59
153
position : Position ,
60
154
context : InlineCompletionContext ,
61
155
token : CancellationToken
62
156
) : Promise < InlineCompletionItem [ ] | InlineCompletionList > {
63
- const requestStartTime = Date . now ( )
64
- const request : InlineCompletionWithReferencesParams = {
65
- textDocument : {
66
- uri : document . uri . toString ( ) ,
67
- } ,
68
- position,
69
- context,
157
+ if ( this . isNewSession ) {
158
+ // make service requests if it's a new session
159
+ await this . recommendationService . getAllRecommendations (
160
+ this . languageClient ,
161
+ document ,
162
+ position ,
163
+ context ,
164
+ token
165
+ )
70
166
}
71
-
72
- const response = await this . languageClient . sendRequest (
73
- inlineCompletionWithReferencesRequestType as any ,
74
- request ,
75
- token
76
- )
77
-
78
- const list : InlineCompletionListWithReferences = response as InlineCompletionListWithReferences
79
- this . languageClient . info ( `Client: Received ${ list . items . length } suggestions` )
80
- const firstCompletionDisplayLatency = Date . now ( ) - requestStartTime
81
-
82
- // Add completion session tracking and attach onAcceptance command to each item to record used decision
83
- for ( const item of list . items ) {
167
+ // get active item from session for displaying
168
+ const items = this . sessionManager . getActiveRecommendation ( )
169
+ const session = this . sessionManager . getActiveSession ( )
170
+ if ( ! session || ! items . length ) {
171
+ return [ ]
172
+ }
173
+ for ( const item of items ) {
84
174
item . command = {
85
- command : 'aws.sample-vscode-ext- amazonq.accept ' ,
175
+ command : 'aws.amazonq.acceptInline ' ,
86
176
title : 'On acceptance' ,
87
- arguments : [ list . sessionId , item . itemId , requestStartTime , firstCompletionDisplayLatency ] ,
177
+ arguments : [
178
+ session . sessionId ,
179
+ item . itemId ,
180
+ session . requestStartTime ,
181
+ session . firstCompletionDisplayLatency ,
182
+ ] ,
88
183
}
89
184
}
90
-
91
- return list as InlineCompletionList
185
+ return items as InlineCompletionItem [ ]
92
186
}
93
187
}
0 commit comments