opensource/kilocode
3
0
Fork 0
mirror of https://github.com/Kilo-Org/kilocode.git synced 2026-09-24 16:02:55 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix(jetbrains): defer turn state across revert instead of dropping it

Browse Source 
Address review: guarding TurnOpen/TurnClose/status with revertOp dropped
the transitions entirely, so a turn that started just before a rollback
could leave the prompt idle while the server turn was still active. Record
the underlying turn/status transition in revertDeferred while the revert is
held and reconcile it when the operation releases: an aborted turn releases
to Idle, a still-active turn releases back to Busy/Retry/Offline. Failed
reverts still surface the error explicitly.
This commit is contained in:
kirillk
2026-07-10 13:49:29 -04:00
parent af0297087d
commit 1fe08892f0
2 changed files with 66 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files
packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/session/controller/SessionController.kt
+43 -6
View File
@@ -155,6 +155,9 @@ class SessionController(
private var revertOp: RevertOp? = null
private var revertSeq = 0L
private var revertWatchdog: UiTimer? = null
// While a revert is in flight, turn/status transitions are deferred here instead of dropped,
// then reconciled when the operation releases so an underlying server turn is not lost.
private var revertDeferred: SessionState? = null
private var creating: CompletableDeferred<String?>? = null
private val childJobs: MutableMap<String, Job> = mutableMapOf()
private val childIds: MutableSet<String> = mutableSetOf()
@@ -1289,14 +1292,20 @@ class SessionController(
is ChatEventDto.TurnOpen -> {
partType = null
tool = null
if (revertOp != null) return
if (revertOp != null) {
revertDeferred = SessionState.Busy(KiloBundle.message("session.status.considering"))
return
}
model.setState(SessionState.Busy(KiloBundle.message("session.status.considering")))
}
is ChatEventDto.TurnClose -> {
partType = null
tool = null
if (revertOp != null) return
if (revertOp != null) {
revertDeferred = SessionState.Idle
return
}
// Keep pending questions visible for follow-up flows that arrive just before close.
val current = model.state
if (current is SessionState.AwaitingQuestion) return
@@ -1469,7 +1478,16 @@ class SessionController(
}
private fun status(dto: SessionStatusDto) {
if (revertOp != null) return
if (revertOp != null) {
revertDeferred = when (dto.type) {
"idle" -> SessionState.Idle
"busy" -> SessionState.Busy(KiloBundle.message("session.status.considering"))
"retry" -> SessionState.Retry(dto.message ?: "", dto.attempt ?: 0, dto.next ?: 0L)
"offline" -> SessionState.Offline(dto.message ?: "", dto.requestID ?: "")
else -> revertDeferred
}
return
}
val state = when (dto.type) {
"idle" -> {
val current = model.state
@@ -1501,6 +1519,10 @@ class SessionController(
if (revertOp != null) return null
val op = RevertOp(++revertSeq)
revertOp = op
// Start with no deferred state; only turn/status transitions that actually arrive while the
// revert is held are recorded, so an aborted turn releases to Idle and a still-active turn
// (e.g. one that opened after the busy check) releases back to Busy.
revertDeferred = null
model.setState(SessionState.Reverting(text, kind, message))
startRevertWatchdog(op)
return op
@@ -1534,7 +1556,7 @@ class SessionController(
stopRevertWatchdog()
revertJob = null
revertOp = null
if (model.state is SessionState.Reverting) model.setState(SessionState.Idle)
reconcileReverting()
}
private fun cancelReverting(op: RevertOp) {
@@ -1543,7 +1565,16 @@ class SessionController(
stopRevertWatchdog()
revertJob = null
revertOp = null
if (model.state is SessionState.Reverting) model.setState(SessionState.Idle)
reconcileReverting()
}
// Release the revert lock back to whatever turn/status transition arrived while it was held,
// so an underlying server turn is restored instead of leaving an idle prompt over an active turn.
private fun reconcileReverting() {
assertEdt()
val next = revertDeferred ?: SessionState.Idle
revertDeferred = null
if (model.state is SessionState.Reverting) model.setState(next)
}
private fun failReverting(op: RevertOp, e: Exception) {
@@ -1552,11 +1583,16 @@ class SessionController(
stopRevertWatchdog()
revertJob = null
revertOp = null
// A failed revert surfaces the error explicitly; drop any deferred turn state.
revertDeferred = null
model.setState(SessionState.Error(e.message ?: KiloBundle.message("session.error.unknown")))
}
private fun idle() {
if (revertOp != null) return
if (revertOp != null) {
revertDeferred = SessionState.Idle
return
}
// Treat session.idle as an explicit signal to return to Idle.
// Only apply if we're not in a more specific non-terminal state.
val current = model.state
@@ -2162,6 +2198,7 @@ class SessionController(
revertJob?.cancel()
revertJob = null
revertOp = null
revertDeferred = null
val callbacks = enhancements.values.toList()
enhancements.clear()
cs.cancel()
packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/session/controller/TurnLifecycleTest.kt
+23
View File
@@ -599,6 +599,29 @@ class TurnLifecycleTest : SessionControllerTestBase() {
flush()
}
fun `test revert reconciles active turn deferred during rollback`() {
val (m, _, _) = prompted()
// Land on idle so the rollback does not abort, then start a gated rollback.
emit(ChatEventDto.TurnClose("ses_test", reason = "completed"))
assertTrue(m.model.state is SessionState.Idle)
val gate = CompletableDeferred<Unit>()
rpc.revertGate = gate
edt { m.revert("msg1") }
settle()
assertTrue(m.model.state is SessionState.Reverting)
// A server turn opens while the rollback is in flight; the transition is deferred, not dropped.
emit(ChatEventDto.TurnOpen("ses_test"))
assertTrue("turn open must not clear reverting", m.model.state is SessionState.Reverting)
gate.complete(Unit)
flush()
// Releasing the revert restores the still-active turn instead of forcing an idle prompt.
assertTrue("expected Busy, was ${m.model.state}", m.model.state is SessionState.Busy)
}
fun `test rollback while busy enters reverting and aborts`() {
val (m, _, _) = prompted()
emit(ChatEventDto.TurnOpen("ses_test"))