Skip to content

Commit c17a0f1

Browse files
tverneytverney
committed
fix(dev): replace gradle dependency for our own cancellation token
1 parent 914b140 commit c17a0f1

File tree

11 files changed

+50
-20
lines changed

11 files changed

+50
-20
lines changed

‎plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/controller/FeatureDevController.kt‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -304,7 +304,7 @@ class FeatureDevController(
304304
messenger.sendChatInputEnabledMessage(tabId = message.tabId, enabled = false)
305305
session = getSessionInfo(message.tabId)
306306

307-
if (session.sessionState.token?.token() !== null) {
307+
if (session.sessionState.token?.token !== null) {
308308
session.sessionState.token?.cancel()
309309
}
310310
}

‎plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/controller/FeatureDevControllerExtensions.kt‎

Lines changed: 6 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@ import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.session.NewFil
2323
import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.session.PrepareCodeGenerationState
2424
import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.session.Session
2525
import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.session.SessionState
26+
import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.util.CancellationTokenSource
2627
import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.util.getFollowUpOptions
2728
import software.aws.toolkits.jetbrains.utils.notifyInfo
2829
import software.aws.toolkits.resources.message
@@ -55,7 +56,7 @@ suspend fun FeatureDevController.onCodeGeneration(
5556
var remainingIterations: Int? = state.codeGenerationRemainingIterationCount
5657
var totalIterations: Int? = state.codeGenerationTotalIterationCount
5758

58-
if (state.token?.token()?.isCancellationRequested == true) {
59+
if (state.token?.token?.isCancellationRequested() == true) {
5960
disposeToken(state, messenger, tabId, state.currentIteration?.let { CODE_GENERATION_RETRY_LIMIT.minus(it) }, CODE_GENERATION_RETRY_LIMIT)
6061
return
6162
}
@@ -80,7 +81,7 @@ suspend fun FeatureDevController.onCodeGeneration(
8081
}
8182
}
8283

83-
if (state.token?.token()?.isCancellationRequested == true) {
84+
if (state.token?.token?.isCancellationRequested() == true) {
8485
disposeToken(state, messenger, tabId, state.currentIteration?.let { CODE_GENERATION_RETRY_LIMIT.minus(it) }, CODE_GENERATION_RETRY_LIMIT)
8586
return
8687
}
@@ -135,10 +136,10 @@ suspend fun FeatureDevController.onCodeGeneration(
135136
messenger.sendUpdatePlaceholder(tabId = tabId, newPlaceholder = message("amazonqFeatureDev.placeholder.after_code_generation"))
136137
} finally {
137138
if (session.sessionState.token
138-
?.token()
139-
?.isCancellationRequested == true
139+
?.token
140+
?.isCancellationRequested() == true
140141
) {
141-
session.sessionState.token = GradleConnector.newCancellationTokenSource()
142+
session.sessionState.token = CancellationTokenSource()
142143
} else {
143144
messenger.sendAsyncEventProgress(tabId = tabId, inProgress = false) // Finish processing the event
144145
messenger.sendChatInputEnabledMessage(tabId = tabId, enabled = false) // Lock chat input until a follow-up is clicked.

‎plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/session/CodeGenerationState.kt‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,6 @@
44
package software.aws.toolkits.jetbrains.services.amazonqFeatureDev.session
55

66
import kotlinx.coroutines.delay
7-
import org.gradle.tooling.CancellationTokenSource
87
import software.amazon.awssdk.services.codewhispererruntime.model.CodeGenerationWorkflowStatus
98
import software.aws.toolkits.core.utils.getLogger
109
import software.aws.toolkits.core.utils.warn
@@ -20,6 +19,7 @@ import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.PromptRefusalE
2019
import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.ThrottlingException
2120
import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.messages.sendAnswerPart
2221
import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.messages.sendUpdatePlaceholder
22+
import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.util.CancellationTokenSource
2323
import software.aws.toolkits.jetbrains.services.cwc.controller.chat.telemetry.getStartUrl
2424
import software.aws.toolkits.resources.message
2525
import software.aws.toolkits.telemetry.AmazonqTelemetry
@@ -63,7 +63,7 @@ class CodeGenerationState(
6363
currentCodeGenerationId = currentCodeGenerationId.toString(),
6464
)
6565

66-
if (action.token?.token()?.isCancellationRequested != true) {
66+
if (action.token?.token?.isCancellationRequested() != true) {
6767
this.currentCodeGenerationId = codeGenerationId.toString()
6868
}
6969

@@ -144,7 +144,7 @@ private suspend fun CodeGenerationState.generateCode(
144144
val requestDelay = 10000L
145145

146146
repeat(pollCount) {
147-
if (token?.token()?.isCancellationRequested == true) {
147+
if (token?.token?.isCancellationRequested() == true) {
148148
return CodeGenerationResult(emptyList(), emptyList(), emptyList())
149149
}
150150
val codeGenerationResultState =

‎plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/session/ConversationNotStartedState.kt‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,8 @@
33

44
package software.aws.toolkits.jetbrains.services.amazonqFeatureDev.session
55

6-
import org.gradle.tooling.CancellationTokenSource
6+
import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.util.CancellationTokenSource
7+
78

89
class ConversationNotStartedState(
910
override var approach: String,

‎plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/session/PrepareCodeGenerationState.kt‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,13 +3,13 @@
33

44
package software.aws.toolkits.jetbrains.services.amazonqFeatureDev.session
55

6-
import org.gradle.tooling.CancellationTokenSource
76
import software.aws.toolkits.core.utils.getLogger
87
import software.aws.toolkits.core.utils.warn
98
import software.aws.toolkits.jetbrains.services.amazonq.messages.MessagePublisher
109
import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.FEATURE_NAME
1110
import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.messages.sendAnswerPart
1211
import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.messages.sendUpdatePlaceholder
12+
import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.util.CancellationTokenSource
1313
import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.util.deleteUploadArtifact
1414
import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.util.uploadArtifactToS3
1515
import software.aws.toolkits.jetbrains.services.cwc.controller.chat.telemetry.getStartUrl

‎plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/session/Session.kt‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@ import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.FEATURE_NAME
1515
import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.MAX_PROJECT_SIZE_BYTES
1616
import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.clients.FeatureDevClient
1717
import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.messages.sendAsyncEventProgress
18+
import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.util.CancellationTokenSource
1819
import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.util.FeatureDevService
1920
import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.util.resolveAndCreateOrUpdateFile
2021
import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.util.resolveAndDeleteFile
@@ -82,7 +83,7 @@ class Session(val tabID: String, val project: Project) {
8283
currentIteration = 1, // first code gen iteration
8384
uploadId = "", // There is no code gen uploadId so far
8485
messenger = messenger,
85-
token = GradleConnector.newCancellationTokenSource()
86+
token = CancellationTokenSource()
8687
)
8788
}
8889

‎plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/session/SessionState.kt‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,8 @@
33

44
package software.aws.toolkits.jetbrains.services.amazonqFeatureDev.session
55

6-
import org.gradle.tooling.CancellationTokenSource
6+
import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.util.CancellationTokenSource
7+
78

89
interface SessionState {
910
val tabID: String

‎plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/session/SessionStateTypes.kt‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,8 +4,8 @@
44
package software.aws.toolkits.jetbrains.services.amazonqFeatureDev.session
55

66
import com.fasterxml.jackson.annotation.JsonValue
7-
import org.gradle.tooling.CancellationTokenSource
87
import software.aws.toolkits.jetbrains.services.amazonq.FeatureDevSessionContext
8+
import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.util.CancellationTokenSource
99
import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.util.FeatureDevService
1010
import software.aws.toolkits.jetbrains.services.cwc.messages.RecommendationContentSpan
1111

‎plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/util/CancellationTokenSource.kt‎

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,24 @@
1+
// Copyright 2024 Amazon.com, Inc. or its affiliates. All Rights Reserved.
2+
// SPDX-License-Identifier: Apache-2.0
3+
4+
package software.aws.toolkits.jetbrains.services.amazonqFeatureDev.util
5+
6+
class CancellationToken {
7+
private var isCancelled = false
8+
9+
fun isCancellationRequested(): Boolean = isCancelled
10+
11+
internal fun cancel() {
12+
isCancelled = true
13+
}
14+
}
15+
16+
class CancellationTokenSource {
17+
private val _token = CancellationToken()
18+
val token: CancellationToken
19+
get() = _token
20+
21+
fun cancel() {
22+
_token.cancel()
23+
}
24+
}

‎plugins/amazonq/chat/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/controller/FeatureDevControllerTest.kt‎

Lines changed: 6 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -58,6 +58,7 @@ import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.session.Sessio
5858
import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.session.SessionStateConfig
5959
import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.session.SessionStatePhase
6060
import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.storage.ChatSessionStorage
61+
import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.util.CancellationTokenSource
6162
import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.util.FeatureDevService
6263
import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.util.getFollowUpOptions
6364
import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.util.selectFolder
@@ -224,7 +225,7 @@ class FeatureDevControllerTest : FeatureDevTestBase() {
224225
whenever(spySession.sessionState).thenReturn(
225226
PrepareCodeGenerationState(
226227
testTabId,
227-
GradleConnector.newCancellationTokenSource(),
228+
CancellationTokenSource(),
228229
"test-command",
229230
sessionStateConfig,
230231
newFileContents,
@@ -285,7 +286,7 @@ class FeatureDevControllerTest : FeatureDevTestBase() {
285286
whenever(mockSession.sessionState).thenReturn(
286287
PrepareCodeGenerationState(
287288
testTabId,
288-
GradleConnector.newCancellationTokenSource(),
289+
CancellationTokenSource(),
289290
"test-command",
290291
sessionStateConfig,
291292
newFileContents,
@@ -344,7 +345,7 @@ class FeatureDevControllerTest : FeatureDevTestBase() {
344345
whenever(mockSession.sessionState).thenReturn(
345346
PrepareCodeGenerationState(
346347
testTabId,
347-
GradleConnector.newCancellationTokenSource(),
348+
CancellationTokenSource(),
348349
"",
349350
mock(),
350351
filePaths,
@@ -383,7 +384,7 @@ class FeatureDevControllerTest : FeatureDevTestBase() {
383384
whenever(mockSession.sessionState).thenReturn(
384385
PrepareCodeGenerationState(
385386
testTabId,
386-
GradleConnector.newCancellationTokenSource(),
387+
CancellationTokenSource(),
387388
"",
388389
mock(),
389390
filePaths,
@@ -415,7 +416,7 @@ class FeatureDevControllerTest : FeatureDevTestBase() {
415416
whenever(spySession.sessionState).thenReturn(
416417
PrepareCodeGenerationState(
417418
testTabId,
418-
GradleConnector.newCancellationTokenSource(),
419+
CancellationTokenSource(),
419420
"",
420421
mock(),
421422
newFileContents,

0 commit comments

Comments
 (0)