Skip to content

Commit 170e77f

Browse files
lkk214lkk214
authored
feat(sketch): Add interrupt handling and process listeners (#300)
1 parent b2a4dd5 commit 170e77f

File tree

6 files changed

+115
-31
lines changed

6 files changed

+115
-31
lines changed

‎core/src/main/kotlin/cc/unitmesh/devti/gui/chat/ChatCodingPanel.kt

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -259,12 +259,13 @@ class ChatCodingPanel(private val chatCodingService: ChatCodingService, val disp
259259
messageView.updateContent(text)
260260
}
261261

262-
if (delaySeconds.isNotEmpty()) {
262+
if (delaySeconds.isNotBlank()) {
263263
val elapsedTime = System.currentTimeMillis() - startTime
264264
withContext(Dispatchers.IO) {
265-
val delaySec = delaySeconds.toLong()
266-
val remainingTime = maxOf(delaySec * 1000 - elapsedTime, 0)
267-
delay(remainingTime)
265+
delaySeconds.toLongOrNull()?.let {
266+
val remainingTime = maxOf(it * 1000 - elapsedTime, 0)
267+
delay(remainingTime)
268+
}
268269
}
269270
}
270271

‎core/src/main/kotlin/cc/unitmesh/devti/inline/AutoDevInlineChatPanel.kt

Lines changed: 47 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@ package cc.unitmesh.devti.inline
22

33
import cc.unitmesh.devti.llms.LlmFactory
44
import cc.unitmesh.devti.llms.cancelHandler
5+
import cc.unitmesh.devti.sketch.SketchProcessListener
56
import cc.unitmesh.devti.sketch.SketchToolWindow
67
import cc.unitmesh.devti.util.AutoDevCoroutineScope
78
import com.intellij.icons.AllIcons
@@ -11,6 +12,7 @@ import com.intellij.openapi.actionSystem.Presentation
1112
import com.intellij.openapi.actionSystem.impl.ActionButton
1213
import com.intellij.openapi.application.ApplicationManager
1314
import com.intellij.openapi.application.invokeLater
15+
import com.intellij.openapi.application.runInEdt
1416
import com.intellij.openapi.editor.Editor
1517
import com.intellij.openapi.editor.EditorCustomElementRenderer
1618
import com.intellij.openapi.editor.Inlay
@@ -35,7 +37,7 @@ import javax.swing.*
3537
class AutoDevInlineChatPanel(val editor: Editor) : JPanel(GridBagLayout()), EditorCustomElementRenderer,
3638
Disposable {
3739
var inlay: Inlay<*>? = null
38-
val inputPanel = AutoDevInlineChatInput(this, onSubmit = { input ->
40+
val inputPanel = AutoDevInlineChatInput(this, onSubmit = { input, onCreated ->
3941
this.centerPanel.isVisible = true
4042
val project = editor.project!!
4143

@@ -45,6 +47,7 @@ class AutoDevInlineChatPanel(val editor: Editor) : JPanel(GridBagLayout()), Edit
4547
val panelView = SketchToolWindow(project, editor)
4648
panelView.minimumSize = Dimension(800, 40)
4749
addToCenter(panelView)
50+
onCreated(panelView) // Add process listener before onStart
4851

4952
AutoDevCoroutineScope.scope(project).launch {
5053
val suggestion = StringBuilder()
@@ -62,6 +65,7 @@ class AutoDevInlineChatPanel(val editor: Editor) : JPanel(GridBagLayout()), Edit
6265
panelView.resize()
6366
panelView.onFinish(suggestion.toString())
6467
}
68+
panelView
6569
})
6670
private var centerPanel: JPanel = JPanel(BorderLayout())
6771
private var container: Container? = null
@@ -172,10 +176,14 @@ class AutoDevInlineChatPanel(val editor: Editor) : JPanel(GridBagLayout()), Edit
172176

173177
class AutoDevInlineChatInput(
174178
val autoDevInlineChatPanel: AutoDevInlineChatPanel,
175-
val onSubmit: (String) -> Unit,
179+
val onSubmit: (String, (SketchToolWindow) -> Unit) -> SketchToolWindow,
176180
) : JPanel(GridBagLayout()), Disposable {
177181
private val textArea: JBTextArea
178182

183+
private var view: SketchToolWindow? = null
184+
185+
private var btnPresentation: Presentation? = null
186+
179187
init {
180188
layout = BorderLayout()
181189
textArea = object : JBTextArea(), KeyboardAwareFocusOwner {
@@ -195,6 +203,7 @@ class AutoDevInlineChatInput(
195203
// escape to close
196204
textArea.actionMap.put("escapeAction", object : AbstractAction() {
197205
override fun actionPerformed(e: ActionEvent) {
206+
cancel()
198207
AutoDevInlineChatService.getInstance().closeInlineChat(autoDevInlineChatPanel.editor)
199208
}
200209
})
@@ -214,21 +223,50 @@ class AutoDevInlineChatInput(
214223
})
215224
textArea.inputMap.put(KeyStroke.getKeyStroke(KeyEvent.VK_ENTER, KeyEvent.SHIFT_DOWN_MASK), "newlineAction")
216225

217-
val submitPresentation = Presentation("Submit")
218-
submitPresentation.icon = AllIcons.Actions.Execute
219-
val submitButton = ActionButton(
220-
DumbAwareAction.create { submit() },
221-
submitPresentation, "", Dimension(40, 20)
226+
btnPresentation = Presentation()
227+
setPresentationTextAndIcon(false)
228+
val actionBtn = ActionButton(
229+
DumbAwareAction.create { onEnter() },
230+
btnPresentation, "", Dimension(40, 20)
222231
)
223232

224233
add(textArea)
225-
add(submitButton, BorderLayout.EAST)
234+
add(actionBtn, BorderLayout.EAST)
235+
}
236+
237+
private fun onEnter() {
238+
if (btnPresentation?.icon == AllIcons.Actions.Execute) submit()
239+
else if (btnPresentation?.icon == AllIcons.Actions.Suspend) cancel()
226240
}
227241

228242
private fun submit() {
243+
view?.cancel("Cancel by resubmit") // Or not allowed to submit at runtime
229244
val trimText = textArea.text.trim()
230245
textArea.text = ""
231-
onSubmit(trimText)
246+
view = onSubmit(trimText) {
247+
it.addProcessListener(object : SketchProcessListener {
248+
override fun onBefore() = setPresentationTextAndIcon(true)
249+
override fun onAfter() = setPresentationTextAndIcon(false)
250+
})
251+
}
252+
253+
}
254+
255+
private fun cancel() {
256+
view?.cancel("Cancel")
257+
setPresentationTextAndIcon(false)
258+
}
259+
260+
private fun setPresentationTextAndIcon(running: Boolean) {
261+
runInEdt {
262+
if (running) {
263+
btnPresentation?.text = "Cancel"
264+
btnPresentation?.icon = AllIcons.Actions.Suspend
265+
} else {
266+
btnPresentation?.text = "Submit"
267+
btnPresentation?.icon = AllIcons.Actions.Execute
268+
}
269+
}
232270
}
233271

234272
fun getInputComponent(): Component = textArea

‎core/src/main/kotlin/cc/unitmesh/devti/llms/custom/CustomSSEProcessor.kt

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@ import cc.unitmesh.devti.coder.recording.JsonlRecording
55
import cc.unitmesh.devti.coder.recording.Recording
66
import cc.unitmesh.devti.coder.recording.RecordingInstruction
77
import cc.unitmesh.devti.gui.chat.message.ChatRole
8+
import cc.unitmesh.devti.llms.CustomFlowWrapper
89
import cc.unitmesh.devti.settings.coder.coderSetting
910
import com.fasterxml.jackson.databind.ObjectMapper
1011
import com.intellij.openapi.components.service
@@ -72,14 +73,15 @@ open class CustomSSEProcessor(private val project: Project) {
7273

7374
@OptIn(kotlinx.coroutines.ExperimentalCoroutinesApi::class)
7475
fun streamSSE(call: Call, promptText: String, keepHistory: Boolean = false, messages: MutableList<Message>): Flow<String> {
76+
var emit: FlowableEmitter<SSE>? = null
7577
val sseFlowable = Flowable
7678
.create({ emitter: FlowableEmitter<SSE> ->
77-
call.enqueue(ResponseBodyCallback(emitter, true))
79+
emit = emitter.apply { call.enqueue(ResponseBodyCallback(emitter, true)) }
7880
}, BackpressureStrategy.BUFFER)
7981

8082
try {
8183
var output = ""
82-
return callbackFlow {
84+
return CustomFlowWrapper(callbackFlow {
8385
withContext(Dispatchers.IO) {
8486
sseFlowable
8587
.doOnError {
@@ -143,7 +145,7 @@ open class CustomSSEProcessor(private val project: Project) {
143145
close()
144146
}
145147
awaitClose()
146-
}
148+
}).also { it.cancelCallback { emit?.onComplete() } }
147149
} catch (e: Exception) {
148150
if (hasSuccessRequest) {
149151
logger.info("Failed to stream", e)

‎core/src/main/kotlin/cc/unitmesh/devti/settings/devops/AutoDevDevOpsConfigurableProvider.kt

Lines changed: 4 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,6 @@
11
package cc.unitmesh.devti.settings.devops
22

33
import cc.unitmesh.devti.AutoDevBundle
4-
import cc.unitmesh.devti.settings.custom.TeamPromptsProjectSettingsService
54
import com.intellij.openapi.components.*
65
import com.intellij.openapi.options.BoundConfigurable
76
import com.intellij.openapi.options.Configurable
@@ -51,9 +50,9 @@ class DevOpsConfigurable(project: Project) : BoundConfigurable(AutoDevBundle.mes
5150

5251
override fun apply() {
5352
settings.modify { state ->
54-
state.githubToken = githubTokenField.password.toString()
53+
state.githubToken = githubTokenField.password.concatToString()
5554
state.gitlabUrl = gitlabUrlField.text
56-
state.gitlabToken = gitlabTokenField.password.toString()
55+
state.gitlabToken = gitlabTokenField.password.concatToString()
5756
}
5857
}
5958

@@ -64,9 +63,9 @@ class DevOpsConfigurable(project: Project) : BoundConfigurable(AutoDevBundle.mes
6463
}
6564

6665
override fun isModified(): Boolean {
67-
return settings.state.githubToken != githubTokenField.password.toString() ||
66+
return settings.state.githubToken != githubTokenField.password.concatToString() ||
6867
settings.state.gitlabUrl != gitlabUrlField.text ||
69-
settings.state.gitlabToken != gitlabTokenField.password.toString()
68+
settings.state.gitlabToken != gitlabTokenField.password.concatToString()
7069
}
7170
}
7271

‎core/src/main/kotlin/cc/unitmesh/devti/sketch/SketchInputListener.kt

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@ import cc.unitmesh.devti.gui.chat.ChatCodingService
55
import cc.unitmesh.devti.gui.chat.ui.AutoDevInputListener
66
import cc.unitmesh.devti.gui.chat.ui.AutoDevInputSection
77
import cc.unitmesh.devti.gui.chat.ui.AutoDevInputTrigger
8+
import cc.unitmesh.devti.llms.cancelHandler
89
import cc.unitmesh.devti.prompting.SimpleDevinPrompter
910
import cc.unitmesh.devti.provider.devins.LanguageProcessor
1011
import cc.unitmesh.devti.template.GENIUS_CODE
@@ -38,6 +39,7 @@ class SketchInputListener(
3839
override fun onStop(component: AutoDevInputSection) {
3940
chatCodingService.stop()
4041
toolWindow.hiddenProgressBar()
42+
toolWindow.stop()
4143
}
4244

4345
override fun onSubmit(component: AutoDevInputSection, trigger: AutoDevInputTrigger) {
@@ -57,14 +59,15 @@ class SketchInputListener(
5759
val devInProcessor = LanguageProcessor.devin()
5860
val compiledInput = runReadAction { devInProcessor?.compile(project, userInput) } ?: userInput
5961

62+
toolWindow.beforeRun()
6063
toolWindow.updateHistoryPanel()
6164
toolWindow.addRequestPrompt(compiledInput)
6265

6366
val flow = chatCodingService.request(systemPrompt, compiledInput)
6467
val suggestion = StringBuilder()
6568

6669
AutoDevCoroutineScope.workerThread().launch {
67-
flow.cancellable().collect { char ->
70+
flow.cancelHandler { toolWindow.handleCancel = it }.cancellable().collect { char ->
6871
suggestion.append(char)
6972

7073
invokeLater {

‎core/src/main/kotlin/cc/unitmesh/devti/sketch/SketchToolWindow.kt

Lines changed: 50 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -46,11 +46,16 @@ import java.awt.event.MouseAdapter
4646
import java.awt.event.MouseEvent
4747
import javax.swing.*
4848

49+
interface SketchProcessListener {
50+
fun onBefore() {}
51+
fun onAfter() {}
52+
}
53+
4954
class SketchToolWindow(val project: Project, val editor: Editor?, private val showInput: Boolean = false) :
5055
SimpleToolWindowPanel(true, true), NullableComponent, Disposable {
5156
private val chatCodingService = ChatCodingService(ChatActionType.SKETCH, project)
5257
private var progressBar: CustomProgressBar = CustomProgressBar(this)
53-
private var shireInput: AutoDevInputSection = AutoDevInputSection(project, this, showAgent = false)
58+
private var inputSection: AutoDevInputSection = AutoDevInputSection(project, this, showAgent = false)
5459

5560
private var myText: String = ""
5661

@@ -63,6 +68,8 @@ class SketchToolWindow(val project: Project, val editor: Editor?, private val sh
6368

6469
private var isUserScrolling: Boolean = false
6570

71+
private var isInterrupted: Boolean = false
72+
6673
private var systemPrompt: JPanel = JPanel(BorderLayout())
6774
private var contentPanel = JPanel(BorderLayout())
6875

@@ -102,6 +109,8 @@ class SketchToolWindow(val project: Project, val editor: Editor?, private val sh
102109

103110
private val listener = SketchInputListener(project, chatCodingService, this)
104111

112+
private val processListeners = mutableListOf<SketchProcessListener>()
113+
105114
init {
106115
if (showInput) {
107116
val header = panel {
@@ -118,7 +127,7 @@ class SketchToolWindow(val project: Project, val editor: Editor?, private val sh
118127

119128
header.border = JBUI.Borders.compound(
120129
JBUI.Borders.customLine(UIUtil.getBoundsColor(), 0, 0, 1, 0),
121-
JBUI.Borders.empty(0, 4, 0, 4)
130+
JBUI.Borders.empty(0, 4)
122131
)
123132

124133
contentPanel.add(header, BorderLayout.NORTH)
@@ -138,18 +147,29 @@ class SketchToolWindow(val project: Project, val editor: Editor?, private val sh
138147
contentPanel.add(progressBar, BorderLayout.SOUTH)
139148

140149
if (showInput) {
141-
shireInput.also {
150+
inputSection.also {
142151
it.border = JBUI.Borders.empty(8)
143152
}
144153

145-
shireInput.addListener(listener)
146-
contentPanel.add(shireInput, BorderLayout.SOUTH)
154+
inputSection.addListener(listener)
155+
contentPanel.add(inputSection, BorderLayout.SOUTH)
156+
157+
addProcessListener(object : SketchProcessListener {
158+
override fun onBefore() {
159+
isInterrupted = false
160+
inputSection.showStopButton()
161+
}
162+
override fun onAfter() {
163+
inputSection.showSendButton()
164+
}
165+
})
147166
}
148167

149168
setContent(contentPanel)
150169
}
151170

152171
fun onStart() {
172+
beforeRun()
153173
initializePreAllocatedBlocks(project)
154174
progressBar.isIndeterminate = true
155175
progressBar.isVisible = !showInput
@@ -159,6 +179,22 @@ class SketchToolWindow(val project: Project, val editor: Editor?, private val sh
159179
progressBar.isVisible = false
160180
}
161181

182+
fun stop() {
183+
cancel("Stop")
184+
inputSection.showSendButton()
185+
}
186+
187+
fun addProcessListener(processorListener: SketchProcessListener) {
188+
processListeners.add(processorListener)
189+
}
190+
191+
fun beforeRun() {
192+
processListeners.forEach { it.onBefore() }
193+
}
194+
fun AfterRun() {
195+
processListeners.forEach { it.onAfter() }
196+
}
197+
162198
private val blockViews: MutableList<LangSketch> = mutableListOf()
163199
private fun initializePreAllocatedBlocks(project: Project) {
164200
repeat(32) {
@@ -287,14 +323,16 @@ class SketchToolWindow(val project: Project, val editor: Editor?, private val sh
287323
progressBar.isVisible = false
288324
scrollToBottom()
289325

290-
if (AutoSketchMode.getInstance(project).isEnable) {
326+
AfterRun()
327+
328+
if (AutoSketchMode.getInstance(project).isEnable && !isInterrupted) {
291329
AutoSketchMode.getInstance(project).start(text, this@SketchToolWindow.listener)
292330
}
293331
}
294332

295333
fun sendInput(text: String) {
296-
shireInput.text += "\n" + text
297-
shireInput.send()
334+
inputSection.text += "\n" + text
335+
inputSection.send()
298336
}
299337

300338
private fun scrollToBottom() {
@@ -318,7 +356,10 @@ class SketchToolWindow(val project: Project, val editor: Editor?, private val sh
318356

319357
override fun isNull(): Boolean = !isVisible
320358

321-
fun cancel(s: String) = runCatching { handleCancel?.invoke(s) }
359+
fun cancel(s: String) = runCatching {
360+
handleCancel?.also { handleCancel = null }?.invoke(s)
361+
isInterrupted = true
362+
}
322363

323364
fun resetSketchSession() {
324365
chatCodingService.clearSession()

0 commit comments

Comments
 (0)