Merge remote-tracking branch 'origin/main' into feat/explain-tool-auto-approval

This commit is contained in:
Bruno Agatao
2026-07-24 12:35:42 +02:00
81 changed files with 4239 additions and 1675 deletions
+5
View File
@@ -0,0 +1,5 @@
---
"@kilocode/cli": patch
---
Restore stream idle timeouts to opt-in provider configuration instead of aborting quiet model streams by default.
@@ -0,0 +1,5 @@
---
"@kilocode/cli": patch
---
Fix compaction failure against strict OpenAI-compatible providers during context compaction. The compaction path no longer leaks `maxOutputTokens` into provider options, which was rejected by strict upstreams with "Unsupported parameter(s)".
@@ -0,0 +1,5 @@
---
"@kilocode/kilo-jetbrains": patch
---
Improve JetBrains diff previews by hiding hunk headers and adding full-path tooltips to clickable file links.
+5
View File
@@ -0,0 +1,5 @@
---
"@kilocode/kilo-jetbrains": patch
---
Render edit tool results with a clickable file target and a highlighted, simplified diff view.
+5
View File
@@ -0,0 +1,5 @@
---
"@kilocode/kilo-jetbrains": patch
---
Open edit tool file links directly when multiple files share the same name.
@@ -0,0 +1,5 @@
---
"@kilocode/kilo-jetbrains": patch
---
Render multi-file apply_patch edits as a "Patch" with a file-count tag and one section per file, each showing a clickable filename link and its own changes badge aligned with the diff.
@@ -0,0 +1,5 @@
---
"@kilocode/kilo-jetbrains": patch
---
Smooth out chat scrolling in large JetBrains sessions by only refreshing hover state for the message under the pointer.
@@ -0,0 +1,5 @@
---
"@kilocode/kilo-jetbrains": patch
---
Improve chat scrolling performance in large JetBrains sessions.
@@ -0,0 +1,5 @@
---
"@kilocode/kilo-jetbrains": patch
---
Size edit and shell preview popovers to their content with a wider maximum width.
@@ -0,0 +1,5 @@
---
"@kilocode/cli": minor
---
`kilo remote` instances now advertise themselves on the relay heartbeat. Each heartbeat carries the host's hostname, the project directory name, and the CLI build version, and each session entry advertises the platform it was created on. The cloud relay learns about a freshly-connected instance immediately (no 10s wait for the first timer tick), and the advertisement is race-safe across the explicit `kilo remote` command and bootstrap auto-enable (`KILO_REMOTE=1` / `remote_control` config). Legacy CLIs that send neither field remain wire-compatible.
+2 -4
View File
@@ -115,11 +115,9 @@ export const Info = Schema.Struct({
description:
"Timeout in milliseconds to wait for response headers. Provider integrations may set defaults. Set to false to disable timeout.",
}),
// kilocode_change: accept `false` so internal callers can disable the
// watchdog. PositiveInt already excludes 0, so a public zero stays invalid.
chunkTimeout: Schema.optional(Schema.Union([PositiveInt, Schema.Literal(false)])).annotate({
chunkTimeout: Schema.optional(PositiveInt).annotate({
description:
"Timeout in milliseconds between streamed SSE chunks for this provider. If no chunk arrives within this window, the request is aborted. Set to false to disable the idle watchdog.",
"Timeout in milliseconds between streamed SSE chunks for this provider. If no chunk arrives within this window, the request is aborted.",
}),
}),
[Schema.Record(Schema.String, Schema.Any)],
@@ -346,7 +346,7 @@ internal class SessionScroll(
@RequiresEdt
private fun layoutScroll() {
root.validate()
component.validate()
}
@RequiresEdt
@@ -8,6 +8,7 @@ import java.awt.Container
import java.awt.Dimension
import java.awt.Insets
import java.awt.LayoutManager
import java.util.IdentityHashMap
/**
* A vertical, width-aware layout manager for the session transcript.
@@ -33,8 +34,12 @@ class SessionLayout(
private val basePad: Insets = JBUI.emptyInsets(),
) : LayoutManager {
private val cache = IdentityHashMap<Component, Measured>()
override fun addLayoutComponent(name: String, comp: Component) = Unit
override fun removeLayoutComponent(comp: Component) = Unit
override fun removeLayoutComponent(comp: Component) {
cache.remove(comp)
}
override fun preferredLayoutSize(parent: Container): Dimension {
val ins = insets(parent)
@@ -46,9 +51,7 @@ class SessionLayout(
if (!first) h += gap(comp)
first = false
val child = bounds(ins, w, comp)
// Pre-size to available width so HTML panes reflow before we measure
comp.setSize(child.width, comp.height.coerceAtLeast(1))
h += comp.preferredSize.height
h += measure(comp, child.width)
}
// w and h are already scaled px (child preferred heights + scaled gaps/insets) and
// match what layoutContainer stacks, so return a plain Dimension. A JBDimension would
@@ -68,14 +71,36 @@ class SessionLayout(
if (!first) y += gap(comp)
first = false
val child = bounds(ins, w, comp)
// Fix width first so HTML reflows, then read the resulting height
comp.setSize(child.width, comp.height.coerceAtLeast(1))
val h = comp.preferredSize.height
val h = measure(comp, child.width)
comp.setBounds(child.left, y, child.width, h)
y += h
}
}
/**
* Drop the cached measurement for [comp] so the next layout pass re-measures it.
*
* [measure] trusts `comp.isValid` as a freshness signal, which is safe only while `comp` is
* invalidated through this container. A child that is its own validate root (see
* [ai.kilocode.client.session.views.TurnView.isValidateRoot]) can be re-validated independently
* by `RepaintManager` — its `isValid` flips back to `true` before this layout re-measures it,
* so a content change that grows/shrinks its height would otherwise return a stale cached value.
* Callers that mutate such a child's content must forget it here so the cache stays honest.
*/
fun forget(comp: Component) {
cache.remove(comp)
}
private fun measure(comp: Component, width: Int): Int {
val hit = cache[comp]
if (comp.isValid && hit?.width == width) return hit.height
// Pre-size to available width so HTML panes reflow before we measure.
comp.setSize(width, comp.height.coerceAtLeast(1))
val h = comp.preferredSize.height
cache[comp] = Measured(width, h)
return h
}
private fun bounds(ins: Insets, width: Int, comp: Component): Bounds {
val view = view(comp) ?: return Bounds(ins.left, width)
if (view.sessionViewKind != SessionView.Kind.UserPrompt) return Bounds(ins.left, width)
@@ -103,6 +128,8 @@ class SessionLayout(
private fun view(comp: Component): SessionView? = comp as? SessionView
private data class Bounds(val left: Int, val width: Int)
private data class Measured(val width: Int, val height: Int)
}
/**
@@ -96,34 +96,37 @@ class SessionMessageListPanel(
is SessionModelEvent.TurnRemoved -> onTurnRemoved(event.id)
is SessionModelEvent.ContentAdded -> {
msgToView[event.messageId]?.upsertPart(event.content)
msgToTurn[event.messageId]?.syncCopyToolbars()
refresh()
if (msgToView[event.messageId]?.upsertPartChanged(event.content) == true) {
onContentChanged(event.messageId)
}
}
is SessionModelEvent.ContentUpdated -> {
msgToView[event.messageId]?.upsertPart(event.content)
msgToTurn[event.messageId]?.syncCopyToolbars()
refresh()
if (msgToView[event.messageId]?.upsertPartChanged(event.content) == true) {
onContentChanged(event.messageId)
}
}
is SessionModelEvent.ContentRemoved -> {
msgToView[event.messageId]?.removePart(event.contentId)
msgToTurn[event.messageId]?.syncCopyToolbars()
refresh()
if (msgToView[event.messageId]?.removePartChanged(event.contentId) == true) {
onContentChanged(event.messageId)
}
}
is SessionModelEvent.ContentDelta -> {
if (event.created) return@addListener
if (event.delta.isEmpty()) return@addListener
val handled = msgToView[event.messageId]?.appendDelta(event.contentId, event.delta) == true
if (handled) {
msgToTurn[event.messageId]?.syncCopyToolbars()
forgetTurn(event.messageId)
return@addListener
}
val content = model.content(event.messageId, event.contentId)
if (content != null) {
msgToView[event.messageId]?.upsertPart(content)
msgToTurn[event.messageId]?.syncCopyToolbars()
if (msgToView[event.messageId]?.upsertPartChanged(content) == true) {
onContentChanged(event.messageId)
}
}
}
@@ -132,6 +135,7 @@ class SessionMessageListPanel(
is SessionModelEvent.StateChanged -> {
syncActive(event.state)
syncSettled(event.state)
syncReverted()
syncReverting(event.state)
anchorFooter()
@@ -222,6 +226,7 @@ class SessionMessageListPanel(
tv.syncCopyToolbars()
syncReverted()
add(tv)
syncSettled()
anchorFooter()
refresh()
}
@@ -234,8 +239,7 @@ class SessionMessageListPanel(
// Remove messages no longer in this turn
for (id in prev) {
if (id !in next) {
tv.removeMessage(id)
unregister(id)
if (tv.removeMessageChanged(id)) unregister(id)
}
}
@@ -248,6 +252,7 @@ class SessionMessageListPanel(
}
tv.syncCopyToolbars()
syncReverted()
syncSettled()
refresh()
}
@@ -257,6 +262,7 @@ class SessionMessageListPanel(
for (msgId in tv.messageIds()) unregister(msgId)
remove(tv)
Disposer.dispose(tv)
syncSettled()
anchorFooter()
refresh()
}
@@ -285,6 +291,7 @@ class SessionMessageListPanel(
}
syncActive(model.state)
syncSettled(model.state)
syncReverted()
syncReverting(model.state)
banner?.update()
@@ -313,6 +320,7 @@ class SessionMessageListPanel(
revertingMessage = null
removeAll()
syncActive(model.state)
syncSettled(model.state)
syncReverting(model.state)
banner?.update()
anchorFooter()
@@ -375,6 +383,11 @@ class SessionMessageListPanel(
for (mv in msgToView.values) mv.setHiddenQuestionTool(ref)
}
private fun syncSettled(state: SessionState = model.state) {
val active = if (state.isBusy()) turnViews.values.lastOrNull() else null
for (view in turnViews.values) view.setSettled(view !== active)
}
/**
* Re-insert [question], [permission], [login], and [progress] as the last children
* so active views always render after all turn views, and progress is last.
@@ -413,6 +426,25 @@ class SessionMessageListPanel(
repaint()
}
/**
* Handle a content mutation that changed an already-rendered message: sync the turn's copy
* toolbars, forget its cached height, then relayout. [forgetTurn] is essential when the update
* lands on a settled turn — a settled [TurnView] is its own validate root, so `RepaintManager`
* re-validates it independently and its `isValid` flag no longer signals the height change to
* [SessionLayout]'s measurement cache.
*/
private fun onContentChanged(messageId: String) {
msgToTurn[messageId]?.syncCopyToolbars()
forgetTurn(messageId)
refresh()
}
/** Drop [SessionLayout]'s cached height for the turn holding [messageId] after its content changes. */
private fun forgetTurn(messageId: String) {
val tv = msgToTurn[messageId] ?: return
(layout as? SessionLayout)?.forget(tv)
}
private fun hover(view: PartView, value: Boolean) {
if (value) {
val prev = hovered
@@ -1,14 +1,20 @@
package ai.kilocode.client.session.ui.popup
import ai.kilocode.client.session.ui.style.SessionUiStyle
import com.intellij.openapi.Disposable
import com.intellij.ui.EditorTextField
import com.intellij.ui.components.JBTextArea
import com.intellij.util.ui.JBUI
import java.awt.BorderLayout
import java.awt.Color
import java.awt.Component
import java.awt.Container
import java.awt.Dimension
import java.awt.Insets
import javax.swing.JComponent
import javax.swing.JEditorPane
import javax.swing.JPanel
import javax.swing.JScrollPane
class HeaderPopupRequest(
val anchor: JComponent,
@@ -20,11 +26,15 @@ class HeaderPopupBody(
component: JComponent,
val disposable: Disposable,
val background: Color,
maxWidth: Int = SessionUiStyle.View.Popup.MAX_WIDTH,
) {
val component: JComponent = HeaderPopupPanel(component)
val component: JComponent = HeaderPopupPanel(component, JBUI.scale(maxWidth))
}
private class HeaderPopupPanel(private val child: JComponent) : JPanel(BorderLayout()) {
private class HeaderPopupPanel(
private val child: JComponent,
private val maxWidth: Int,
) : JPanel(BorderLayout()) {
init {
// Transparent so the balloon fill shows uniformly behind nested popup content.
isOpaque = false
@@ -32,14 +42,32 @@ private class HeaderPopupPanel(private val child: JComponent) : JPanel(BorderLay
}
override fun getPreferredSize(): Dimension {
val size = super.getPreferredSize()
val cap = JBUI.scale(350)
val width = size.width.takeIf { it > 0 }?.coerceAtMost(cap) ?: cap
val width = contentWidth(child).takeIf { it > 0 }?.coerceAtMost(maxWidth) ?: maxWidth
fit(child, width)
val height = super.getPreferredSize().height.coerceAtMost(JBUI.scale(450))
val height = super.getPreferredSize().height.coerceAtMost(JBUI.scale(SessionUiStyle.View.Popup.MAX_HEIGHT))
return Dimension(width, height)
}
private fun contentWidth(item: Component): Int = when (item) {
is EditorTextField -> item.preferredSize.width
is JBTextArea -> item.preferredSize.width
is JEditorPane -> item.preferredSize.width
is JScrollPane -> {
val view = item.viewport?.view?.let(::contentWidth) ?: 0
view + horiz(item.insets) + horiz(item.viewportBorder?.getBorderInsets(item))
}
// JComponent is a Container, so leaf components (labels, buttons, icons) reach here with no
// children — fall back to their own preferred width instead of measuring an empty child set.
is Container -> {
val kids = item.components
if (kids.isEmpty()) (item as? JComponent)?.preferredSize?.width ?: 0
else (kids.maxOfOrNull(::contentWidth) ?: 0) + horiz((item as? JComponent)?.insets)
}
else -> 0
}
private fun horiz(insets: Insets?): Int = (insets?.left ?: 0) + (insets?.right ?: 0)
private fun fit(item: JComponent, width: Int) {
if (width <= 0) return
// JBHtmlPane derives wrapped preferred height from the current width, not just HTML content.
@@ -37,6 +37,12 @@ object SessionUiStyle {
const val BODY_EXTRA_HEIGHT = 16
}
object Popup {
const val MAX_WIDTH = 350
const val WIDE_MAX_WIDTH = MAX_WIDTH * 2
const val MAX_HEIGHT = 450
}
internal const val BORDER_DELTA = 80
internal const val HOVER_BORDER_ALPHA = 0.18f
internal const val HOVER_FILL_ALPHA = 0.10f
@@ -169,6 +175,7 @@ object SessionUiStyle {
object Tool {
const val BODY_LINES = 15
const val TASK_LINES = 10
const val DIFF_LINES = 20
const val PREVIEW_LIMIT = 20_000
fun pending(): Color = UiStyle.Colors.weak()
@@ -7,6 +7,7 @@ import ai.kilocode.client.session.model.FileAttachment
import ai.kilocode.client.session.model.Message
import ai.kilocode.client.session.model.Reasoning
import ai.kilocode.client.session.model.StepFinish
import ai.kilocode.client.session.model.Text
import ai.kilocode.client.session.model.Tool
import ai.kilocode.client.session.model.ToolCallRef
import ai.kilocode.client.session.model.ToolExecState
@@ -110,7 +111,12 @@ class MessageView(
/** Add or update the renderer for [content]. */
@RequiresEdt
fun upsertPart(content: Content) {
if (content is StepFinish) return
upsertPartChanged(content)
}
@RequiresEdt
fun upsertPartChanged(content: Content): Boolean {
if (content is StepFinish) return false
if (isHidden(content)) {
if (isPromptMention(content)) syncPromptMentions()
// Remove any stale view for this content so it disappears when suppressed
@@ -122,7 +128,7 @@ class MessageView(
stale.remove(content.id)
if (!stale.isEmpty()) {
refresh()
return
return true
}
attachments = null
}
@@ -131,14 +137,15 @@ class MessageView(
Disposer.dispose(stale)
syncBorder()
refresh()
return true
}
return
return false
}
val id = aliases[content.id]
if (id != null && content is Reasoning) {
updateAlias(content, id)
if (!updateAlias(content, id)) return false
refresh()
return
return true
}
if (id != null) {
aliases.remove(content.id)
@@ -149,20 +156,24 @@ class MessageView(
if (existing is PromptAttachmentView && content is FileAttachment) {
existing.upsert(content)
refresh()
return
return true
}
if (ViewFactory.shouldReplace(existing, content)) {
replacePart(content, existing)
return
return true
}
if (content is Text && existing is TextView && existing !is PromptView && existing.markdown() == content.content.toString()) {
return false
}
existing.update(content)
syncPromptToolbar()
refresh()
return
return true
}
addPart(content)
syncBorder()
refresh()
return true
}
@RequiresEdt
@@ -203,14 +214,15 @@ class MessageView(
}
@RequiresEdt
private fun updateAlias(content: Reasoning, id: String) {
val view = parts[id] as? ReasoningView ?: return
private fun updateAlias(content: Reasoning, id: String): Boolean {
val view = parts[id] as? ReasoningView ?: return false
val prev = sources[content.id].orEmpty()
val next = content.content.toString()
val delta = if (next.startsWith(prev)) next.removePrefix(prev) else next
sources[content.id] = next
if (delta.isEmpty()) return
if (delta.isEmpty()) return false
view.update(merged(view, content, delta))
return true
}
private fun merged(view: ReasoningView, content: Reasoning, delta: String) = Reasoning(view.contentId).also {
@@ -242,16 +254,21 @@ class MessageView(
/** Remove the renderer for [contentId] if present. */
@RequiresEdt
fun removePart(contentId: String) {
removePartChanged(contentId)
}
@RequiresEdt
fun removePartChanged(contentId: String): Boolean {
if (aliases.remove(contentId) != null) {
sources.remove(contentId)
return
return true
}
val view = parts.remove(contentId) ?: return
val view = parts.remove(contentId) ?: return false
if (view is PromptAttachmentView) {
view.remove(contentId)
if (!view.isEmpty()) {
refresh()
return
return true
}
attachments = null
}
@@ -262,6 +279,7 @@ class MessageView(
Disposer.dispose(view)
syncBorder()
refresh()
return true
}
/**
@@ -335,6 +353,7 @@ class MessageView(
/** Append a streaming delta to the renderer for [contentId]. */
@RequiresEdt
fun appendDelta(contentId: String, delta: String): Boolean {
if (delta.isEmpty()) return false
val id = aliases[contentId]
if (id != null) sources[contentId] = sources[contentId].orEmpty() + delta
val part = parts[id ?: contentId] ?: return false
@@ -12,6 +12,7 @@ import ai.kilocode.client.session.ui.style.SessionUiStyle
import ai.kilocode.client.session.views.base.PartView
import com.intellij.openapi.Disposable
import com.intellij.openapi.util.Disposer
import com.intellij.openapi.util.registry.Registry
import com.intellij.util.concurrency.annotations.RequiresEdt
import javax.swing.JComponent
@@ -38,6 +39,7 @@ class TurnView(
) : SessionLayoutPanel(SessionUiStyle.SessionLayout.GAP), Disposable, SessionEditorStyleTarget, SessionView {
private val messages = LinkedHashMap<String, MessageView>()
private var settled = true
override val sessionViewKind = SessionView.Kind.Default
@@ -48,6 +50,17 @@ class TurnView(
isOpaque = false
}
@RequiresEdt
fun setSettled(value: Boolean) {
if (settled == value) return
settled = value
revalidate()
}
override fun isValidateRoot(): Boolean {
return Registry.`is`("kilo.session.validateRoots", true) && settled
}
/** Add a new [MessageView] for [msg] at the end of this turn. */
fun addMessage(msg: Message): MessageView {
val view = MessageView(msg, openFile, style, openUrl, selection, openAttachment, resize, repo, hover, revert)
@@ -60,11 +73,17 @@ class TurnView(
/** Remove the [MessageView] for [msgId] if present. */
fun removeMessage(msgId: String) {
val view = messages.remove(msgId) ?: return
removeMessageChanged(msgId)
}
@RequiresEdt
fun removeMessageChanged(msgId: String): Boolean {
val view = messages.remove(msgId) ?: return false
remove(view)
Disposer.dispose(view)
syncCopyToolbars()
revalidate()
return true
}
@RequiresEdt
@@ -4,6 +4,7 @@ import ai.kilocode.client.session.SessionFileOpener
import ai.kilocode.client.session.views.base.GenericView
import ai.kilocode.client.session.views.base.PartView
import ai.kilocode.client.session.views.question.QuestionResultView
import ai.kilocode.client.session.views.tool.EditToolView
import ai.kilocode.client.session.views.tool.GlobToolView
import ai.kilocode.client.session.views.tool.ReadToolView
import ai.kilocode.client.session.views.tool.SearchToolView
@@ -54,6 +55,7 @@ object ViewFactory {
GlobToolView.canRender(content) -> GlobToolView(content, selection = selection, repo = repo)
SearchToolView.canRender(content) -> SearchToolView(content, selection = selection, repo = repo)
ReadToolView.canRender(content) -> ReadToolView(content, openFile, selection = selection)
EditToolView.canRender(content) -> EditToolView(content, openFile, selection = selection)
TaskToolView.canRender(content) -> TaskToolView(content, selection = selection)
else -> ToolView(content, selection = selection)
}
@@ -100,6 +102,8 @@ object ViewFactory {
if (view !is SearchToolView && SearchToolView.canRender(content)) return true
if (view is ReadToolView) return !ReadToolView.canRender(content) || QuestionResultView.canRender(content)
if (view is ToolView && ReadToolView.canRender(content)) return true
if (view is EditToolView) return !EditToolView.canRender(content) || QuestionResultView.canRender(content)
if (view is ToolView && EditToolView.canRender(content)) return true
if (view is TaskToolView) return !TaskToolView.canRender(content) || QuestionResultView.canRender(content)
if (view !is TaskToolView && TaskToolView.canRender(content)) return true
if (view is ToolView) return QuestionResultView.canRender(content)
@@ -0,0 +1,276 @@
package ai.kilocode.client.session.views.tool
import ai.kilocode.client.plugin.KiloBundle
import ai.kilocode.client.session.SessionFileOpener
import ai.kilocode.client.session.model.Content
import ai.kilocode.client.session.model.Tool
import ai.kilocode.client.session.model.ToolKind
import ai.kilocode.client.session.ui.popup.HeaderPopupBody
import ai.kilocode.client.session.ui.popup.HeaderPopupRequest
import ai.kilocode.client.session.ui.selection.SessionSelection
import ai.kilocode.client.session.ui.style.SessionEditorStyle
import ai.kilocode.client.session.ui.style.SessionUiStyle
import ai.kilocode.client.session.views.base.SecondarySessionPartView
import ai.kilocode.client.telemetry.Telemetry
import ai.kilocode.client.ui.DiffStatBadge
import ai.kilocode.client.ui.UiStyle
import ai.kilocode.client.ui.md.MdCodeBlockBorder
import ai.kilocode.client.ui.md.MdCodeBlockOptions
import com.intellij.openapi.actionSystem.DataSink
import com.intellij.openapi.actionSystem.UiDataProvider
import com.intellij.openapi.util.Disposer
import com.intellij.ui.EditorTextField
import com.intellij.ui.components.JBLabel
import com.intellij.util.concurrency.annotations.RequiresEdt
import com.intellij.util.ui.JBFont
import com.intellij.util.ui.JBUI
import java.awt.Dimension
import javax.swing.ScrollPaneConstants
/**
* Renders write tools (edit/write/apply_patch) with a Read-style header — an "Edit" title and a
* clickable file link — plus a diff-stat changes tag. The expandable body and the collapsed hover
* popup both render the unified diff via the shared markdown code editor, which colors it as a diff.
*/
class EditToolView(
tool: Tool,
private val openFile: SessionFileOpener = { _, _ -> },
private val selection: SessionSelection? = null,
private val parts: ToolParts = toolParts(tool, openFile),
private var body: EditBody = editBody(tool, selection, openFile),
) : SecondarySessionPartView(parts.header, { body.mount(tool) }), UiDataProvider {
override val contentId: String = tool.id
private var item = tool
private var style = SessionEditorStyle.current()
private var multi = editFiles(tool).size > 1
private val badge = DiffStatBadge(0, 0)
private val filesTag = JBLabel().apply {
foreground = UiStyle.Colors.weak()
font = JBFont.small()
border = JBUI.Borders.emptyRight(SessionUiStyle.View.Layout.HORIZONTAL_PADDING)
isVisible = false
}
init {
body.parent = this
parts.controls.add(filesTag)
parts.controls.add(badge)
bindHeader(parts.glyph, parts.title, parts.sub, parts.state, parts.center, parts.controls, parts.slot, filesTag, badge)
applyStyle(style)
sync()
}
override fun uiDataSnapshot(sink: DataSink) {
selection?.provideCopy(sink) { body.markdown() ?: diffMarkdown(item) }
}
@RequiresEdt
override fun expand(): Boolean {
val changed = super.expand()
if (!changed) return false
syncBody()
body.applyStyle(style)
return true
}
@RequiresEdt
override fun getPreferredSize(): Dimension {
val size = super.getPreferredSize()
if (!bodyVisible()) return size
val height = row.preferredSize.height + (body.panel()?.preferredSize?.height ?: 0)
return Dimension(size.width, minOf(size.height, height))
}
@RequiresEdt
override fun update(content: Content) {
if (content !is Tool) return
item = content
var changed = if (!expandable()) collapse() else false
changed = swapBody() || changed
changed = sync() || changed
changed = syncBody() || changed
if (changed) refresh()
}
/** Rebuild the body delegate when a streaming tool crosses the single/multi-file boundary. */
@RequiresEdt
private fun swapBody(): Boolean {
val next = editFiles(item).size > 1
if (next == multi) return false
multi = next
val expanded = isExpanded()
discardBody()
body.disposeBody()
body = editBody(item, selection, openFile).also { it.parent = this }
if (expanded) expand()
return true
}
@RequiresEdt
fun labelText(): String = listOf(parts.title.text, subtitleText(parts), parts.state.text)
.filter { it.isNotBlank() }
.joinToString(" ")
@RequiresEdt
fun bodyText(): String = editDiff(item)
@RequiresEdt
fun hasToggle(): Boolean = arrow.isVisible
@RequiresEdt
fun diffStat(): Pair<Int, Int> = diffStat(item)
@RequiresEdt
internal fun badgeVisible() = badge.isVisible
@RequiresEdt
internal fun filesTagVisible() = filesTag.isVisible
@RequiresEdt
internal fun filesTagText() = filesTag.text
@RequiresEdt
internal fun linkVisible() = parts.link.isVisible
@RequiresEdt
internal fun linkLabel() = parts.label
@RequiresEdt
internal fun linkHref() = parts.href
@RequiresEdt
internal fun linkTooltip() = parts.link.toolTipText
@RequiresEdt
internal fun openLink() = parts.openLink()
@RequiresEdt
internal fun bodyCreated() = body.created()
@RequiresEdt
internal fun bodyVisible() = body.attached(this)
@RequiresEdt
internal fun markdown() = body.markdown() ?: diffMarkdown(item)
@RequiresEdt
internal fun codeEditors(): List<EditorTextField> = body.codeEditors()
@RequiresEdt
override fun headerPopup(): HeaderPopupRequest? {
if (isExpanded()) return null
if (editDiff(item).isBlank()) return null
return HeaderPopupRequest(row, build = { buildPopupBody() }) {
Telemetry.send("Header Popup Shown", mapOf("surface" to "session", "tool" to "edit"))
}
}
@RequiresEdt
override fun applyStyle(style: SessionEditorStyle) {
this.style = style
var changed = false
changed = setFont(parts.title, style.boldEditorFont) || changed
changed = setFont(parts.sub, style.transcriptFont) || changed
changed = setFont(parts.link, style.transcriptFont) || changed
changed = setFont(parts.state, style.smallEditorFont) || changed
changed = body.applyStyle(style) || changed
if (changed) refresh()
}
private fun expandable(): Boolean =
editDiff(item).isNotBlank() || output(item).isNotBlank() || !item.error.isNullOrBlank()
private fun sync(): Boolean {
val expand = expandable()
var changed = false
changed = syncExpandable(expand) || changed
changed = setVisible(parts.state, !expand) || changed
changed = setIcon(parts.glyph, icon(item)) || changed
changed = setForeground(parts.glyph, color(item)) || changed
val count = editFiles(item).size
val titleText = if (count > 1) KiloBundle.message("session.part.tool.patch") else title(item)
changed = setText(parts.title, titleText) || changed
val path = if (count > 1) null else editPath(item)
changed = setFileTarget(parts, path, if (path == null) "" else tail(path)) || changed
changed = setForeground(parts.title, titleColor(item)) || changed
changed = setForeground(parts.link, UiStyle.Colors.fg()) || changed
changed = setText(parts.state, stateText(item)) || changed
changed = setForeground(parts.state, color(item)) || changed
changed = syncFilesTag(count) || changed
changed = syncBadge() || changed
return changed
}
private fun syncFilesTag(count: Int): Boolean {
val show = count > 1
var changed = setVisible(filesTag, show)
if (show) changed = setText(filesTag, KiloBundle.message("session.part.tool.edit.files", count)) || changed
return changed
}
private fun syncBadge(): Boolean {
val (added, removed) = diffStat(item)
val show = added > 0 || removed > 0
val changed = setVisible(badge, show)
if (show) badge.update(added, removed)
return changed
}
private fun syncBody(): Boolean = body.update(item)
@RequiresEdt
private fun buildPopupBody(): HeaderPopupBody {
val owner = Disposer.newDisposable("Edit popup body")
val popup = popupBody(item, selection, openFile).also { it.parent = owner }
// mount() already renders the current item (ToolMarkdownBody.mount calls update; PatchBody.mount
// calls rebuild and sets its signature), so a follow-up update() here would be a no-op.
val panel = popup.mount(item)
popup.applyStyle(style)
return HeaderPopupBody(panel, owner, style.editorBackground, SessionUiStyle.View.Popup.WIDE_MAX_WIDTH)
}
override fun dumpLabel() = "EditToolView#$contentId(${labelText()})"
companion object {
fun canRender(tool: Tool) = tool.kind == ToolKind.WRITE
}
}
/** Picks the multi-file patch body for apply_patch spanning several files, else the single diff. */
private fun editBody(tool: Tool, selection: SessionSelection?, openFile: SessionFileOpener): EditBody =
if (editFiles(tool).size > 1) PatchBody(selection, openFile) else diffBody(selection)
private fun popupBody(tool: Tool, selection: SessionSelection?, openFile: SessionFileOpener): EditBody =
if (editFiles(tool).size > 1) PatchBody(selection, openFile, POPUP_OPTS) else popupDiffBody(selection)
private fun diffBody(selection: SessionSelection?) = ToolMarkdownBody(
MdCodeBlockOptions(
border = MdCodeBlockBorder.Bottom,
maxLines = SessionUiStyle.View.Tool.DIFF_LINES,
verticalPolicy = ScrollPaneConstants.VERTICAL_SCROLLBAR_AS_NEEDED,
editorOnly = true,
),
selection,
render = ::diffMarkdown,
)
private fun popupDiffBody(selection: SessionSelection?) = ToolMarkdownBody(
POPUP_OPTS,
selection,
render = ::diffMarkdown,
)
private val POPUP_OPTS = MdCodeBlockOptions(
border = MdCodeBlockBorder.None,
verticalPolicy = ScrollPaneConstants.VERTICAL_SCROLLBAR_AS_NEEDED,
editorOnly = true,
)
/**
* Diff body markdown: per-file sections when an apply_patch touched multiple files, otherwise the
* single unified patch, falling back to the tool output/error when no diff is available.
*/
@RequiresEdt
internal fun diffMarkdown(tool: Tool): String {
val files = editFiles(tool)
if (files.count { it.patch.isNotBlank() } > 1) return multiFileDiffMarkdown(files)
val diff = editDiff(tool)
if (diff.isNotBlank()) return patchMarkdown(diff)
val body = plainBody(tool)
if (body.isBlank()) return ""
val fence = fence(body)
return buildString {
append(fence).append('\n')
append(body)
if (!body.endsWith('\n')) append('\n')
append(fence)
}
}
@@ -0,0 +1,191 @@
package ai.kilocode.client.session.views.tool
import ai.kilocode.client.session.SessionFileOpener
import ai.kilocode.client.session.model.Tool
import ai.kilocode.client.session.ui.selection.SessionSelection
import ai.kilocode.client.session.ui.style.SessionEditorStyle
import ai.kilocode.client.session.ui.style.SessionUiStyle
import ai.kilocode.client.ui.DiffStatBadge
import ai.kilocode.client.ui.UiStyle
import ai.kilocode.client.ui.layout.Stack
import ai.kilocode.client.ui.md.MdCodeBlockBorder
import ai.kilocode.client.ui.md.MdCodeBlockFactory
import ai.kilocode.client.ui.md.MdCodeBlockOptions
import ai.kilocode.client.ui.md.MdView
import ai.kilocode.client.ui.md.MdViewFactory
import com.intellij.openapi.Disposable
import com.intellij.openapi.util.Disposer
import com.intellij.ui.EditorTextField
import com.intellij.ui.components.JBScrollPane
import com.intellij.util.concurrency.annotations.RequiresEdt
import com.intellij.util.ui.JBUI
import java.awt.Component
import javax.swing.JComponent
import javax.swing.JPanel
import javax.swing.ScrollPaneConstants
/**
* Body surface shared by the single-file markdown diff ([ToolMarkdownBody]) and the multi-file
* apply_patch view ([PatchBody]), so [EditToolView] can hold either behind one type and swap between
* them when a streaming tool crosses the single/multi boundary.
*/
interface EditBody {
var parent: Disposable?
@RequiresEdt fun mount(tool: Tool): JComponent
@RequiresEdt fun created(): Boolean
@RequiresEdt fun panel(): JComponent?
@RequiresEdt fun attached(host: Component): Boolean
@RequiresEdt fun update(tool: Tool): Boolean
@RequiresEdt fun applyStyle(style: SessionEditorStyle): Boolean
@RequiresEdt fun markdown(): String?
@RequiresEdt fun codeEditors(): List<EditorTextField>
@RequiresEdt fun disposeBody()
}
/**
* Renders an apply_patch that touched several files as one section per file: a clickable filename
* link (same chrome as the Read/Edit header link) plus a per-file changes badge, left-aligned to the
* diff's own text inset, followed by that file's unified diff. Sections are rebuilt as a group when
* the underlying file set changes, matching the retained-Swing rebuild-on-add/remove convention.
*/
class PatchBody(
private val selection: SessionSelection?,
private val openFile: SessionFileOpener,
private val opts: MdCodeBlockOptions = DIFF_OPTS,
) : EditBody {
override var parent: Disposable? = null
private var root: Stack? = null
private var owner: Disposable? = null
private val views = mutableListOf<MdView>()
private val links = mutableListOf<FileLinkLabel>()
private var style = SessionEditorStyle.current()
private var signature = ""
@RequiresEdt
override fun mount(tool: Tool): JComponent {
root?.let { return it }
val panel = Stack.vertical()
root = panel
rebuild(tool)
return panel
}
@RequiresEdt
override fun created(): Boolean = root != null
@RequiresEdt
override fun panel(): JComponent? = root
@RequiresEdt
override fun attached(host: Component): Boolean = root?.parent === host
@RequiresEdt
override fun update(tool: Tool): Boolean {
if (root == null) return false
if (signatureOf(tool) == signature) return false
rebuild(tool)
return true
}
@RequiresEdt
override fun applyStyle(style: SessionEditorStyle): Boolean {
this.style = style
var changed = false
views.forEach { changed = applyMd(it) || changed }
links.forEach { if (it.font != style.transcriptFont) { it.font = style.transcriptFont; changed = true } }
return changed
}
@RequiresEdt
override fun markdown(): String? {
if (views.isEmpty()) return null
return views.joinToString("\n\n") { it.markdown() }
}
@RequiresEdt
override fun codeEditors(): List<EditorTextField> = views.flatMap { view ->
(view.component as? JPanel)?.components
?.filterIsInstance<JBScrollPane>()
?.mapNotNull { it.viewport.view as? EditorTextField }
?: emptyList()
}
@RequiresEdt
override fun disposeBody() {
val panel = root
owner?.let(Disposer::dispose)
owner = null
views.clear()
links.clear()
panel?.removeAll()
signature = ""
}
@RequiresEdt
private fun rebuild(tool: Tool) {
val panel = root ?: return
val parent = parent ?: error("Patch body has no parent")
disposeBody()
val disposable = Disposer.newDisposable("Patch body")
Disposer.register(parent, disposable)
owner = disposable
editFiles(tool).filter { it.patch.isNotBlank() }.forEachIndexed { index, file ->
if (index > 0) panel.gap(JBUI.scale(SessionUiStyle.View.Code.BLOCK_GAP))
panel.next(header(file))
panel.gap(UiStyle.Gap.sm())
val md = MdViewFactory.create(style, selection, MdCodeBlockFactory.default(opts))
Disposer.register(disposable, md)
applyMd(md)
md.set(patchMarkdown(file.patch))
views.add(md)
panel.next(md.component)
}
signature = signatureOf(tool)
panel.revalidate()
panel.repaint()
}
private fun signatureOf(tool: Tool): String = editFiles(tool)
.joinToString("\u0000") { "${it.path}\u0001${it.additions}\u0001${it.deletions}\u0001${it.patch}" }
@RequiresEdt
private fun header(file: EditFileChange): JComponent {
val link = FileLinkLabel(openFile).apply {
foreground = UiStyle.Colors.fg()
font = style.transcriptFont
setTarget(file.path, tail(file.path))
isVisible = true
}
links.add(link)
val row = Stack.horizontal(UiStyle.Gap.sm())
.next(link)
.next(DiffStatBadge(file.additions, file.deletions))
return JBUI.Panels.simplePanel(row).apply {
isOpaque = false
border = JBUI.Borders.emptyLeft(SessionUiStyle.View.Code.VIEWPORT_HORIZONTAL_PADDING)
}
}
private fun applyMd(md: MdView): Boolean {
val before = md.font
md.applyStyle(style)
md.font = style.editorFont
md.foreground = style.editorForeground
md.background = style.editorBackground
md.preBg = style.editorBackground
md.codeFont = style.editorFamily
md.component.border = JBUI.Borders.empty()
return before != md.font
}
private companion object {
val DIFF_OPTS = MdCodeBlockOptions(
border = MdCodeBlockBorder.Bottom,
maxLines = SessionUiStyle.View.Tool.DIFF_LINES,
verticalPolicy = ScrollPaneConstants.VERTICAL_SCROLLBAR_AS_NEEDED,
editorOnly = true,
)
}
}
@@ -96,6 +96,8 @@ class ReadToolView(
@RequiresEdt
internal fun linkHref() = parts.href
@RequiresEdt
internal fun linkTooltip() = parts.link.toolTipText
@RequiresEdt
internal fun openLink() = parts.openLink()
@RequiresEdt
@@ -129,24 +131,9 @@ class ReadToolView(
private fun syncSubtitle(): Boolean {
val target = target(item)?.takeIf { it.type == "file" }
if (target != null) {
var changed = false
if (parts.href != target.path) {
parts.href = target.path
changed = true
}
changed = setLinkText(parts, tail(target.path).ifBlank { target.path }) || changed
changed = show(parts, true) || changed
return changed
}
var changed = false
if (parts.href != null) {
parts.href = null
changed = true
}
if (target != null) return setFileTarget(parts, target.path, tail(target.path))
var changed = setFileTarget(parts, null, "")
changed = setText(parts.sub, subtitle(item)) || changed
changed = show(parts, false) || changed
return changed
}
@@ -14,12 +14,11 @@ import ai.kilocode.client.ui.UiStyle
import ai.kilocode.client.ui.md.MdCodeBlockBorder
import ai.kilocode.client.ui.md.MdCodeBlockFactory
import ai.kilocode.client.ui.md.MdCodeBlockOptions
import ai.kilocode.client.ui.md.MdView
import ai.kilocode.client.ui.md.MdViewFactory
import ai.kilocode.client.ui.md.hybrid.MdTerminal
import com.intellij.openapi.actionSystem.DataSink
import com.intellij.openapi.actionSystem.UiDataProvider
import com.intellij.openapi.Disposable
import com.intellij.openapi.util.Disposer
import com.intellij.ui.EditorTextField
import com.intellij.ui.components.JBHtmlPane
import com.intellij.ui.components.JBScrollPane
@@ -34,8 +33,8 @@ class ShellToolView(
tool: Tool,
private val selection: SessionSelection? = null,
private val parts: ToolParts = toolParts(tool),
private val holder: ShellHolder = ShellHolder(tool, selection),
) : SecondarySessionPartView(parts.header, { holder.body().panel }), UiDataProvider {
private val body: ToolMarkdownBody = shellBody(selection),
) : SecondarySessionPartView(parts.header, { body.mount(tool) }), UiDataProvider {
override val contentId: String = tool.id
@@ -43,14 +42,14 @@ class ShellToolView(
private var style = SessionEditorStyle.current()
init {
holder.parent = this
body.parent = this
bindHeader(parts.glyph, parts.title, parts.sub, parts.state, parts.center, parts.controls, parts.slot)
applyStyle(style)
sync()
}
override fun uiDataSnapshot(sink: DataSink) {
selection?.provideCopy(sink) { holder.shell?.markdown() ?: fallbackText() }
selection?.provideCopy(sink) { body.markdown() ?: fallbackText() }
}
private fun fallbackText() = ShellContent(item).body
@@ -60,7 +59,7 @@ class ShellToolView(
val changed = super.expand()
if (!changed) return false
syncBody()
holder.shell?.applyStyle(style)
body.applyStyle(style)
return true
}
@@ -68,7 +67,7 @@ class ShellToolView(
override fun getPreferredSize(): Dimension {
val size = super.getPreferredSize()
if (!bodyVisible()) return size
val height = row.preferredSize.height + (holder.shell?.panel?.preferredSize?.height ?: 0)
val height = row.preferredSize.height + (body.panel()?.preferredSize?.height ?: 0)
return Dimension(size.width, minOf(size.height, height))
}
@@ -105,16 +104,16 @@ class ShellToolView(
fun hasToggle(): Boolean = arrow.isVisible
@RequiresEdt
internal fun bodyCreated() = holder.shell != null
internal fun bodyCreated() = body.created()
@RequiresEdt
internal fun bodyVisible() = holder.shell?.panel?.parent === this
internal fun bodyVisible() = body.attached(this)
@RequiresEdt
internal fun markdown() = holder.shell?.markdown() ?: ShellContent(item).markdown
internal fun markdown() = body.markdown() ?: ShellContent(item).markdown
@RequiresEdt
internal fun codeEditors(): List<EditorTextField> = holder.shell?.codeEditors() ?: emptyList()
internal fun codeEditors(): List<EditorTextField> = body.codeEditors()
@RequiresEdt
internal fun commandFont() = codeEditors().firstOrNull()?.font ?: style.editorFont
@@ -138,10 +137,10 @@ class ShellToolView(
internal fun controlCount() = if (arrow.isVisible) 1 else 0
@RequiresEdt
internal fun mdComponent() = holder.shell?.mdComponent()
internal fun mdComponent() = body.panel()
@RequiresEdt
internal fun horizontalPolicy() = holder.shell?.scrolls()?.firstOrNull()?.horizontalScrollBarPolicy
internal fun horizontalPolicy() = body.scrolls().firstOrNull()?.horizontalScrollBarPolicy
?: ScrollPaneConstants.HORIZONTAL_SCROLLBAR_NEVER
@RequiresEdt
@@ -161,7 +160,7 @@ class ShellToolView(
changed = setFont(parts.sub, style.transcriptFont) || changed
changed = setFont(parts.link, style.smallEditorFont) || changed
changed = setFont(parts.state, style.smallEditorFont) || changed
holder.shell?.let { changed = it.applyStyle(style) || changed }
changed = body.applyStyle(style) || changed
if (changed) refresh()
}
@@ -181,10 +180,7 @@ class ShellToolView(
return changed
}
private fun syncBody(): Boolean {
val body = holder.shell ?: return false
return body.update(item)
}
private fun syncBody(): Boolean = body.update(item)
@RequiresEdt
private fun buildPopupBody(cmd: String): HeaderPopupBody {
@@ -208,7 +204,7 @@ class ShellToolView(
md.component.border = JBUI.Borders.empty()
md.set(popupMd(formatCommand(cmd)))
padPopup(md.component)
return HeaderPopupBody(md.component, md, style.editorBackground)
return HeaderPopupBody(md.component, md, style.editorBackground, SessionUiStyle.View.Popup.WIDE_MAX_WIDTH)
}
override fun dumpLabel() = "ShellToolView#$contentId(${labelText()})"
@@ -234,96 +230,28 @@ private fun padPopup(root: JComponent) {
private fun grow(size: Dimension, pad: Int) = Dimension(size.width, size.height + pad)
class ShellHolder(
private val tool: Tool,
private val selection: SessionSelection?,
) {
var parent: Disposable? = null
var shell: ShellBody? = null
private fun shellBody(selection: SessionSelection?) = ToolMarkdownBody(
MdCodeBlockOptions(
border = MdCodeBlockBorder.Bottom,
maxLines = 15,
verticalPolicy = ScrollPaneConstants.VERTICAL_SCROLLBAR_AS_NEEDED,
editorOnly = true,
),
selection,
render = { ShellContent(it).markdown },
font = SessionEditorStyle::transcriptFont,
chrome = ::styleShellHtml,
)
@RequiresEdt
fun body(): ShellBody {
val current = shell
if (current != null) return current
val owner = parent ?: error("Shell holder has no parent")
return ShellBody(tool, selection, owner).also {
shell = it
Disposer.register(owner, it)
}
/** Pads the left edge of shell section headers ("Command"/"Output") to line up with code text. */
@RequiresEdt
private fun styleShellHtml(md: MdView) {
val root = md.component as? JPanel ?: return
root.components.filterIsInstance<JBHtmlPane>().forEach {
it.border = JBUI.Borders.emptyLeft(SessionUiStyle.View.Code.VIEWPORT_HORIZONTAL_PADDING)
}
}
class ShellBody(
tool: Tool,
selection: SessionSelection?,
parent: Disposable,
) : Disposable {
private val md = MdViewFactory.create(
SessionEditorStyle.current(),
selection,
MdCodeBlockFactory.default(
MdCodeBlockOptions(
border = MdCodeBlockBorder.Bottom,
maxLines = 15,
verticalPolicy = ScrollPaneConstants.VERTICAL_SCROLLBAR_AS_NEEDED,
editorOnly = true,
),
),
)
val panel = md.component
init {
Disposer.register(parent, md)
applyStyle(SessionEditorStyle.current())
update(tool)
}
@RequiresEdt
fun update(tool: Tool): Boolean {
val content = ShellContent(tool)
if (md.markdown() == content.markdown) return false
md.set(content.markdown)
styleShell()
return true
}
@RequiresEdt
fun applyStyle(style: SessionEditorStyle): Boolean {
val before = md.font
md.applyStyle(style)
md.font = style.transcriptFont
md.foreground = style.editorForeground
md.background = style.editorBackground
md.preBg = style.editorBackground
md.codeFont = style.editorFamily
md.component.border = JBUI.Borders.empty()
styleShell()
return before != md.font
}
@RequiresEdt
private fun styleShell() {
val root = md.component as? JPanel ?: return
root.components.filterIsInstance<JBHtmlPane>().forEach {
it.border = JBUI.Borders.emptyLeft(SessionUiStyle.View.Code.VIEWPORT_HORIZONTAL_PADDING)
}
}
@RequiresEdt
fun markdown() = md.markdown()
@RequiresEdt
fun mdComponent() = md.component
@RequiresEdt
fun scrolls(): List<JBScrollPane> = (md.component as? JPanel)?.components?.filterIsInstance<JBScrollPane>() ?: emptyList()
@RequiresEdt
fun codeEditors(): List<EditorTextField> = scrolls().mapNotNull { it.viewport.view as? EditorTextField }
override fun dispose() = Unit
}
private data class ShellContent(
val command: String,
val output: String,
@@ -406,9 +334,4 @@ private fun StringBuilder.section(title: String, text: String, lang: String) {
append(fence)
}
private fun fence(text: String): String {
val size = Regex("`+").findAll(text).maxOfOrNull { it.value.length } ?: 0
return "`".repeat(maxOf(3, size + 1))
}
private fun clean(text: String): String = MdTerminal.strip(MdTerminal.reduce(text, keepSgr = false))
@@ -0,0 +1,102 @@
package ai.kilocode.client.session.views.tool
import ai.kilocode.client.session.model.Tool
import ai.kilocode.client.session.ui.selection.SessionSelection
import ai.kilocode.client.session.ui.style.SessionEditorStyle
import ai.kilocode.client.ui.md.MdCodeBlockFactory
import ai.kilocode.client.ui.md.MdCodeBlockOptions
import ai.kilocode.client.ui.md.MdView
import ai.kilocode.client.ui.md.MdViewFactory
import com.intellij.openapi.Disposable
import com.intellij.openapi.util.Disposer
import com.intellij.ui.EditorTextField
import com.intellij.ui.components.JBScrollPane
import com.intellij.util.concurrency.annotations.RequiresEdt
import com.intellij.util.ui.JBUI
import java.awt.Component
import java.awt.Font
import javax.swing.JComponent
import javax.swing.JPanel
/**
* A markdown-backed tool body (unified diff, shell transcript, ...) that is built lazily on first
* expansion and then mutated in place. Shared by [ShellToolView] and [EditToolView] so the
* lazy-init, styling, disposal, and editor-lookup logic lives in one place instead of being
* duplicated per tool.
*
* [render] turns the current [Tool] into the markdown to display, [font] picks the body font from
* the active style, and [chrome] applies any per-view tweaks after the markdown is (re)built.
*/
class ToolMarkdownBody(
private val opts: MdCodeBlockOptions,
private val selection: SessionSelection?,
private val render: (Tool) -> String,
private val font: (SessionEditorStyle) -> Font = SessionEditorStyle::editorFont,
private val chrome: (MdView) -> Unit = {},
) : EditBody {
override var parent: Disposable? = null
private var view: MdView? = null
/** Builds the body on first call, wiring it into [parent]'s disposable tree, then returns it. */
@RequiresEdt
override fun mount(tool: Tool): JComponent {
view?.let { return it.component }
val owner = parent ?: error("Tool markdown body has no parent")
val md = MdViewFactory.create(SessionEditorStyle.current(), selection, MdCodeBlockFactory.default(opts))
Disposer.register(owner, md)
view = md
applyStyle(SessionEditorStyle.current())
update(tool)
return md.component
}
@RequiresEdt
override fun created(): Boolean = view != null
@RequiresEdt
override fun panel(): JComponent? = view?.component
@RequiresEdt
override fun attached(host: Component): Boolean = view?.component?.parent === host
@RequiresEdt
override fun update(tool: Tool): Boolean {
val md = view ?: return false
val value = render(tool)
if (md.markdown() == value) return false
md.set(value)
chrome(md)
return true
}
@RequiresEdt
override fun applyStyle(style: SessionEditorStyle): Boolean {
val md = view ?: return false
val before = md.font
md.applyStyle(style)
md.font = font(style)
md.foreground = style.editorForeground
md.background = style.editorBackground
md.preBg = style.editorBackground
md.codeFont = style.editorFamily
md.component.border = JBUI.Borders.empty()
chrome(md)
return before != md.font
}
@RequiresEdt
override fun markdown(): String? = view?.markdown()
@RequiresEdt
fun scrolls(): List<JBScrollPane> =
(view?.component as? JPanel)?.components?.filterIsInstance<JBScrollPane>() ?: emptyList()
@RequiresEdt
override fun codeEditors(): List<EditorTextField> = scrolls().mapNotNull { it.viewport.view as? EditorTextField }
@RequiresEdt
override fun disposeBody() {
view?.let(Disposer::dispose)
view = null
}
}
@@ -6,6 +6,7 @@ import ai.kilocode.client.plugin.KiloBundle
import ai.kilocode.client.session.SessionFileOpener
import ai.kilocode.client.session.model.Tool
import ai.kilocode.client.session.model.ToolExecState
import ai.kilocode.client.session.model.ToolKind
import ai.kilocode.client.session.ui.selection.SessionSelection
import ai.kilocode.client.session.ui.selection.SessionCopyTarget
import ai.kilocode.client.session.ui.style.SessionEditorStyle
@@ -36,8 +37,14 @@ import com.intellij.ui.components.JBTextArea
import com.intellij.util.concurrency.annotations.RequiresEdt
import com.intellij.util.ui.JBUI
import com.intellij.xml.util.XmlStringUtil
import kotlinx.serialization.json.Json
import kotlinx.serialization.json.JsonArray
import kotlinx.serialization.json.JsonObject
import kotlinx.serialization.json.contentOrNull
import kotlinx.serialization.json.intOrNull
import kotlinx.serialization.json.jsonObject
import kotlinx.serialization.json.jsonPrimitive
import java.awt.BorderLayout
import java.awt.CardLayout
import java.awt.Color
import java.awt.Cursor
import java.awt.Dimension
@@ -59,18 +66,17 @@ class ToolParts(
val glyph: JBLabel,
val title: JBLabel,
val sub: JBLabel,
val link: JBLabel,
val link: FileLinkLabel,
val slot: JPanel,
val state: JBLabel,
val center: JPanel,
val controls: JComponent,
private val open: SessionFileOpener? = null,
val extra: JBLabel? = null,
val targets: List<JBLabel> = emptyList(),
private val mode: ToolBodyMode = ToolBodyMode.EDITOR,
) {
var href: String? = null
var label: String = ""
val href: String? get() = link.href
val label: String get() = link.label
private var body: ToolBody? = null
val text: JBTextArea?
@@ -93,8 +99,7 @@ class ToolParts(
@RequiresEdt
fun openLink(anchor: RelativePoint? = null) {
val value = href ?: return
open?.invoke(value, anchor)
link.openLink(anchor)
}
@RequiresEdt
@@ -109,6 +114,52 @@ class ToolParts(
}
}
class FileLinkLabel(
private val open: SessionFileOpener? = null,
) : JBLabel() {
var href: String? = null
private set
var label: String = ""
private set
init {
isVisible = false
isFocusable = false
foreground = UiStyle.Colors.fg()
cursor = Cursor.getPredefinedCursor(Cursor.HAND_CURSOR)
setRequestFocusEnabled(false)
addMouseListener(object : MouseAdapter() {
override fun mouseClicked(e: MouseEvent) {
openLink(RelativePoint(this@FileLinkLabel, Point(width / 2, height)))
}
})
}
@RequiresEdt
fun setTarget(path: String?, text: String): Boolean {
val next = single(text.ifBlank { path.orEmpty() })
val value = if (next.isBlank()) "" else XmlStringUtil.wrapInHtml("<nobr><u>${XmlStringUtil.escapeString(next)}</u></nobr>")
var changed = false
if (href != path) {
href = path
toolTipText = path
changed = true
}
if (label != next || this.text != value) {
label = next
this.text = value
changed = true
}
return changed
}
@RequiresEdt
fun openLink(anchor: RelativePoint? = null) {
val value = href ?: return
open?.invoke(value, anchor)
}
}
class ToolBody private constructor(
val area: JBTextArea?,
val ed: EditorTextField?,
@@ -345,36 +396,20 @@ private class ToolField(value: String, private var style: SessionEditorStyle, pr
}
}
private const val SUB_CARD = "sub"
private const val LINK_CARD = "link"
@RequiresEdt
internal fun toolParts(
tool: Tool,
openFile: SessionFileOpener? = null,
mode: ToolBodyMode = ToolBodyMode.TEXT,
): ToolParts {
lateinit var parts: ToolParts
val glyph = JBLabel()
val title = clip(JBLabel())
val sub = clip(JBLabel()).apply { foreground = UiStyle.Colors.weak() }
val link = clip(JBLabel()).apply {
isVisible = false
isFocusable = false
foreground = UiStyle.Colors.fg()
cursor = Cursor.getPredefinedCursor(Cursor.HAND_CURSOR)
setRequestFocusEnabled(false)
addMouseListener(object : MouseAdapter() {
override fun mouseClicked(e: MouseEvent) {
parts.openLink(RelativePoint(this@apply, Point(width / 2, 0)))
}
})
}
val slot = JPanel(CardLayout()).apply {
isOpaque = false
val link = clip(FileLinkLabel(openFile))
val slot = Stack.fitHorizontal().apply {
minimumSize = Dimension(0, minimumSize.height)
add(sub, SUB_CARD)
add(link, LINK_CARD)
next(sub)
next(link)
}
val state = clip(JBLabel()).apply { foreground = UiStyle.Colors.weak() }
val center = JPanel(BorderLayout(UiStyle.Gap.md(), 0)).apply {
@@ -390,7 +425,7 @@ internal fun toolParts(
add(center, BorderLayout.CENTER)
add(controls, BorderLayout.EAST)
}
parts = ToolParts(header, glyph, title, sub, link, slot, state, center, controls, openFile, mode = mode)
val parts = ToolParts(header, glyph, title, sub, link, slot, state, center, controls, mode = mode)
return parts.also {
controls.add(it.state)
}
@@ -406,12 +441,11 @@ internal fun searchParts(count: Int): ToolParts {
foreground = UiStyle.Colors.fg()
}
}
val link = clip(JBLabel()).apply { isVisible = false }
val slot = JPanel(CardLayout()).apply {
isOpaque = false
val link = clip(FileLinkLabel())
val slot = Stack.fitHorizontal().apply {
minimumSize = Dimension(0, minimumSize.height)
add(sub, SUB_CARD)
add(link, LINK_CARD)
next(sub)
next(link)
}
val state = clip(JBLabel()).apply { foreground = UiStyle.Colors.weak() }
val stack = Stack.fitHorizontal(UiStyle.Gap.md()).apply { targets.forEach { next(it) } }
@@ -449,9 +483,10 @@ internal fun icon(tool: Tool) = when (tool.name) {
else -> SessionViewIcons.mcp
}
internal fun title(tool: Tool) = when (tool.name) {
"read" -> KiloBundle.message("session.part.tool.read")
"bash" -> KiloBundle.message("session.part.tool.shell")
internal fun title(tool: Tool) = when {
tool.name == "read" -> KiloBundle.message("session.part.tool.read")
tool.name == "bash" -> KiloBundle.message("session.part.tool.shell")
tool.kind == ToolKind.WRITE -> KiloBundle.message("session.part.tool.edit")
else -> toolTitle(tool)
}
@@ -477,17 +512,18 @@ internal fun setTargetText(label: JBLabel, text: String): Boolean {
return true
}
/**
* Shows [path] as a clickable file link in the header slot, or clears the link when [path] is null.
* Shared by [ai.kilocode.client.session.views.tool.ReadToolView] and
* [ai.kilocode.client.session.views.tool.EditToolView] so both render file targets identically.
*/
@RequiresEdt
internal fun setLinkText(parts: ToolParts, text: String): Boolean {
val label = single(text)
val value = if (label.isBlank()) "" else XmlStringUtil.wrapInHtml("<nobr><u>${XmlStringUtil.escapeString(label)}</u></nobr>")
if (parts.label == label && parts.link.text == value) return false
parts.label = label
parts.link.text = value
return true
internal fun setFileTarget(parts: ToolParts, path: String?, label: String): Boolean {
val changed = parts.link.setTarget(path, label)
return show(parts, path != null) || changed
}
private fun clip(label: JBLabel): JBLabel = label.apply {
private fun <T : JBLabel> clip(label: T): T = label.apply {
minimumSize = Dimension(0, minimumSize.height)
}
@@ -504,9 +540,10 @@ private fun single(text: String): String = text.lineSequence()
@RequiresEdt
internal fun show(parts: ToolParts, link: Boolean): Boolean {
if (parts.link.isVisible == link && parts.sub.isVisible != link) return false
(parts.slot.layout as CardLayout).show(parts.slot, if (link) LINK_CARD else SUB_CARD)
return true
var changed = false
changed = setVisible(parts.link, link) || changed
changed = setVisible(parts.sub, !link) || changed
return changed
}
internal fun subtitleText(parts: ToolParts): String = if (parts.link.isVisible) parts.label else parts.sub.text
@@ -704,6 +741,161 @@ private fun toolSubtitle(tool: Tool): String {
return listOfNotNull(base).plus(args).joinToString(" ")
}
/** File path targeted by a write tool, preferring the most specific resolvable path. */
internal fun editPath(tool: Tool): String = editPaths(tool).maxWithOrNull(
compareBy<String>({ OSAgnosticPathUtil.isAbsolute(it) }, { depth(it) }),
) ?: tool.name
private fun editPaths(tool: Tool): List<String> {
val direct = listOf(tool.input["filePath"], tool.input["path"])
val diff = listOfNotNull(editFile(parseJsonObject(tool.metadata["filediff"])))
val files = parseJsonArray(tool.metadata["files"])?.mapNotNull { editFile(it.jsonObject) } ?: emptyList()
return (direct + diff + files + listOf(tool.title, tool.name))
.mapNotNull { it?.takeIf { value -> value.isNotBlank() } }
}
private fun editFile(obj: JsonObject?): String? = listOf("filePath", "path", "file", "relativePath")
.firstNotNullOfOrNull { key -> obj?.get(key)?.jsonPrimitive?.contentOrNull?.takeIf { it.isNotBlank() } }
private fun depth(path: String): Int = path.count { it == '/' || it == '\\' }
private val DIFF_JSON = Json { ignoreUnknownKeys = true; isLenient = true }
private fun parseJsonObject(raw: String?): JsonObject? =
raw?.takeIf { it.isNotBlank() }?.let { runCatching { DIFF_JSON.parseToJsonElement(it).jsonObject }.getOrNull() }
private fun parseJsonArray(raw: String?): JsonArray? =
raw?.takeIf { it.isNotBlank() }?.let { runCatching { DIFF_JSON.parseToJsonElement(it) as? JsonArray }.getOrNull() }
private fun patchOf(obj: JsonObject?): String? =
obj?.get("patch")?.jsonPrimitive?.contentOrNull?.takeIf { it.isNotBlank() }
/**
* Unified diff patch produced by a write tool, or empty when none is available. Kilo strips the raw
* `diff` field from stored parts (see stripPartMetadata) but keeps `filediff.patch` (edit/write) and
* per-file `files[].patch` (apply_patch) when under the size cap, so read those first.
*/
internal fun editDiff(tool: Tool): String {
tool.metadata["diff"]?.takeIf { it.isNotBlank() }?.let { return it }
patchOf(parseJsonObject(tool.metadata["filediff"]))?.let { return it }
parseJsonArray(tool.metadata["files"])?.let { files ->
val joined = files.mapNotNull { patchOf(it.jsonObject) }.joinToString("\n")
if (joined.isNotBlank()) return joined
}
return ""
}
/** One file touched by an apply_patch call, parsed from the tool's `files[]` metadata. */
internal data class EditFileChange(
val path: String,
val type: String,
val additions: Int,
val deletions: Int,
val patch: String,
)
/** Per-file changes from an apply_patch tool; empty for single-file edit/write tools (`filediff`). */
internal fun editFiles(tool: Tool): List<EditFileChange> =
parseJsonArray(tool.metadata["files"])?.mapNotNull { element ->
val obj = element.jsonObject
val path = editFile(obj) ?: return@mapNotNull null
EditFileChange(
path = path,
type = obj["type"]?.jsonPrimitive?.contentOrNull.orEmpty(),
additions = obj["additions"]?.jsonPrimitive?.intOrNull ?: 0,
deletions = obj["deletions"]?.jsonPrimitive?.intOrNull ?: 0,
patch = patchOf(obj).orEmpty(),
)
} ?: emptyList()
/**
* Sectioned markdown for a multi-file patch: each file gets a labeled header line (path plus its own
* add/remove counts) followed by its own fenced diff, so the joined apply_patch diff no longer runs
* together into one indistinguishable block. The path is wrapped in inline code so characters like
* underscores are not parsed as markdown emphasis.
*/
internal fun multiFileDiffMarkdown(files: List<EditFileChange>): String =
files.filter { it.patch.isNotBlank() }.joinToString("\n\n") { file ->
buildString {
append('`').append(tail(file.path)).append('`')
append(" +").append(file.additions).append(" -").append(file.deletions)
append("\n\n")
append(patchMarkdown(file.patch))
}
}
/** Added/removed line counts, preferring the counts computed by the CLI, else counting patch lines. */
internal fun diffStat(tool: Tool): Pair<Int, Int> {
parseJsonObject(tool.metadata["filediff"])?.let { fd ->
val add = fd["additions"]?.jsonPrimitive?.intOrNull
val del = fd["deletions"]?.jsonPrimitive?.intOrNull
if (add != null || del != null) return (add ?: 0) to (del ?: 0)
}
parseJsonArray(tool.metadata["files"])?.let { files ->
var add = 0
var del = 0
var found = false
files.forEach {
it.jsonObject["additions"]?.jsonPrimitive?.intOrNull?.let { v -> add += v; found = true }
it.jsonObject["deletions"]?.jsonPrimitive?.intOrNull?.let { v -> del += v; found = true }
}
if (found) return add to del
}
val patch = editDiff(tool)
if (patch.isBlank()) return 0 to 0
var added = 0
var removed = 0
for (line in patch.lineSequence()) {
when {
line.startsWith("+++") || line.startsWith("---") -> Unit
line.startsWith("+") -> added++
line.startsWith("-") -> removed++
}
}
return added to removed
}
/** Display-only diff body without VCS/file metadata headers (Index, diff --git, ---, +++, etc.). */
internal fun pureDiff(diff: String): String = diff.lineSequence()
.filterNot(::diffMeta)
.joinToString("\n")
.trim('\n')
private fun diffMeta(line: String): Boolean = line.startsWith("Index:") ||
line.startsWith("====") ||
line.startsWith("diff --git ") ||
line.startsWith("@@") ||
line.startsWith("index ") ||
line.startsWith("--- ") ||
line.startsWith("+++ ") ||
line.startsWith("new file mode ") ||
line.startsWith("deleted file mode ") ||
line.startsWith("old mode ") ||
line.startsWith("new mode ") ||
line.startsWith("similarity index ") ||
line.startsWith("dissimilarity index ") ||
line.startsWith("rename from ") ||
line.startsWith("rename to ") ||
line.startsWith("copy from ") ||
line.startsWith("copy to ")
/** Wraps a unified patch in a fenced `patch` block so the markdown code editor highlights it. */
internal fun patchMarkdown(diff: String): String = buildString {
// Fall back to the raw patch when stripping metadata leaves nothing (e.g. a pure rename or
// mode-only change with no +/-/context lines) so we never render an empty fenced block.
val body = pureDiff(diff).ifBlank { diff.trim('\n') }
val fence = fence(body)
append(fence).append("patch-pure\n")
append(body)
if (!body.endsWith('\n')) append('\n')
append(fence)
}
internal fun fence(text: String): String {
val size = Regex("`+").findAll(text).maxOfOrNull { it.value.length } ?: 0
return "`".repeat(maxOf(3, size + 1))
}
internal fun tail(path: String): String {
val value = path.trimEnd('/', '\\')
val index = maxOf(value.lastIndexOf('/'), value.lastIndexOf('\\'))
@@ -0,0 +1,90 @@
package ai.kilocode.client.ui.md.hybrid
import com.intellij.openapi.diff.DiffColors
import com.intellij.openapi.editor.DefaultLanguageHighlighterColors
import com.intellij.openapi.editor.colors.TextAttributesKey
import com.intellij.openapi.editor.ex.EditorEx
import com.intellij.openapi.editor.markup.HighlighterLayer
import com.intellij.openapi.editor.markup.HighlighterTargetArea
/**
* Overlays unified-diff coloring on a plain-text code editor: added lines get the theme's diff
* "inserted" background, removed lines the "deleted" background, hunk headers a keyword color, and
* file/index headers a dimmed comment color. Colors come from the active scheme via [DiffColors]
* and [DefaultLanguageHighlighterColors], so the result tracks the IDE theme like the diff viewer.
*/
internal object MdDiffHighlight {
data class Span(val key: TextAttributesKey, val area: HighlighterTargetArea)
data class Display(val text: String, val spans: List<Range>)
data class Range(val start: Int, val end: Int, val span: Span)
fun apply(editor: EditorEx, text: String) {
editor.markupModel.removeAllHighlighters()
val doc = editor.document
val size = doc.textLength
for (n in 0 until doc.lineCount) {
val start = doc.getLineStartOffset(n).coerceAtMost(size)
val end = doc.getLineEndOffset(n).coerceAtMost(size)
if (start >= end) continue
val span = classify(doc.charsSequence.subSequence(start, end).toString()) ?: continue
editor.markupModel.addRangeHighlighter(span.key, start, end, HighlighterLayer.SYNTAX + 1, span.area)
}
}
fun applyPure(editor: EditorEx, text: String) {
editor.markupModel.removeAllHighlighters()
val doc = editor.document
for (range in display(text).spans) {
val start = range.start.coerceAtMost(doc.textLength)
val end = range.end.coerceAtMost(doc.textLength)
if (start >= end) continue
editor.markupModel.addRangeHighlighter(range.span.key, start, end, HighlighterLayer.SYNTAX + 1, range.span.area)
}
}
fun display(text: String): Display {
val out = StringBuilder()
val ranges = mutableListOf<Range>()
text.lineSequence().forEachIndexed { i, line ->
if (i > 0) out.append('\n')
val span = classify(line)
val body = when {
line.startsWith("+") || line.startsWith("-") || line.startsWith(" ") -> line.drop(1)
else -> line
}
val start = out.length
out.append(body)
if (span != null) ranges.add(Range(start, out.length, span))
}
return Display(out.toString(), ranges)
}
private fun classify(line: String): Span? = when {
fileHeader(line) || meta(line) -> comment
line.startsWith("@@") -> hunk
line.startsWith("+") -> inserted
line.startsWith("-") -> deleted
else -> null
}
// Unified-diff file headers are the marker followed by a space (or the bare marker), e.g. "+++ b/f".
// Guarding on that shape keeps content lines like "++x;" (an inserted "+x;") from being dimmed.
private fun fileHeader(line: String): Boolean =
(line.startsWith("+++") || line.startsWith("---")) &&
(line.length == 3 || line[3] == ' ' || line[3] == '\t')
private fun meta(line: String): Boolean = line.startsWith("diff ") ||
line.startsWith("index ") ||
line.startsWith("Index:") ||
line.startsWith("===") ||
line.startsWith("new file") ||
line.startsWith("deleted file") ||
line.startsWith("rename ") ||
line.startsWith("similarity ") ||
line.startsWith("\\ No newline")
private val inserted = Span(DiffColors.DIFF_INSERTED, HighlighterTargetArea.LINES_IN_RANGE)
private val deleted = Span(DiffColors.DIFF_DELETED, HighlighterTargetArea.LINES_IN_RANGE)
private val hunk = Span(DefaultLanguageHighlighterColors.KEYWORD, HighlighterTargetArea.EXACT_RANGE)
private val comment = Span(DefaultLanguageHighlighterColors.LINE_COMMENT, HighlighterTargetArea.EXACT_RANGE)
}
@@ -6,7 +6,7 @@ import com.intellij.openapi.fileTypes.PlainTextFileType
import com.intellij.openapi.fileTypes.UnknownFileType
internal sealed class Kind {
data class Source(val file: FileType) : Kind()
data class Source(val file: FileType, val highlight: Highlight = Highlight.None) : Kind()
data class Terminal(val stream: Stream, val mode: Mode) : Kind()
}
@@ -14,6 +14,9 @@ internal enum class Stream { Stdout, Stderr }
internal enum class Mode { Ansi, Shell, Command }
/** Extra overlay highlighting applied on top of a source code block. */
internal enum class Highlight { None, Diff, DiffPure }
internal object MdLanguage {
/** Internal terminal fence tags produced by ShellToolView shell transcript markdown. */
private val terms = mapOf(
@@ -58,11 +61,16 @@ internal object MdLanguage {
"terraform" to "tf",
)
private val diffs = setOf("diff", "patch", "udiff")
private val pure = setOf("diff-pure", "patch-pure")
fun kind(lang: String?): Kind {
val key = lang?.trim()?.split(Regex("\\s+"))?.take(2)?.joinToString(" ")?.lowercase().orEmpty()
terms[key]?.let { return it }
if (key == "shell script") return Kind.Source(type("sh"))
val single = key.substringBefore(' ')
if (key in pure || single in pure) return Kind.Source(PlainTextFileType.INSTANCE, Highlight.DiffPure)
if (key in diffs || single in diffs) return Kind.Source(PlainTextFileType.INSTANCE, Highlight.Diff)
terms[single]?.let { return it }
files[key]?.let { return Kind.Source(type(it)) }
files[single]?.let { return Kind.Source(type(it)) }
@@ -314,7 +314,7 @@ internal open class MdViewHybrid(
is Desc.Html -> HtmlView(desc, htmlBlock(desc.body, disposable), disposable)
is Desc.Table -> TableView(desc, tableBlock(desc.body, disposable), disposable)
is Desc.Code -> when (val kind = desc.kind) {
is Kind.Source -> CodeView(desc, codeBlock(desc.text, kind.file, disposable), disposable)
is Kind.Source -> CodeView(desc, codeBlock(desc.text, kind, disposable), disposable)
is Kind.Terminal -> TermView(desc, terminalBlock(desc.text, kind, disposable), disposable)
}
}
@@ -331,27 +331,38 @@ internal open class MdViewHybrid(
customStyleSheetProvider { sheet() }
},
), UiDataProvider {
// A stationary pointer over scrolling content must keep this pane's hovered link and
// cursor fresh, so we replay a synthetic mouse move whenever the enclosing viewport
// scrolls. Only the pane under the pointer subscribes — otherwise every prose block in a
// large transcript would run a native pointer query + event dispatch on every scroll tick.
private var viewport: JViewport? = null
private var listening = false
private val scroll = ChangeListener { hover() }
private val pointer = object : java.awt.event.MouseAdapter() {
override fun mouseEntered(e: MouseEvent) = listen(true)
override fun mouseExited(e: MouseEvent) = listen(false)
}
private val hierarchy = java.awt.event.HierarchyListener { event ->
if (event.changeFlags and HierarchyEvent.PARENT_CHANGED.toLong() != 0L) attach()
if (event.changeFlags and HierarchyEvent.PARENT_CHANGED.toLong() != 0L) retarget()
}
init {
addMouseListener(pointer)
addHierarchyListener(hierarchy)
Disposer.register(disposable) {
viewport?.removeChangeListener(scroll)
listen(false)
removeMouseListener(pointer)
removeHierarchyListener(hierarchy)
}
}
override fun addNotify() {
super.addNotify()
attach()
retarget()
}
override fun removeNotify() {
viewport?.removeChangeListener(scroll)
listen(false)
viewport = null
super.removeNotify()
}
@@ -360,12 +371,20 @@ internal open class MdViewHybrid(
selection?.provideCopy(sink) { document.getText(0, document.length).trim() }
}
private fun attach() {
// Follow the enclosing viewport as this pane is reparented, keeping any live subscription.
private fun retarget() {
val next = SwingUtilities.getAncestorOfClass(JViewport::class.java, this) as? JViewport
if (viewport === next) return
viewport?.removeChangeListener(scroll)
if (listening) viewport?.removeChangeListener(scroll)
viewport = next
next?.addChangeListener(scroll)
if (listening) viewport?.addChangeListener(scroll)
}
// Track viewport scrolls only while the pointer is over this pane.
private fun listen(on: Boolean) {
if (listening == on) return
listening = on
if (on) viewport?.addChangeListener(scroll) else viewport?.removeChangeListener(scroll)
}
private fun hover() {
@@ -424,20 +443,20 @@ internal open class MdViewHybrid(
return pane
}
private fun codeBlock(text: String, file: FileType, disposable: Disposable): JBScrollPane {
private fun codeBlock(text: String, kind: Kind.Source, disposable: Disposable): JBScrollPane {
val opts = opts()
val value = text.trimEnd('\n')
val value = sourceText(text, kind)
val field = runCatching {
codeField(file, opts, text, false, disposable)
codeField(kind.file, opts, value, false, disposable)
}.getOrElse { err ->
LOG.warn("kind=markdown codeEditor=true failed message=${err.message}", err)
if (code.opts.editorOnly) runCatching {
codeField(PlainTextFileType.INSTANCE, opts, text, false, disposable)
codeField(PlainTextFileType.INSTANCE, opts, value, false, disposable)
}.getOrElse { fallback ->
LOG.warn("kind=markdown codeEditor=true fallback=plain failed message=${fallback.message}", fallback)
throw fallback
} else {
textArea(text, opts, disposable)
textArea(value, opts, disposable)
}
}
sizeCodeField(field, value)
@@ -451,6 +470,12 @@ internal open class MdViewHybrid(
return pane
}
private fun sourceText(text: String, kind: Kind.Source): String {
val value = text.trimEnd('\n')
if (kind.highlight == Highlight.DiffPure) return MdDiffHighlight.display(value).text
return value
}
private fun terminalBlock(text: String, kind: Kind.Terminal, disposable: Disposable): JBScrollPane {
val opts = opts()
val term = MdTerminal.decode(text, kind.stream)
@@ -823,12 +848,18 @@ internal open class MdViewHybrid(
private inner class CodeView(desc: Desc.Code, private val pane: JBScrollPane, disposable: Disposable) :
View(desc, pane, disposable) {
init {
overlay()
}
override fun compatible(desc: Desc) = desc is Desc.Code && (this.desc as Desc.Code).kind == desc.kind
override fun update(desc: Desc) {
if (this.desc == desc) return
this.desc = desc
val value = (desc as Desc.Code).text.trimEnd('\n')
val item = desc as Desc.Code
val kind = item.kind as? Kind.Source
val value = if (kind == null) item.text.trimEnd('\n') else sourceText(item.text, kind)
val view = pane.viewport.view
when (view) {
is CodeField -> view.text = value
@@ -838,6 +869,20 @@ internal open class MdViewHybrid(
sizeCodeField(view, value)
sizeCodePane(pane, view)
}
overlay()
}
/** Applies unified-diff coloring on top of a `diff`/`patch` block; a no-op otherwise. */
private fun overlay() {
val kind = (desc as Desc.Code).kind
if (kind !is Kind.Source || kind.highlight == Highlight.None) return
val field = pane.viewport.view as? CodeField ?: return
val editor = field.getEditor(true) ?: return
if (kind.highlight == Highlight.DiffPure) {
MdDiffHighlight.applyPure(editor, (desc as Desc.Code).text.trimEnd('\n'))
return
}
MdDiffHighlight.apply(editor, field.text)
}
override fun grow(delta: String) {
@@ -861,6 +906,7 @@ internal open class MdViewHybrid(
sizeCodeField(view, text)
sizeCodePane(pane, view)
}
overlay()
}
}
@@ -121,6 +121,11 @@
defaultValue="180000"
restartRequired="false"
overrides="false"/>
<registryKey key="kilo.session.validateRoots"
description="Treat settled (non-streaming) transcript turns as Swing validate roots so their internal repaints do not relayout the whole transcript."
defaultValue="true"
restartRequired="false"
overrides="false"/>
</extensions>
<applicationListeners>
@@ -126,6 +126,9 @@ session.part.tool.error=Error
session.part.tool.agent={0} Agent
session.part.tool.pending=Pending
session.part.tool.read=Read
session.part.tool.edit=Edit
session.part.tool.edit.files={0} files
session.part.tool.patch=Patch
session.part.tool.glob=Glob
session.part.tool.search=Search
session.part.tool.running=Running
@@ -57,6 +57,7 @@ session.part.tool.copy=نسخ
session.part.tool.error=خطأ
session.part.tool.pending=معلق
session.part.tool.read=قراءة
session.part.tool.edit=تحرير
session.part.tool.running=قيد التشغيل
session.part.tool.shell=Shell
session.part.tool.truncated=المخرجات مختصرة في المعاينة المسبقة. المخرجات الكاملة لا تزال في بيانات الجلسة.
@@ -57,6 +57,7 @@ session.part.tool.copy=Kopiraj
session.part.tool.error=Greška
session.part.tool.pending=Na čekanju
session.part.tool.read=Čita
session.part.tool.edit=Uredi
session.part.tool.running=Pokrenuto
session.part.tool.shell=Shell
session.part.tool.truncated=Izlaz skraćen u pregledu. Potpuni izlaz ostaje u podacima sesije.
@@ -57,6 +57,7 @@ session.part.tool.copy=Kopiér
session.part.tool.error=Fejl
session.part.tool.pending=Afventer
session.part.tool.read=Læs
session.part.tool.edit=Rediger
session.part.tool.running=Kører
session.part.tool.shell=Shell
session.part.tool.truncated=Output afkortet i forhåndsvisning. Fuldt output forbliver i sessionsdata.
@@ -57,6 +57,7 @@ session.part.tool.copy=Kopieren
session.part.tool.error=Fehler
session.part.tool.pending=Ausstehend
session.part.tool.read=Lesen
session.part.tool.edit=Bearbeiten
session.part.tool.running=Läuft
session.part.tool.shell=Shell
session.part.tool.truncated=Ausgabe in der Vorschau gekürzt. Vollständige Ausgabe verbleibt in den Sitzungsdaten.
@@ -57,6 +57,7 @@ session.part.tool.copy=Copiar
session.part.tool.error=Error
session.part.tool.pending=Pendiente
session.part.tool.read=Leer
session.part.tool.edit=Editar
session.part.tool.running=Ejecutando
session.part.tool.shell=Shell
session.part.tool.truncated=Salida truncada en la vista previa. La salida completa permanece en los datos de la sesión.
@@ -57,6 +57,7 @@ session.part.tool.copy=Copier
session.part.tool.error=Erreur
session.part.tool.pending=En attente
session.part.tool.read=Lire
session.part.tool.edit=Modifier
session.part.tool.running=En cours
session.part.tool.shell=Shell
session.part.tool.truncated=Sortie tronquée dans l'aperçu. La sortie complète reste dans les données de session.
@@ -57,6 +57,7 @@ session.part.tool.copy=コピー
session.part.tool.error=エラー
session.part.tool.pending=保留中
session.part.tool.read=読み取り
session.part.tool.edit=編集
session.part.tool.running=実行中
session.part.tool.shell=シェル
session.part.tool.truncated=プレビューでは出力が切り詰められています。完全な出力はセッションデータに残っています。
@@ -57,6 +57,7 @@ session.part.tool.copy=복사
session.part.tool.error=오류
session.part.tool.pending=대기 중
session.part.tool.read=읽기
session.part.tool.edit=편집
session.part.tool.running=실행 중
session.part.tool.shell=
session.part.tool.truncated=미리보기에서 출력이 잘렸습니다. 전체 출력은 세션 데이터에 남아 있습니다.
@@ -57,6 +57,7 @@ session.part.tool.copy=Kopiëren
session.part.tool.error=Fout
session.part.tool.pending=In afwachting
session.part.tool.read=Lezen
session.part.tool.edit=Bewerken
session.part.tool.running=Actief
session.part.tool.shell=Shell
session.part.tool.truncated=Uitvoer ingekort in voorvertoning. Volledige uitvoer blijft beschikbaar in sessiegegevens.
@@ -57,6 +57,7 @@ session.part.tool.copy=Kopier
session.part.tool.error=Feil
session.part.tool.pending=Venter
session.part.tool.read=Les
session.part.tool.edit=Rediger
session.part.tool.running=Kjører
session.part.tool.shell=Shell
session.part.tool.truncated=Utdata avkortet i forhåndsvisning. Fullstendig utdata finnes fortsatt i øktdata.
@@ -57,6 +57,7 @@ session.part.tool.copy=Kopiuj
session.part.tool.error=Błąd
session.part.tool.pending=Oczekuje
session.part.tool.read=Odczyt
session.part.tool.edit=Edycja
session.part.tool.running=Uruchomione
session.part.tool.shell=Powłoka
session.part.tool.truncated=Wyjście skrócone w podglądzie. Pełne wyjście pozostaje w danych sesji.
@@ -57,6 +57,7 @@ session.part.tool.copy=Copiar
session.part.tool.error=Erro
session.part.tool.pending=Pendente
session.part.tool.read=Ler
session.part.tool.edit=Editar
session.part.tool.running=Executando
session.part.tool.shell=Shell
session.part.tool.truncated=Saída truncada na pré-visualização. A saída completa permanece nos dados da sessão.
@@ -57,6 +57,7 @@ session.part.tool.copy=Копировать
session.part.tool.error=Ошибка
session.part.tool.pending=Ожидание
session.part.tool.read=Чтение
session.part.tool.edit=Редактирование
session.part.tool.running=Выполняется
session.part.tool.shell=Shell
session.part.tool.truncated=Вывод усечён в предпросмотре. Полный вывод сохраняется в данных сессии.
@@ -57,6 +57,7 @@ session.part.tool.copy=คัดลอก
session.part.tool.error=ข้อผิดพลาด
session.part.tool.pending=รอดำเนินการ
session.part.tool.read=อ่าน
session.part.tool.edit=แก้ไข
session.part.tool.running=กำลังทำงาน
session.part.tool.shell=Shell
session.part.tool.truncated=ผลลัพธ์ถูกตัดทอนในส่วนตัวอย่าง ผลลัพธ์ทั้งหมดยังคงอยู่ในข้อมูลเซสชัน
@@ -57,6 +57,7 @@ session.part.tool.copy=Kopyala
session.part.tool.error=Hata
session.part.tool.pending=Bekliyor
session.part.tool.read=Oku
session.part.tool.edit=Düzenle
session.part.tool.running=Çalışıyor
session.part.tool.shell=Kabuk
session.part.tool.truncated=Önizlemede çıktı kısaltıldı. Tam çıktı oturum verilerinde kalıyor.
@@ -57,6 +57,7 @@ session.part.tool.copy=Копіювати
session.part.tool.error=Помилка
session.part.tool.pending=Очікується
session.part.tool.read=Читання
session.part.tool.edit=Редагування
session.part.tool.running=Виконується
session.part.tool.shell=Shell
session.part.tool.truncated=Вивід у попередньому перегляді усічено. Повний вивід зберігається в даних сесії.
@@ -57,6 +57,7 @@ session.part.tool.copy=复制
session.part.tool.error=错误
session.part.tool.pending=待处理
session.part.tool.read=读取
session.part.tool.edit=编辑
session.part.tool.running=运行中
session.part.tool.shell=Shell
session.part.tool.truncated=预览中的输出已截断。完整输出仍保留在会话数据中。
@@ -57,6 +57,7 @@ session.part.tool.copy=複製
session.part.tool.error=錯誤
session.part.tool.pending=待處理
session.part.tool.read=讀取
session.part.tool.edit=編輯
session.part.tool.running=執行中
session.part.tool.shell=Shell
session.part.tool.truncated=預覽中的輸出已截斷。完整輸出仍保留在工作階段資料中。
@@ -7,6 +7,7 @@ import com.intellij.util.ui.JBUI
import com.intellij.util.ui.components.BorderLayoutPanel
import java.awt.Dimension
import java.awt.Insets
import javax.swing.JPanel
import javax.swing.JLabel
/**
@@ -281,6 +282,66 @@ class SessionLayoutTest : BasePlatformTestCase() {
assertEquals(20 + JBUI.scale(8), c2.y)
}
fun `test valid child reuses cached preferred height`() {
val p = panel(width = 300)
val child = probe(height = 20)
p.add(child)
p.doLayout()
child.markValid()
val count = child.count
p.doLayout()
assertEquals(count, child.count)
assertEquals(20, child.height)
}
fun `test invalid child is measured again`() {
val p = panel(width = 300)
val child = probe(height = 20)
p.add(child)
p.doLayout()
child.markValid()
val count = child.count
child.invalidate()
p.doLayout()
assertEquals(count + 1, child.count)
}
fun `test width change forces cached child remeasure`() {
val p = panel(width = 300)
val child = probe(height = 20)
p.add(child)
p.doLayout()
child.markValid()
val count = child.count
p.setSize(320, 2000)
p.doLayout()
assertEquals(count + 1, child.count)
assertEquals(320, child.width)
}
fun `test forget re-measures a valid child`() {
val p = panel(width = 300)
val child = probe(height = 20)
p.add(child)
p.doLayout()
child.markValid()
val count = child.count
// A settled turn is its own validate root, so it can be re-validated independently and its
// isValid flag flips back to true even after its content (and height) changed. forget()
// drops the stale cached height so the next layout pass re-measures the child.
(p.layout as SessionLayout).forget(child)
p.doLayout()
assertEquals(count + 1, child.count)
}
// ---- helpers ------
/** A fixed-height JLabel. The width is reported as 0 until layout sets it. */
@@ -293,4 +354,25 @@ class SessionLayoutTest : BasePlatformTestCase() {
override fun getPreferredSize(): Dimension = Dimension(0, height)
}
private fun probe(height: Int) = object : JPanel() {
var count = 0
private var valid = false
override fun isValid() = valid
override fun invalidate() {
valid = false
super.invalidate()
}
fun markValid() {
valid = true
}
override fun getPreferredSize(): Dimension {
count++
return Dimension(0, height)
}
}
}
@@ -39,6 +39,8 @@ import ai.kilocode.rpc.dto.TodoDto
import com.intellij.ide.ui.laf.darcula.ui.DarculaButtonUI
import com.intellij.openapi.Disposable
import com.intellij.openapi.util.Disposer
import com.intellij.openapi.util.registry.Registry
import com.intellij.openapi.util.registry.RegistryKeyDescriptor
import com.intellij.testFramework.fixtures.BasePlatformTestCase
import com.intellij.ui.components.ActionLink
import com.intellij.ui.components.JBLabel
@@ -53,7 +55,9 @@ import java.awt.Point
import java.awt.event.MouseEvent
import java.awt.image.BufferedImage
import javax.swing.JButton
import javax.swing.JComponent
import javax.swing.JPanel
import javax.swing.RepaintManager
import javax.swing.SwingUtilities
import javax.swing.border.Border
@@ -401,6 +405,151 @@ class SessionMessageListPanelTest : BasePlatformTestCase() {
assertEquals("hello world", tv.markdown())
}
fun `test empty ContentDelta does not refresh panel`() {
model.upsertMessage(msg("a1", "assistant"))
model.updateContent("a1", part("p1", "a1", "text", text = "hello"))
val mv = panel.findMessage("a1")!!
val tv = mv.part("p1") as TextView
val repaint = TrackingRepaintManager(setOf(panel, mv, tv))
val old = RepaintManager.currentManager(panel)
try {
RepaintManager.setCurrentManager(repaint)
model.appendDelta("a1", "p1", "")
assertEquals("hello", tv.markdown())
assertTrue(repaint.dirty.isEmpty())
assertTrue(repaint.invalid.isEmpty())
} finally {
RepaintManager.setCurrentManager(old)
}
}
fun `test identical ContentUpdated does not refresh panel`() {
model.upsertMessage(msg("a1", "assistant"))
model.updateContent("a1", part("p1", "a1", "text", text = "hello"))
val mv = panel.findMessage("a1")!!
val tv = mv.part("p1") as TextView
val comp = tv.md.component
val repaint = TrackingRepaintManager(setOf(panel, mv, tv))
val old = RepaintManager.currentManager(panel)
try {
RepaintManager.setCurrentManager(repaint)
model.updateContent("a1", part("p1", "a1", "text", text = "hello"))
assertSame(tv, mv.part("p1"))
assertSame(comp, tv.md.component)
assertTrue(repaint.dirty.isEmpty())
assertTrue(repaint.invalid.isEmpty())
} finally {
RepaintManager.setCurrentManager(old)
}
}
// ------ settled turns / validate roots (B) ------
fun `test turns are validate roots when idle`() {
model.upsertMessage(msg("u1", "user"))
model.upsertMessage(msg("a1", "assistant"))
model.upsertMessage(msg("u2", "user"))
assertTrue(panel.findTurn("u1")!!.isValidateRoot())
assertTrue(panel.findTurn("u2")!!.isValidateRoot())
}
fun `test streaming turn is not a validate root while busy`() {
model.upsertMessage(msg("u1", "user"))
model.upsertMessage(msg("a1", "assistant"))
model.upsertMessage(msg("u2", "user"))
model.upsertMessage(msg("a2", "assistant"))
model.setState(SessionState.Busy("thinking"))
assertTrue("prior turn stays a validate root", panel.findTurn("u1")!!.isValidateRoot())
assertFalse("streaming turn must not be a validate root", panel.findTurn("u2")!!.isValidateRoot())
}
fun `test turns settle again when idle`() {
model.upsertMessage(msg("u1", "user"))
model.upsertMessage(msg("u2", "user"))
model.setState(SessionState.Busy("thinking"))
model.setState(SessionState.Idle)
assertTrue(panel.findTurn("u1")!!.isValidateRoot())
assertTrue(panel.findTurn("u2")!!.isValidateRoot())
}
fun `test turn added while busy becomes the active non-root turn`() {
model.upsertMessage(msg("u1", "user"))
model.setState(SessionState.Busy("thinking"))
assertFalse(panel.findTurn("u1")!!.isValidateRoot())
model.upsertMessage(msg("u2", "user"))
assertTrue("previous turn settles once a newer turn is active", panel.findTurn("u1")!!.isValidateRoot())
assertFalse("newest turn is the active streaming turn", panel.findTurn("u2")!!.isValidateRoot())
}
fun `test validate roots flag disables turn isolation`() {
disableValidateRoots()
model.upsertMessage(msg("u1", "user"))
assertFalse(panel.findTurn("u1")!!.isValidateRoot())
}
fun `test settled turns still follow panel width top down`() {
model.upsertMessage(msg("a1", "assistant"))
model.updateContent("a1", part("p1", "a1", "text", text = "answer"))
val turn = panel.findTurn("a1")!!
assertTrue("idle turn is a validate root", turn.isValidateRoot())
panel.setSize(600, 2000)
layout(panel)
val wide = turn.width
panel.setSize(500, 2000)
layout(panel)
assertTrue("validate-root turns must still relayout top-down", turn.width < wide)
assertTrue(turn.isValidateRoot())
}
// ------ streaming stress / teardown ------
fun `test many streamed turns stay bounded and fully tear down`() {
val empty = count(panel)
repeat(40) { i ->
model.upsertMessage(msg("u$i", "user"))
model.updateContent("u$i", part("up$i", "u$i", "text", text = "q$i"))
model.upsertMessage(msg("a$i", "assistant"))
model.updateContent("a$i", part("ap$i", "a$i", "text", text = "```kotlin\nval x = $i\n```"))
repeat(20) { j -> model.appendDelta("a$i", "ap$i", " tok$j") }
}
assertEquals(40, panel.turnCount())
// Retained instances stay identical while streaming into an earlier message,
// and streaming deltas must not grow the component tree.
val tv = panel.findMessage("a0")!!.part("ap0") as TextView
val comp = tv.md.component
val count = count(panel)
repeat(50) { model.appendDelta("a0", "ap0", " x$it") }
assertSame(tv, panel.findMessage("a0")!!.part("ap0"))
assertSame(comp, tv.md.component)
assertEquals(count, count(panel))
model.clear()
assertEquals(0, panel.turnCount())
assertTrue("transcript turns must be removed on clear", panel.components.none { it is TurnView })
assertEquals("clear must return the transcript to its empty component tree", empty, count(panel))
}
fun `test ContentDelta preserves TextView and markdown component`() {
model.upsertMessage(msg("a1", "assistant"))
model.updateContent("a1", part("p1", "a1", "text", text = "first\n\nsecond"))
@@ -1152,6 +1301,18 @@ class SessionMessageListPanelTest : BasePlatformTestCase() {
for (child in root.components) if (child is Container) layout(child)
}
/** The plugin's `<registryKey>` extensions are not loaded in tests, so contribute the key here. */
private fun disableValidateRoots() {
val key = "kilo.session.validateRoots"
Registry.mutateContributedKeys {
it + (key to RegistryKeyDescriptor(key, "test", "true", false, false, null, null))
}
Disposer.register(testRootDisposable) {
Registry.mutateContributedKeys { it - key }
}
Registry.get(key).setValue(false, testRootDisposable)
}
private fun promptBox(root: MessageView): Component {
return components(root).first { it.parent != root && it is JPanel && it.componentCount == 1 && it.components.single() is TextView }
}
@@ -1175,4 +1336,19 @@ class SessionMessageListPanelTest : BasePlatformTestCase() {
.joinToString(" ")
}
}
private class TrackingRepaintManager(private val watched: Set<JComponent>) : RepaintManager() {
val dirty = mutableListOf<JComponent>()
val invalid = mutableListOf<JComponent>()
override fun addDirtyRegion(c: JComponent, x: Int, y: Int, w: Int, h: Int) {
if (c in watched) dirty.add(c)
super.addDirtyRegion(c, x, y, w, h)
}
override fun addInvalidComponent(invalidComponent: JComponent) {
if (invalidComponent in watched) invalid.add(invalidComponent)
super.addInvalidComponent(invalidComponent)
}
}
}
@@ -0,0 +1,476 @@
package ai.kilocode.client.session.views
import ai.kilocode.client.session.model.Tool
import ai.kilocode.client.session.model.ToolExecState
import ai.kilocode.client.session.model.toolKind
import ai.kilocode.client.session.ui.style.SessionUiStyle
import ai.kilocode.client.session.views.base.SecondarySessionPartView
import ai.kilocode.client.session.views.tool.EditToolView
import ai.kilocode.client.session.views.tool.ReadToolView
import ai.kilocode.client.session.views.tool.ToolView
import ai.kilocode.client.ui.DiffStatBadge
import com.intellij.openapi.diff.DiffColors
import com.intellij.openapi.editor.EditorFactory
import com.intellij.openapi.util.Disposer
import com.intellij.testFramework.fixtures.BasePlatformTestCase
import com.intellij.ui.components.JBLabel
import com.intellij.util.ui.JBUI
import com.intellij.util.ui.UIUtil
import kotlinx.serialization.json.addJsonObject
import kotlinx.serialization.json.buildJsonArray
import kotlinx.serialization.json.buildJsonObject
import kotlinx.serialization.json.put
import java.awt.Component
import java.awt.Container
import java.awt.event.MouseEvent
@Suppress("UnstableApiUsage")
class EditToolViewTest : BasePlatformTestCase() {
private val views = mutableListOf<EditToolView>()
override fun tearDown() {
views.forEach { Disposer.dispose(it) }
views.clear()
super.tearDown()
}
fun `test edit tool shows Edit title and clickable file link`() {
val opened = mutableListOf<String>()
val view = track(EditToolView(tool(), openFile = { href, _ -> opened.add(href) }))
val base: Any = view
assertTrue(base is SecondarySessionPartView)
assertTrue(view.labelText().contains("Edit"))
assertTrue(view.linkVisible())
assertEquals("App.kt", view.linkLabel())
assertEquals("/repo/src/App.kt", view.linkHref())
assertEquals("/repo/src/App.kt", view.linkTooltip())
assertTrue(view.labelText().contains("App.kt"))
view.openLink()
assertEquals(listOf("/repo/src/App.kt"), opened)
}
fun `test edit link uses metadata path when input is only filename`() {
val opened = mutableListOf<String>()
val path = "backend/src/com/kirillk/watcher/dao/GameApi.java"
val view = track(EditToolView(tool().also {
it.title = "GameApi.java"
it.input = mapOf("filePath" to "GameApi.java")
it.metadata = mapOf("filediff" to fileDiff(1, 0, PATCH, path))
}, openFile = { href, _ -> opened.add(href) }))
assertEquals("GameApi.java", view.linkLabel())
assertEquals(path, view.linkHref())
view.openLink()
assertEquals(listOf(path), opened)
}
fun `test changes tag shows additions and deletions`() {
val view = track(EditToolView(tool()))
assertTrue(view.badgeVisible())
assertEquals(2 to 1, view.diffStat())
}
fun `test changes tag hidden without diff`() {
val view = track(EditToolView(tool().also { it.metadata = emptyMap() }))
assertFalse(view.badgeVisible())
assertEquals(0 to 0, view.diffStat())
}
fun `test multi file apply_patch shows file count tag and aggregated changes`() {
val view = track(EditToolView(tool().also {
it.input = emptyMap()
it.metadata = mapOf("files" to filesMeta(
FileChange("src/A.kt", 2, 0, ADD_HUNK),
FileChange("src/B.kt", 1, 1, UPDATE_HUNK),
))
}))
assertTrue(view.labelText().contains("Patch"))
assertFalse(view.labelText().contains("Edit"))
assertTrue(view.filesTagVisible())
assertTrue(view.filesTagText()!!.contains("2 files"))
assertFalse(view.linkVisible())
assertTrue(view.badgeVisible())
assertEquals(3 to 1, view.diffStat())
}
fun `test multi file patch body renders a link and diff per file`() {
val opened = mutableListOf<String>()
val view = track(EditToolView(tool().also {
it.input = emptyMap()
it.metadata = mapOf("files" to filesMeta(
FileChange("src/A.kt", 2, 0, ADD_HUNK),
FileChange("pkg/B.kt", 1, 1, UPDATE_HUNK),
))
}, openFile = { href, _ -> opened.add(href) }))
view.toggle()
assertTrue(view.isExpanded())
assertEquals(2, view.codeEditors().size)
val fileLinks = labels(view).filter { it.text?.contains("<u>") == true }
assertTrue(fileLinks.any { it.text!!.contains("A.kt") && !it.text!!.contains("src/") })
assertTrue(fileLinks.any { it.text!!.contains("B.kt") && !it.text!!.contains("pkg/") })
assertTrue(fileLinks.any { it.text!!.contains("A.kt") && it.toolTipText == "src/A.kt" })
assertTrue(fileLinks.any { it.text!!.contains("B.kt") && it.toolTipText == "pkg/B.kt" })
// The per-file header renders one changes badge per file (plus the aggregate header badge).
assertEquals(3, badges(view).size)
click(fileLinks.first { it.text!!.contains("A.kt") }, 1)
assertEquals(listOf("src/A.kt"), opened)
}
fun `test single file apply_patch keeps link and hides count tag`() {
val view = track(EditToolView(tool().also {
it.input = emptyMap()
it.title = "src/Only.kt"
it.metadata = mapOf("files" to filesMeta(FileChange("src/Only.kt", 1, 1, UPDATE_HUNK)))
}))
assertFalse(view.filesTagVisible())
assertTrue(view.linkVisible())
assertEquals(1 to 1, view.diffStat())
assertFalse(view.markdown().contains("src/Only.kt"))
assertEquals(1, Regex("```patch-pure").findAll(view.markdown()).count())
}
fun `test edit body renders unified diff and expands`() {
val view = track(EditToolView(tool()))
assertTrue(view.hasToggle())
assertFalse(view.isExpanded())
assertFalse(view.bodyVisible())
assertTrue(view.markdown().contains("```patch-pure"))
assertTrue(view.markdown().contains("+new1"))
view.toggle()
assertTrue(view.isExpanded())
assertTrue(view.bodyVisible())
assertTrue(view.bodyCreated())
assertTrue(view.codeEditors().single().text.contains("new1"))
assertFalse(view.codeEditors().single().text.contains("+new1"))
assertFalse(view.codeEditors().single().text.contains("-old"))
}
fun `test edit body strips patch metadata headers`() {
// Relative-path headers so the `--- `/`+++ ` file-header assertions below actually exercise
// stripping: the header text (`--- src/App.kt`) shares its prefix with nothing in the body.
val patch = """
Index: src/App.kt
===================================================================
--- src/App.kt
+++ src/App.kt
@@ -1,2 +1,2 @@
keep
-old
+new
""".trimIndent()
val view = track(EditToolView(tool().also { it.metadata = mapOf("filediff" to fileDiff(1, 1, patch)) }))
assertFalse(view.markdown().contains("@@ -1,2 +1,2 @@"))
assertTrue(view.markdown().contains("-old"))
assertTrue(view.markdown().contains("+new"))
assertFalse(view.markdown().contains("Index:"))
assertFalse(view.markdown().contains("--- src/App.kt"))
assertFalse(view.markdown().contains("+++ src/App.kt"))
assertFalse(view.markdown().contains("===="))
view.toggle()
assertTrue(view.codeEditors().single().text.contains("old"))
assertTrue(view.codeEditors().single().text.contains("new"))
assertFalse(view.codeEditors().single().text.contains("-old"))
assertFalse(view.codeEditors().single().text.contains("+new"))
}
fun `test edit body colors added and removed diff lines`() {
val view = track(EditToolView(tool()))
view.toggle()
val editor = view.codeEditors().single().getEditor(true)!!
val chars = editor.document.charsSequence
val spans = editor.markupModel.allHighlighters.mapNotNull { h ->
val key = h.textAttributesKey ?: return@mapNotNull null
key to chars.subSequence(h.startOffset, h.endOffset).toString()
}
assertTrue(spans.any { it.first == DiffColors.DIFF_INSERTED && it.second.startsWith("new1") })
assertTrue(spans.any { it.first == DiffColors.DIFF_DELETED && it.second.startsWith("old") })
}
fun `test clicking link text opens file but empty slot toggles body`() {
val opened = mutableListOf<String>()
val view = track(EditToolView(tool(), openFile = { href, _ -> opened.add(href) }))
val link = linkLabel(view)
val slot = link.parent
click(slot, link.preferredSize.width + 50)
assertTrue(opened.isEmpty())
assertTrue(view.isExpanded())
click(link, 0)
assertEquals(listOf("/repo/src/App.kt"), opened)
}
fun `test metadata only patch falls back to raw text`() {
// A pure rename (no +/-/context lines) is entirely metadata: stripping it leaves nothing, so
// the raw patch must survive rather than render an empty fenced block.
val patch = """
diff --git a/src/Old.kt b/src/New.kt
similarity index 100%
rename from src/Old.kt
rename to src/New.kt
""".trimIndent()
val view = track(EditToolView(tool().also { it.metadata = mapOf("filediff" to fileDiff(0, 0, patch)) }))
assertTrue(view.markdown().contains("rename from src/Old.kt"))
assertTrue(view.markdown().contains("rename to src/New.kt"))
}
fun `test collapsed hover popup shows diff and none when expanded`() {
val view = track(EditToolView(tool()))
assertNotNull(view.headerPopup())
view.toggle()
assertNull(view.headerPopup())
}
fun `test edit header popup widens to diff content`() {
val patch = """
--- src/App.kt
+++ src/App.kt
@@ -1 +1 @@
-old
+${"x".repeat(180)}
""".trimIndent()
val view = track(EditToolView(tool().also {
it.metadata = mapOf("filediff" to fileDiff(1, 1, patch))
}))
val body = view.headerPopup()!!.build()
try {
assertTrue(body.component.preferredSize.width > JBUI.scale(SessionUiStyle.View.Popup.MAX_WIDTH))
assertTrue(body.component.preferredSize.width <= JBUI.scale(SessionUiStyle.View.Popup.WIDE_MAX_WIDTH))
} finally {
Disposer.dispose(body.disposable)
}
}
fun `test edit header popup stays narrow for short diff`() {
val patch = """
--- src/App.kt
+++ src/App.kt
@@ -1 +1 @@
-old
+new
""".trimIndent()
val view = track(EditToolView(tool().also {
it.metadata = mapOf("filediff" to fileDiff(1, 1, patch))
}))
val body = view.headerPopup()!!.build()
try {
assertTrue(body.component.preferredSize.width < JBUI.scale(SessionUiStyle.View.Popup.WIDE_MAX_WIDTH))
} finally {
Disposer.dispose(body.disposable)
}
}
fun `test multi file patch popup reuses patch body links`() {
val opened = mutableListOf<String>()
val view = track(EditToolView(tool().also {
it.input = emptyMap()
it.metadata = mapOf("files" to filesMeta(
FileChange("src/A.kt", 2, 0, ADD_HUNK),
FileChange("pkg/B.kt", 1, 1, UPDATE_HUNK),
))
}, openFile = { href, _ -> opened.add(href) }))
val body = view.headerPopup()!!.build()
try {
val fileLinks = labels(body.component).filter { it.text?.contains("<u>") == true }
assertTrue(fileLinks.any { it.text!!.contains("A.kt") && it.toolTipText == "src/A.kt" })
assertTrue(fileLinks.any { it.text!!.contains("B.kt") && it.toolTipText == "pkg/B.kt" })
click(fileLinks.first { it.text!!.contains("A.kt") }, 1)
assertEquals(listOf("src/A.kt"), opened)
} finally {
Disposer.dispose(body.disposable)
}
}
fun `test no hover popup without diff`() {
val view = track(EditToolView(tool().also { it.metadata = emptyMap() }))
assertNull(view.headerPopup())
}
fun `test view factory routes write tools to edit tool view`() {
assertTrue(ViewFactory.create(tool(), openFile = { _, _ -> }) is EditToolView)
assertTrue(ViewFactory.create(write("write"), openFile = { _, _ -> }) is EditToolView)
assertTrue(ViewFactory.create(write("apply_patch"), openFile = { _, _ -> }) is EditToolView)
}
fun `test canRender matches write kind tools only`() {
assertTrue(EditToolView.canRender(tool()))
assertTrue(EditToolView.canRender(write("write")))
assertFalse(EditToolView.canRender(Tool("p2", "read", toolKind("read"))))
assertFalse(EditToolView.canRender(Tool("p3", "bash", toolKind("bash"))))
}
fun `test shouldReplace swaps generic and edit views`() {
val edit = tool()
val other = Tool("p9", "mystery", toolKind("mystery")).also { it.state = ToolExecState.COMPLETED }
assertTrue(ViewFactory.shouldReplace(ToolView(edit), edit))
assertTrue(ViewFactory.shouldReplace(EditToolView(edit), other))
assertFalse(ViewFactory.shouldReplace(EditToolView(edit), edit))
}
fun `test edit editors are disposed after churn`() {
val base = EditorFactory.getInstance().allEditors.size
repeat(40) { i ->
val view = EditToolView(tool().also { it.metadata = mapOf("diff" to patch(i)) })
view.toggle()
view.codeEditors().forEach { it.getEditor(true) }
Disposer.dispose(view)
}
UIUtil.dispatchAllInvocationEvents()
assertEquals(base, EditorFactory.getInstance().allEditors.size)
}
fun `test multi file patch editors are disposed after churn`() {
val base = EditorFactory.getInstance().allEditors.size
repeat(20) { i ->
val view = EditToolView(tool().also {
it.input = emptyMap()
it.metadata = mapOf("files" to filesMeta(
FileChange("src/A$i.kt", 2, 0, ADD_HUNK),
FileChange("src/B$i.kt", 1, 1, UPDATE_HUNK),
))
})
view.toggle()
view.codeEditors().forEach { it.getEditor(true) }
Disposer.dispose(view)
}
UIUtil.dispatchAllInvocationEvents()
assertEquals(base, EditorFactory.getInstance().allEditors.size)
}
private fun track(view: EditToolView): EditToolView {
views.add(view)
return view
}
private fun click(component: Component, x: Int) {
component.dispatchEvent(MouseEvent(component, MouseEvent.MOUSE_CLICKED, System.currentTimeMillis(), 0, x, 1, 1, false))
}
private fun linkLabel(view: EditToolView): JBLabel =
labels(view).first { it.text?.contains("<u>") == true }
private fun labels(root: Container): List<JBLabel> = root.components.flatMap { child ->
val nested = if (child is Container) labels(child) else emptyList()
if (child is JBLabel) nested + child else nested
}
private fun badges(root: Container): List<DiffStatBadge> = root.components.flatMap { child ->
val nested = if (child is Container) badges(child) else emptyList()
if (child is DiffStatBadge) nested + child else nested
}
private fun tool() = Tool("p1", "edit", toolKind("edit")).also {
it.state = ToolExecState.COMPLETED
it.title = "src/App.kt"
it.input = mapOf("filePath" to "/repo/src/App.kt")
it.output = "Edit applied successfully."
it.metadata = mapOf("filediff" to fileDiff(2, 1, PATCH))
}
private fun write(name: String) = Tool("p1", name, toolKind(name)).also {
it.state = ToolExecState.COMPLETED
it.input = mapOf("filePath" to "/repo/src/App.kt")
it.metadata = mapOf("filediff" to fileDiff(2, 1, PATCH))
}
private fun patch(i: Int) = """
--- src/App.kt
+++ src/App.kt
@@ -1,2 +1,2 @@
line$i
-old$i
+new$i
""".trimIndent()
private data class FileChange(val path: String, val additions: Int, val deletions: Int, val patch: String)
// Mirrors how the CLI serializes metadata.files (a JsonArray of per-file changes rendered to string).
private fun filesMeta(vararg files: FileChange): String = buildJsonArray {
files.forEach { file ->
addJsonObject {
put("relativePath", file.path)
put("type", "update")
put("additions", file.additions)
put("deletions", file.deletions)
put("patch", file.patch)
}
}
}.toString()
// Mirrors how the CLI serializes metadata.filediff (a JsonObject rendered to string).
private fun fileDiff(
additions: Int,
deletions: Int,
patch: String,
path: String = "src/App.kt",
): String = buildJsonObject {
put("file", path)
put("additions", additions)
put("deletions", deletions)
put("patch", patch)
}.toString()
companion object {
private val PATCH = """
--- src/App.kt
+++ src/App.kt
@@ -1,3 +1,4 @@
line1
-old
+new1
+new2
line3
""".trimIndent()
private val ADD_HUNK = """
@@ -0,0 +1,2 @@
+alpha
+beta
""".trimIndent()
private val UPDATE_HUNK = """
@@ -1,2 +1,2 @@
keep
-old
+new
""".trimIndent()
}
}
@@ -51,6 +51,7 @@ class ReadToolViewTest : BasePlatformTestCase() {
assertTrue(view.linkVisible())
assertEquals("SessionUiLayoutTest.kt", view.linkText())
assertEquals(path, view.linkHref())
assertEquals(path, view.linkTooltip())
assertTrue(view.linkMarkup().contains("<nobr><u>SessionUiLayoutTest.kt</u></nobr>"))
assertEquals(UiStyle.Colors.fg().rgb, view.linkForeground().rgb)
assertEquals(view.linkFont(), view.bodyFont())
@@ -75,6 +76,7 @@ class ReadToolViewTest : BasePlatformTestCase() {
assertFalse(view.linkVisible())
assertNull(view.linkHref())
assertNull(view.linkTooltip())
assertEquals(UiStyle.Colors.fg().rgb, view.subtitleForeground().rgb)
assertEquals(view.subtitleFont(), view.bodyFont())
assertTrue(view.labelText().contains(path))
@@ -290,8 +290,8 @@ class ReasoningViewTest : BasePlatformTestCase() {
val panel = scroll.viewport.view as JPanel
assertEquals(1, panel.components.filterIsInstance<JComponent>().size)
assertTrue(body.component.preferredSize.width in 1..JBUI.scale(350))
assertEquals(JBUI.scale(450), body.component.preferredSize.height)
assertTrue(body.component.preferredSize.width in 1..JBUI.scale(SessionUiStyle.View.Popup.MAX_WIDTH))
assertEquals(JBUI.scale(SessionUiStyle.View.Popup.MAX_HEIGHT), body.component.preferredSize.height)
} finally {
Disposer.dispose(body.disposable)
}
@@ -416,9 +416,9 @@ class ShellToolViewTest : BasePlatformTestCase() {
assertTrue(field.preferredSize.height - border.top >= editor.lineHeight * lines)
assertTrue(field.minimumSize.height - border.top >= editor.lineHeight * lines)
assertTrue(pane.preferredSize.height >= field.preferredSize.height + pad.top + pad.bottom)
assertTrue(body.component.preferredSize.width in 1..JBUI.scale(350))
assertTrue(body.component.preferredSize.width in 1..JBUI.scale(SessionUiStyle.View.Popup.WIDE_MAX_WIDTH))
assertTrue(body.component.preferredSize.height > 0)
assertTrue(body.component.preferredSize.height <= JBUI.scale(450))
assertTrue(body.component.preferredSize.height <= JBUI.scale(SessionUiStyle.View.Popup.MAX_HEIGHT))
} finally {
Disposer.dispose(body.disposable)
}
@@ -427,6 +427,33 @@ class ShellToolViewTest : BasePlatformTestCase() {
assertEquals(base, EditorFactory.getInstance().allEditors.size)
}
fun `test shell header popup widens to command content`() {
val view = track(ShellToolView(tool().also {
it.input = mapOf("command" to "echo ${"x".repeat(180)}")
}))
val body = view.headerPopup()!!.build()
try {
assertTrue(body.component.preferredSize.width > JBUI.scale(SessionUiStyle.View.Popup.MAX_WIDTH))
assertTrue(body.component.preferredSize.width <= JBUI.scale(SessionUiStyle.View.Popup.WIDE_MAX_WIDTH))
} finally {
Disposer.dispose(body.disposable)
}
}
fun `test shell header popup stays narrow for short command`() {
val view = track(ShellToolView(tool().also {
it.input = mapOf("command" to "ls")
}))
val body = view.headerPopup()!!.build()
try {
assertTrue(body.component.preferredSize.width < JBUI.scale(SessionUiStyle.View.Popup.WIDE_MAX_WIDTH))
} finally {
Disposer.dispose(body.disposable)
}
}
fun `test shell header popup breaks chained operators outside quotes`() {
val view = track(ShellToolView(tool().also {
it.input = mapOf(
@@ -3,6 +3,7 @@ package ai.kilocode.client.session.views
import ai.kilocode.client.session.model.Tool
import ai.kilocode.client.session.model.ToolExecState
import ai.kilocode.client.session.model.toolKind
import ai.kilocode.client.session.views.tool.EditToolView
import ai.kilocode.client.session.views.tool.GlobToolView
import ai.kilocode.client.session.views.tool.SearchToolView
import ai.kilocode.client.session.views.tool.ShellToolView
@@ -11,6 +12,8 @@ import com.intellij.openapi.editor.EditorFactory
import com.intellij.openapi.util.Disposer
import com.intellij.testFramework.fixtures.BasePlatformTestCase
import com.intellij.util.ui.UIUtil
import kotlinx.serialization.json.buildJsonObject
import kotlinx.serialization.json.put
@Suppress("UnstableApiUsage")
class ToolBodyStressTest : BasePlatformTestCase() {
@@ -62,11 +65,46 @@ class ToolBodyStressTest : BasePlatformTestCase() {
assertEquals(base, EditorFactory.getInstance().allEditors.size)
}
fun `test expanded edit tool editors are disposed after churn`() {
val base = EditorFactory.getInstance().allEditors.size
repeat(60) { i ->
val view = EditToolView(edit(i))
view.toggle()
view.codeEditors().forEach { it.getEditor(true) }
Disposer.dispose(view)
}
drainEdt()
assertEquals(base, EditorFactory.getInstance().allEditors.size)
}
private fun tool(index: Int) = Tool("p$index", "mystery", toolKind("mystery")).also {
it.state = ToolExecState.COMPLETED
it.output = (1..20).joinToString("\n") { line -> "line $index/$line" }
}
private fun edit(index: Int) = Tool("e$index", "edit", toolKind("edit")).also {
it.state = ToolExecState.COMPLETED
it.input = mapOf("filePath" to "/repo/src/File$index.kt")
val patch = buildString {
append("--- src/File$index.kt\n")
append("+++ src/File$index.kt\n")
append("@@ -1,3 +1,4 @@\n")
append(" line1\n")
append("-old$index\n")
append("+new$index\n")
}
it.metadata = mapOf(
"filediff" to buildJsonObject {
put("file", "src/File$index.kt")
put("additions", 1)
put("deletions", 1)
put("patch", patch)
}.toString(),
)
}
private fun shell(index: Int) = Tool("p$index", "bash", toolKind("bash")).also {
it.state = ToolExecState.COMPLETED
it.input = mapOf("command" to "log $index")
@@ -273,6 +273,8 @@ class MdViewHybridTest : BasePlatformTestCase() {
assertTrue(iter.isValid)
val rect = pane.modelToView2D(iter.startOffset)!!.bounds
// Real AWT delivers MOUSE_ENTERED before MOUSE_MOVED; the enter arms scroll tracking.
pane.dispatchEvent(MouseEvent(pane, MouseEvent.MOUSE_ENTERED, System.currentTimeMillis(), 0, rect.x + 1, rect.y + rect.height / 2, 0, false, MouseEvent.NOBUTTON))
pane.dispatchEvent(MouseEvent(pane, MouseEvent.MOUSE_MOVED, System.currentTimeMillis(), 0, rect.x + 1, rect.y + rect.height / 2, 0, false, MouseEvent.NOBUTTON))
host.viewport.viewPosition = Point(0, 32)
drainEdt()
@@ -281,6 +283,24 @@ class MdViewHybridTest : BasePlatformTestCase() {
assertTrue(events.contains(HyperlinkEvent.EventType.EXITED))
}
fun `test prose pane tracks viewport scrolls only while hovered`() {
view.set("See [docs](https://example.com)\n\n" + (1..20).joinToString("\n") { "line $it" })
val pane = htmls().single()
val host = JBScrollPane(view.component)
host.setSize(420, 64)
view.component.setSize(420, view.component.preferredSize.height)
host.doLayout()
view.component.doLayout()
drainEdt()
val base = host.viewport.changeListeners.size
pane.dispatchEvent(MouseEvent(pane, MouseEvent.MOUSE_ENTERED, System.currentTimeMillis(), 0, 1, 1, 0, false, MouseEvent.NOBUTTON))
assertEquals("hovered prose pane must follow viewport scrolls", base + 1, host.viewport.changeListeners.size)
pane.dispatchEvent(MouseEvent(pane, MouseEvent.MOUSE_EXITED, System.currentTimeMillis(), 0, -1, -1, 0, false, MouseEvent.NOBUTTON))
assertEquals("pane must stop following scrolls once the pointer leaves", base, host.viewport.changeListeners.size)
}
fun `test file ref links include line suffix and exclude punctuation`() {
view.set("See kilocode/session/prompt.ts:302, native-plan-prompt.txt:37-38.")
val html = view.html()
@@ -0,0 +1,31 @@
package ai.kilocode.client.ui.md.hybrid
import com.intellij.openapi.diff.DiffColors
import com.intellij.openapi.editor.DefaultLanguageHighlighterColors
import com.intellij.testFramework.fixtures.BasePlatformTestCase
class MdDiffHighlightTest : BasePlatformTestCase() {
fun `test inserted line whose content starts with plus plus is not dimmed as a header`() {
// "++x;" is an inserted line ("+" marker + "+x;" content), not a "+++" file header.
val out = MdDiffHighlight.display("++x;")
assertEquals(1, out.spans.size)
assertEquals(DiffColors.DIFF_INSERTED, out.spans.single().span.key)
}
fun `test deleted line whose content starts with a dash is not dimmed as a header`() {
// "--x" is a deleted line ("-" marker + "-x" content), not a "---" file header.
val out = MdDiffHighlight.display("--x")
assertEquals(1, out.spans.size)
assertEquals(DiffColors.DIFF_DELETED, out.spans.single().span.key)
}
fun `test real file headers are dimmed as comments`() {
val out = MdDiffHighlight.display("--- a/File.kt\n+++ b/File.kt")
assertEquals(2, out.spans.size)
assertTrue(out.spans.all { it.span.key == DefaultLanguageHighlighterColors.LINE_COMMENT })
}
}
+32
View File
@@ -4,6 +4,31 @@ import { bootstrap } from "../bootstrap"
import { KiloSessions } from "@/kilo-sessions/kilo-sessions"
import { context } from "@/project/instance-context"
import { InstanceRuntime } from "@/project/instance-runtime"
import { Instance } from "@/kilocode/instance"
import { InstallationVersion } from "@opencode-ai/core/installation/version"
import os from "node:os"
import path from "node:path"
function truncate(value: string, max: number) {
return value.length > max ? value.slice(0, max) : value
}
// kilocode_change start - K1 W1: extracted so the advertisement payload shape
// is unit-testable as real behavior, rather than only through a source-text/
// regex assertion on this file (the handler itself can't be driven end-to-end
// — see the doc comment on `handler` below).
export function buildInstanceAdvertisement(directory: string): {
name: string
projectName: string
version: string
} {
return {
name: truncate(os.hostname(), 64),
projectName: truncate(path.basename(directory) || directory, 64),
version: truncate(InstallationVersion, 32),
}
}
// kilocode_change end
export const RemoteCommand = cmd({
command: "remote",
@@ -11,6 +36,13 @@ export const RemoteCommand = cmd({
builder: (yargs) => yargs,
handler: async () => {
await bootstrap(process.cwd(), async () => {
// kilocode_change - K1 W1: advertise this instance on the relay
// heartbeat so the cloud side can show it as a spawn-capable instance.
// The process-wide `KILO_REMOTE_ATTACH_SESSION` guard was removed in K1
// (in-process sessions only; no spawned children), so this is always
// advertised for the explicit `kilo remote` command path.
KiloSessions.setInstanceAdvertisement(buildInstanceAdvertisement(Instance.directory))
await KiloSessions.enableRemote()
console.log("Remote connection enabled.")
@@ -35,8 +35,13 @@ export namespace AttachedState {
* `opts.requireSessionId` is forwarded by `announce(id)` so the relay
* only resolves the attach promise when a fresh heartbeat whose
* payload contains that id was actually sent. Presence fire-and-
* forget heartbeats call without an id and resolve on any fresh send. */
heartbeat: (opts?: { requireSessionId?: string }) => Promise<void>
* forget heartbeats call without an id and resolve on any fresh send.
*
* `opts.detachSessionId` is forwarded by `detach(id)` so the relay
* only resolves the detach promise when a fresh heartbeat whose
* payload DOES NOT contain that id was actually sent (the negative-
* containment fence). */
heartbeat: (opts?: { requireSessionId?: string; detachSessionId?: string }) => Promise<void>
log?: { warn: (msg: string, meta?: unknown) => void }
}
@@ -54,11 +59,34 @@ export namespace AttachedState {
* case presence is authoritative and the attach resolves successfully.
* On success advances `lastSentKey` to the current union. */
announce(id: string): Promise<void>
/**
* Awaitable session-detach. Removes the id from BOTH the presence and
* pending sets and awaits a fresh heartbeat whose payload no longer
* contains the id (id-containment fence so a stale "still contains"
* cycle cannot falsely report the detach as complete).
*
* On heartbeat failure: rolls back by restoring the prior ownership
* (presence add + pending add as appropriate), re-throws, and the
* caller is responsible for NOT sending the success response so the
* CLI can keep the session attached and the process alive.
*
* The id is also added to a suppression tombstone: until presence
* itself stops reporting the id, subsequent `setPresence` calls
* will NOT re-adopt it (this prevents an immediately-following
* presence replacement from instantly re-attaching a session that
* the remote just exited). The tombstone is released the moment
* `setPresence` receives a set that does not contain the id (i.e.
* presence has genuinely dropped it), so a later real reopen
* (a fresh announce after a legitimate re-open) is not blocked.
*/
detach(id: string): Promise<void>
/** Current union of presence pending for the next heartbeat payload. */
union(): ReadonlySet<string>
/** True iff the id is in either the presence or pending set. */
has(id: string): boolean
/** Clear both sets across a connection lifecycle. The next setPresence
* call after reset will fire a heartbeat because the baseline key is
* empty. */
* empty. Also clears the suppressions. */
reset(): void
}
@@ -78,6 +106,13 @@ export namespace AttachedState {
export function create(options: Options): Interface {
const presence = new Set<string>()
const pending = new Set<string>()
// kilocode_change - K1 W1: tombstones for ids that have been remotely
// detached but are still being reported by presence. While an id is in
// this set, setPresence MUST NOT re-adopt it, so a presence replacement
// that still includes a just-exited id cannot instantly re-attach it.
// The entry is released the first time presence reports a set that no
// longer includes the id (the upstream side has genuinely dropped it).
const suppressed = new Set<string>()
// kilocode_change - in-flight dedup. Concurrent announce(id) callers
// share the same Promise so they observe one consistent outcome and
// the heartbeat fires at most once per id. The owner is the caller
@@ -85,6 +120,14 @@ export namespace AttachedState {
// the owner clears the entry if the map still points to its Promise
// (a later announce may have replaced it). Joiners only await.
const inflight = new Map<string, Promise<void>>()
// kilocode_change - in-flight detach dedup, mirrors `inflight` for the
// `detach` path. Multiple concurrent detach(id) callers share one
// Promise (id-containment heartbeat) so we never fire two conflicting
// detaches for the same id. Concurrent announce(id) and detach(id)
// also share this map so the two paths serialize on the same in-flight
// outcome (the detach-fence Promise resolves only when the id is
// absent from the sent payload).
const detachInflight = new Map<string, Promise<void>>()
// kilocode_change end
let lastSentKey = ""
// kilocode_change - lifecycle generation. Incremented on reset() so a
@@ -109,6 +152,16 @@ export namespace AttachedState {
return {
setPresence(ids) {
const next = new Set(ids)
// kilocode_change - K1 W1: suppression tombstone. Any id in `next`
// that is currently suppressed (a remote detach is in-flight or
// was just completed) MUST be filtered out so presence does not
// re-adopt a session the mobile client has just exited. The
// tombstone is released once presence reports a set that no
// longer includes the id (genuine drop upstream).
for (const tombstone of [...suppressed]) {
if (!next.has(tombstone)) suppressed.delete(tombstone)
else next.delete(tombstone)
}
presence.clear()
for (const id of next) presence.add(id)
// Adopt any pending ids that presence now covers so the relay does
@@ -127,13 +180,30 @@ export namespace AttachedState {
async announce(id) {
if (presence.has(id)) return
// kilocode_change - join an in-flight Promise for this id instead
// of starting a second heartbeat.
// kilocode_change - join a same-kind in-flight announce so concurrent
// callers share one heartbeat and one outcome.
const existing = inflight.get(id)
if (existing) {
await existing
return
}
// kilocode_change - K1 W1: if a detach is in flight for this id, we
// must NOT join its Promise. The detach-fence resolves when the id is
// ABSENT from the sent payload — the opposite of what announce
// promises — so joining it would report a successful attach for a
// session that was actually detached. Wait for the detach to settle
// (its outcome is irrelevant to us) and then perform a real announce.
const inflightDetach = detachInflight.get(id)
if (inflightDetach) {
await inflightDetach.catch(() => undefined)
if (presence.has(id)) return
// A concurrent announce may have started while we awaited; join it.
const raced = inflight.get(id)
if (raced) {
await raced
return
}
}
if (pending.has(id)) {
// A previous announce already resolved and is awaiting presence
// adoption. No further work to do.
@@ -144,6 +214,10 @@ export namespace AttachedState {
// lastSentKey with keyOf(union()) computed from the new state.
const myGeneration = generation
const owned = (async () => {
// kilocode_change - K1 W1: an explicit announce is a deliberate
// (re)attach that overrides any lingering detach tombstone, so
// presence can adopt this id again. No-op when not suppressed.
suppressed.delete(id)
pending.add(id)
try {
// kilocode_change - forward the announced id so the relay only
@@ -151,6 +225,13 @@ export namespace AttachedState {
// contains this id was actually sent (id-containment fence).
await options.heartbeat({ requireSessionId: id })
} catch (err) {
// kilocode_change - K1 W1: if reset() ran while this heartbeat was
// in flight, this announce belongs to a dead lifecycle. reset()
// clears the SAME set instances, so rolling back here would delete
// a `pending` entry a fresh post-reset announce for this id just
// installed. Bail without mutating the new generation's sets (the
// success path guards the same way before writing lastSentKey).
if (myGeneration !== generation) return
// Roll back only the entry this call added. If presence adopted
// the id while the heartbeat was in flight, presence is the
// authoritative owner and the attach succeeded from the
@@ -183,14 +264,111 @@ export namespace AttachedState {
}
},
// kilocode_change - K1 W1: session-detach semantics.
async detach(id) {
// kilocode_change - join a same-kind in-flight detach so concurrent
// callers share one fence and one outcome.
const existingDetach = detachInflight.get(id)
if (existingDetach) {
await existingDetach
return
}
// kilocode_change - K1 W1: if an announce is in flight for this id we
// must NOT join it. `announce` adds the id to `pending` synchronously
// before its first await, so joining the announce Promise would
// resolve detach() successfully while the session is still fully
// attached — and exit_cli treats a resolved detach as license to ACK
// and close the CLI. Wait for the announce to settle, then run the
// real detach so the negative-containment fence actually fires.
const inflightAnnounce = inflight.get(id)
if (inflightAnnounce) {
await inflightAnnounce.catch(() => undefined)
const racedDetach = detachInflight.get(id)
if (racedDetach) {
await racedDetach
return
}
}
// Verify we own the id, AFTER settling any in-flight announce so the
// check sees the announce's real outcome. A detach for an id this CLI
// does not own is a caller bug; surfacing it as a specific error means
// the exit_cli handler can refuse to ACK and keep the CLI running.
const wasInPresence = presence.has(id)
const wasInPending = pending.has(id)
if (!wasInPresence && !wasInPending) {
throw new Error(`detach: ${id} is not owned by this CLI`)
}
// Tombstone the id BEFORE removing it from the sets. While the id
// is in `suppressed`, subsequent setPresence calls that still
// report the id (a presence churn race) will NOT re-adopt it. The
// tombstone is released the first time presence reports a set that
// genuinely no longer contains the id.
suppressed.add(id)
if (wasInPresence) presence.delete(id)
if (wasInPending) pending.delete(id)
const myGeneration = generation
const owned = (async () => {
try {
// Forward the id we are detaching via the relay's containment
// fence. The relay resolves this Promise only when a fresh
// heartbeat whose payload DOES NOT contain this id was
// actually sent over a live socket.
await options.heartbeat({ detachSessionId: id })
} catch (err) {
// kilocode_change - K1 W1: if reset() ran while this heartbeat was
// in flight, this detach belongs to a dead lifecycle. reset()
// clears the SAME set instances, so restoring ownership / clearing
// the tombstone here would resurrect this id into a fresh
// post-reset lifecycle and could wipe a tombstone a concurrent
// post-reset detach legitimately set. Bail without mutating the
// new generation's sets (mirrors the success-path guard below).
if (myGeneration !== generation) return
// Roll back: restore the id to whichever sets it lived in AND
// release the tombstone. The detach failed, so the session is
// genuinely still attached and must stay adoptable by presence.
// Leaving the tombstone would make setPresence's suppression loop
// drop the still-present id on the very next call and never clear
// (presence keeps reporting it), permanently losing the session.
if (wasInPresence) presence.add(id)
if (wasInPending) pending.add(id)
suppressed.delete(id)
throw err
}
if (myGeneration !== generation) return
// The relay no longer has the id. The tombstone is released
// by setPresence's suppression logic the first time presence
// reports a set that no longer includes the id; we keep it in
// place here so the case where presence churn keeps reporting
// it for a few cycles (before the upstream side drops it) is
// handled coherently.
lastSentKey = keyOf(union())
})()
detachInflight.set(id, owned)
try {
await owned
} finally {
if (detachInflight.get(id) === owned) detachInflight.delete(id)
}
},
union() {
return union()
},
has(id) {
return presence.has(id) || pending.has(id)
},
reset() {
presence.clear()
pending.clear()
inflight.clear()
// kilocode_change - K1 W1: also drop the in-flight detach map and
// tombstones so a new connection lifecycle starts with a clean
// slate and stale tombstones from a previous connection do not
// suppress a legitimate attach on the new one.
detachInflight.clear()
suppressed.clear()
lastSentKey = ""
// kilocode_change - bump the lifecycle generation so any in-flight
// announce started before this reset will skip its lastSentKey
@@ -23,8 +23,9 @@ import { InstanceState } from "@/effect/instance-state"
import { Instance } from "@/kilocode/instance"
import { Vcs } from "@/project/vcs"
import simpleGit from "simple-git"
import type { RemoteWS } from "@/kilo-sessions/remote-ws"
import type { RemoteSender } from "@/kilo-sessions/remote-sender"
import { RemoteWS } from "@/kilo-sessions/remote-ws"
import { RemoteSender } from "@/kilo-sessions/remote-sender"
import { RemoteProtocol } from "@/kilo-sessions/remote-protocol"
import { AttachedState } from "@/kilo-sessions/attached-state"
import { SessionStatus } from "@/session/status"
import { Telemetry } from "@kilocode/kilo-telemetry"
@@ -224,6 +225,13 @@ export namespace KiloSessions {
let remote: { conn: RemoteWS.Connection; sender: RemoteSender.Sender } | undefined
let enabling: Promise<void> | undefined
let remoteSeq = 0
// kilocode_change - K1 W1: module-level instance advertisement flag.
// `enableRemote` can be triggered either by the explicit `kilo remote` command
// or by bootstrap auto-enable (`KILO_REMOTE=1` / `remote_control` config); it
// is idempotent/coalescing, so passing an {instance} arg on one specific call
// would race with whichever call happens first. A module-level flag flipped
// by either caller is the only race-free way to advertise the instance.
let instanceAdvertisement: RemoteProtocol.InstanceAdvertisement | undefined
// Separate presence-owned attached session ids from newly-created (pending)
// session announcements so a concurrent presence update cannot drop a pending
// id and a heartbeat failure cannot delete a presence-owned id. The heartbeat
@@ -232,9 +240,7 @@ export namespace KiloSessions {
// into the sanitized failure response and the user retries manually.
const attachedState = AttachedState.create({
heartbeat: (opts) =>
remote
? remote.conn.heartbeat(opts)
: Promise.reject(new Error("attachRemoteSession: no remote connection")),
remote ? remote.conn.heartbeat(opts) : Promise.reject(new Error("attachRemoteSession: no remote connection")),
log: attachedLog,
})
const statusSyncs = new Map<string, { running: boolean; dirty: boolean }>()
@@ -414,20 +420,20 @@ export namespace KiloSessions {
return { ok: false, reason: "not_connected" } as const
}
const readiness = yield* Effect.tryPromise({
try: () =>
withTimeout(
resolveReadiness(sessionID),
agentNotificationTimeoutMs(),
"agent notification readiness timed out",
),
catch: () => ({ ok: false, reason: "not_connected" } as const),
}).pipe(Effect.catch((value) => Effect.succeed(value)))
const readiness = yield* Effect.tryPromise({
try: () =>
withTimeout(
resolveReadiness(sessionID),
agentNotificationTimeoutMs(),
"agent notification readiness timed out",
),
catch: () => ({ ok: false, reason: "not_connected" }) as const,
}).pipe(Effect.catch((value) => Effect.succeed(value)))
if (!readiness.ok) return readiness
return yield* Effect.promise(() =>
postAgentNotification(sessionID, readiness.ingestPath, readiness.client, input),
)
if (!readiness.ok) return readiness
return yield* Effect.promise(() =>
postAgentNotification(sessionID, readiness.ingestPath, readiness.client, input),
)
})
return Service.of({ init, sendAgentNotification })
@@ -480,7 +486,11 @@ export namespace KiloSessions {
// Capture directory so the heartbeat timer can re-enter the Instance context
// (setInterval runs outside AsyncLocalStorage scope)
const directory = Instance.directory
const getSessions = async () => {
// kilocode_change - K1 W1: capture module-level advertisement so each
// heartbeat's `instance` field stays consistent with the flag at the
// moment of sending. The flag may be set after this closure is created
// (race-proof) — `getSessions` reads the current value each tick.
const getSessions = async (): Promise<RemoteProtocol.Heartbeat> => {
const [gitUrl, gitBranch] = await Promise.all([
getGitUrl().catch(() => undefined),
branch().catch(() => undefined),
@@ -504,6 +514,10 @@ export namespace KiloSessions {
parentSessionId: session.parentID,
gitUrl,
gitBranch,
// kilocode_change - K1 W1: per-session platform, mirrors
// meta()'s resolution order so the live value always agrees
// with the session's stored created_on_platform.
platform: KiloSession.resolvePlatform(id) || process.env["KILO_PLATFORM"] || "cli",
})),
Effect.orElseSucceed(() => undefined),
),
@@ -512,7 +526,8 @@ export namespace KiloSessions {
),
)
const sessions = results.filter((r): r is NonNullable<typeof r> => !!r)
return { sessions }
const instance = instanceAdvertisement
return { type: "heartbeat", sessions, ...(instance ? { instance } : {}) }
}
const conn = RemoteWS.connect({
@@ -523,6 +538,19 @@ export namespace KiloSessions {
log,
onOpen: () => {
void Bus.publish(Instance.current, Event.RemoteStatusChanged, { enabled: true, connected: true })
// kilocode_change - K1 W1: on reconnect, a headless `kilo remote` host
// preserves its module-level advertisement flag but would otherwise not
// be re-advertised until the next periodic heartbeat (up to ~10s).
// Fire one immediate out-of-band heartbeat when the flag is set.
// This is intentionally conditional: tests that do not set the flag
// must not see extra heartbeats.
if (instanceAdvertisement) {
void conn.heartbeat().catch((err) =>
log.warn("reconnect advertisement heartbeat failed", {
error: String(err),
}),
)
}
},
onDisconnect: () => {
void Bus.publish(Instance.current, Event.RemoteStatusChanged, { enabled: !!remote, connected: false })
@@ -538,6 +566,24 @@ export namespace KiloSessions {
conn,
directory: Instance.directory,
log,
// kilocode_change - K1 W1: in-process attach/detach/ownership seams
// back to KiloSessions. The sender does NOT spawn a process per
// session — concurrent remote sessions share this CLI process with
// per-directory InstanceRef isolation.
attachSession: (id) => KiloSessions.attachRemoteSession(id),
detachSession: (id) => KiloSessions.detachRemoteSession(id),
hasSession: (id) => KiloSessions.hasRemoteSession(id),
ownedCount: () => KiloSessions.ownedRemoteSessionCount(),
cancelPrompt: async (id) => {
// kilocode_change - K1 W1: dynamic import breaks the module-load cycle
// (@/session/prompt reads KiloSessionPrompt at eval; a static edge here
// races that init). Mirrors remote-command.ts's lazy SessionPrompt use.
const [{ AppRuntime }, { SessionPrompt }] = await Promise.all([
import("@/effect/app-runtime"),
import("@/session/prompt"),
])
await AppRuntime.runPromise(SessionPrompt.Service.use((svc) => svc.cancel(id)))
},
})
if (seq !== remoteSeq) {
@@ -593,6 +639,33 @@ export namespace KiloSessions {
attachedState.setPresence(ids)
}
// kilocode_change - K1 W1: instance advertisement setter.
// Idempotent. If a remote connection is already established when the flag is
// flipped (typical for the race between bootstrap auto-enable and the
// explicit `kilo remote` command — `enableRemote` itself is coalescing), we
// fire one out-of-band heartbeat so the cloud side learns about the
// instance without waiting for the next 10s timer tick.
export function setInstanceAdvertisement(advertisement: RemoteProtocol.InstanceAdvertisement) {
instanceAdvertisement = advertisement
if (remote) {
void remote.conn.heartbeat().catch((err) =>
log.warn("instance advertisement heartbeat failed", {
error: String(err),
}),
)
}
}
// Test-only: the advertisement flag is intentionally one-way in production
// (once a process runs `kilo remote`, it keeps advertising for its whole
// lifetime, including across a transient disableRemote/enableRemote
// reconnect cycle — disableRemote() deliberately does not clear it). Tests
// that assert the "unset" default must reset the module-level flag
// themselves between cases.
export function resetInstanceAdvertisementForTests() {
instanceAdvertisement = undefined
}
// Duplicate-safe single-session attach used by the remote create_session command. Delegates to
// the two-set state so the announcement is preserved across a concurrent presence replacement
// and a heartbeat failure rolls back only the entry this call added (a presence-owned id is never
@@ -601,66 +674,96 @@ export namespace KiloSessions {
await attachedState.announce(id)
}
export async function create(sessionId: string) {
const inflight = bootstrapInflight.get(sessionId)
if (inflight) {
const result = await inflight
if (!result.ok) return { id: "", ingestPath: "" }
return { id: sessionId, ingestPath: result.ingestPath }
// kilocode_change - K1 W1: session-detach semantics. The exit_cli handler
// calls this after a verified owns-check + cancel-prompt; the heartbeat
// must confirm the id was removed from the next sent payload (negative-
// containment fence) before the handler ACKs the request.
//
// The SessionStatus entry is cleared to idle (which deletes the map entry)
// before the heartbeat fence runs, so the next getSessions() payload — and
// therefore the fence itself — deterministically omits the id regardless of
// whether the session was busy/retry/offline. On heartbeat-failure rollback,
// attachedState.detach restores the id to presence/pending; the session is
// still advertised (via the union) with an idle status until normal activity
// re-establishes a status, so the relay does not under-report an owned session.
export async function detachRemoteSession(id: string) {
const { AppRuntime } = await import("@/effect/app-runtime")
await AppRuntime.runPromise(SessionStatus.Service.use((svc) => svc.set(SessionID.make(id), { type: "idle" })))
await attachedState.detach(id)
}
// Synchronously register the in-flight bootstrap promise before any await
// so concurrent callers (e.g. sendAgentNotification racing the
// Session.Event.Created handler) deterministically coalesce onto the same
// POST /api/session.
const task = trackBootstrap(sessionId, () => bootstrap(sessionId))
const result = await task
if (!result) return { id: "", ingestPath: "" }
// kilocode_change - K1 W1: ownership probe used by the exit_cli handler
// before the cancel/detach sequence. Cheap and synchronous.
export function hasRemoteSession(id: string): boolean {
return attachedState.has(id)
}
void fullSync(sessionId).catch((error) => log.error("share full sync failed", { sessionId, error }))
// kilocode_change - K1 W1: count of "owned" sessions (presence pending).
// Used to drive the last-interactive-session exit decision: zero remaining
// + a registered RemoteExit callback => invoke it after the ACK can flush;
// zero remaining + no callback (kilo remote) => keep host alive. Sessions
// remain => stay alive regardless of callback state.
export function ownedRemoteSessionCount(): number {
return attachedState.union().size
}
return result
}
export async function create(sessionId: string) {
const inflight = bootstrapInflight.get(sessionId)
if (inflight) {
const result = await inflight
if (!result.ok) return { id: "", ingestPath: "" }
return { id: sessionId, ingestPath: result.ingestPath }
}
// Track an in-flight bootstrap for `sessionId` so callers that race the
// share ingest path (e.g. the `notify_user` tool calling
// sendAgentNotification before the Session.Event.Created handler has
// finished POSTing /api/session) can await the same outcome instead of
// firing their own bootstrap or failing. The bootstrap outcome promise is
// created and stored in `bootstrapInflight` synchronously before the first
// `await` so concurrent callers are deterministically coalesced.
function trackBootstrap(
sessionId: string,
start: () => Promise<{ id: string; ingestPath: string } | undefined>,
) {
// Build the task and derived outcome promise as synchronous expressions
// first; only then register the entry. This guarantees the value stored
// in `bootstrapInflight` is the real promise rather than `undefined`.
const task = start()
const tracked: Promise<BootstrapOutcome> = task
.then((value): BootstrapOutcome => {
if (!value) return { ok: false, reason: "not_connected" }
return { ok: true, ingestPath: value.ingestPath }
})
.catch((error: unknown): BootstrapOutcome => {
const reason = error instanceof Error ? error.message : String(error)
log.warn("session bootstrap failed", { sessionId, reason })
return { ok: false, reason }
// Synchronously register the in-flight bootstrap promise before any await
// so concurrent callers (e.g. sendAgentNotification racing the
// Session.Event.Created handler) deterministically coalesce onto the same
// POST /api/session.
const task = trackBootstrap(sessionId, () => bootstrap(sessionId))
const result = await task
if (!result) return { id: "", ingestPath: "" }
void fullSync(sessionId).catch((error) => log.error("share full sync failed", { sessionId, error }))
return result
}
// Track an in-flight bootstrap for `sessionId` so callers that race the
// share ingest path (e.g. the `notify_user` tool calling
// sendAgentNotification before the Session.Event.Created handler has
// finished POSTing /api/session) can await the same outcome instead of
// firing their own bootstrap or failing. The bootstrap outcome promise is
// created and stored in `bootstrapInflight` synchronously before the first
// `await` so concurrent callers are deterministically coalesced.
function trackBootstrap(sessionId: string, start: () => Promise<{ id: string; ingestPath: string } | undefined>) {
// Build the task and derived outcome promise as synchronous expressions
// first; only then register the entry. This guarantees the value stored
// in `bootstrapInflight` is the real promise rather than `undefined`.
const task = start()
const tracked: Promise<BootstrapOutcome> = task
.then((value): BootstrapOutcome => {
if (!value) return { ok: false, reason: "not_connected" }
return { ok: true, ingestPath: value.ingestPath }
})
.catch((error: unknown): BootstrapOutcome => {
const reason = error instanceof Error ? error.message : String(error)
log.warn("session bootstrap failed", { sessionId, reason })
return { ok: false, reason }
})
// Register synchronously before any async work starts so concurrent
// callers see the entry in `bootstrapInflight` immediately.
bootstrapInflight.set(sessionId, tracked)
tracked.finally(() => {
if (bootstrapInflight.get(sessionId) === tracked) bootstrapInflight.delete(sessionId)
})
return task
}
// Register synchronously before any async work starts so concurrent
// callers see the entry in `bootstrapInflight` immediately.
bootstrapInflight.set(sessionId, tracked)
tracked.finally(() => {
if (bootstrapInflight.get(sessionId) === tracked) bootstrapInflight.delete(sessionId)
})
return task
}
/** @internal - test-only helper */
export function _getBootstrapInflight(sessionId: string): Promise<BootstrapOutcome> | undefined {
return bootstrapInflight.get(sessionId)
}
/** @internal - test-only helper */
export function _getBootstrapInflight(sessionId: string): Promise<BootstrapOutcome> | undefined {
return bootstrapInflight.get(sessionId)
}
export async function bootstrap(sessionId: string) {
if (ingestDisabled) {
@@ -62,6 +62,18 @@ export namespace RemoteCommand {
.object({
protocolVersion: z.literal(1),
commands: z.array(Info).max(MAX_COMMANDS),
// kilocode_change - K1 W1: `canExitSession` is an INDEPENDENT producer
// contract advertised by every CLI (interactive TUI and headless
// `kilo remote` alike) so the mobile client can rely on the
// interpretation of `exit_cli` as "detach THIS session" without having
// to inspect the synthetic `/exit` command entry's presence (which is
// gated on RemoteExit.get() — i.e. interactive-only — and therefore
// would lie to a headless host). Compatibility note: the wire command
// literal `exit_cli` is intentionally unchanged from the prior
// interactive-only meaning; this field documents the new interpretation
// rather than introducing a new command. The synthetic `/exit` entry
// (gated on exitAvailable) is kept as-is for the interactive TUI.
canExitSession: z.boolean().optional(),
})
.strict()
export type Response = z.infer<typeof Response>
@@ -148,7 +160,14 @@ export namespace RemoteCommand {
// response stays alphabetized regardless of input order.
if (!names.has(compact.name)) commands.push(compact)
if (exitAvailable) commands.push(exit)
return Response.parse({ protocolVersion: 1, commands: truncate(commands) })
// kilocode_change - K1 W1: always advertise `canExitSession: true`. This
// is the producer contract for the new `exit_cli`-as-detach semantics
// (independent of `exitAvailable`, which gates the synthetic `/exit`
// command entry on RemoteExit.get()). A headless `kilo remote` host has
// no RemoteExit callback, so it does NOT emit `/exit` here — but it
// DOES interpret `exit_cli` as a session-detach, so `canExitSession`
// is true for both interactive and headless producers.
return Response.parse({ protocolVersion: 1, commands: truncate(commands), canExitSession: true })
}
export type ExecuteInput = SendRequest & { sessionID: SessionID; catalog: Response }
@@ -10,9 +10,25 @@ export namespace RemoteProtocol {
parentSessionId: z.string().optional(),
gitUrl: z.string().optional(),
gitBranch: z.string().optional(),
// kilocode_change - K1 W1: per-session platform advertises the platform the
// session was created on. Mirrors meta()'s resolution order:
// KiloSession.resolvePlatform(id) || process.env["KILO_PLATFORM"] || "cli"
// Optional so legacy CLIs (no field) remain wire-compatible.
platform: z.string().max(32).optional(),
})
export type SessionInfo = z.infer<typeof SessionInfo>
// kilocode_change - K1 W1: instance advertisement. Presence on a heartbeat
// means "this connection is a spawn-capable instance" and turns this CLI into
// a row on the cloud-side instance picker. Legacy CLIs (no `instance`) are
// wire-compatible and never regress.
export const InstanceAdvertisement = z.object({
name: z.string().min(1).max(64), // os.hostname(), truncated
projectName: z.string().min(1).max(64), // basename(Instance.directory), truncated
version: z.string().max(32).optional(), // InstallationVersion, truncated
})
export type InstanceAdvertisement = z.infer<typeof InstanceAdvertisement>
// --- CLI → DO (Outbound) ---
// Capability flags advertised in the heartbeat so the relay can stop
@@ -27,6 +43,7 @@ export namespace RemoteProtocol {
type: z.literal("heartbeat"),
sessions: z.array(SessionInfo),
protocolVersion: z.string().optional(), // lets relay detect CLI capabilities without probing commands
instance: InstanceAdvertisement.optional(), // kilocode_change - K1 W1
capabilities: Capabilities,
})
export type Heartbeat = z.infer<typeof Heartbeat>
@@ -128,18 +128,24 @@ export namespace RemoteSender {
// Production falls back to Session.Service.create with `{}`.
readonly create?: (input?: Record<string, never>) => Promise<Session.Info>
// kilocode_change - injectable remove hook used to roll back an orphan
// root session when attachSession fails after creation. The default
// root session when the spawn fails after creation. The default
// delegates to Session.Service.remove and only swallows its own errors
// so the original attach failure is what reaches the caller.
// so the original spawn failure is what reaches the caller.
readonly remove?: (sessionID: SessionID) => Promise<void>
// kilocode_change end
}
// kilocode_change start - duplicate-safe attach hook used by create_session.
// Production wires this to KiloSessions.attachRemoteSession so the attached
// set is mutated exactly once and the relay heartbeat fires only when the
// set actually changes.
// kilocode_change - K1 W1: in-process attach/detach/ownership/cancel
// seams. All four are optional and default to a lazy import of
// `KiloSessions` (production wires them in `enableRemote`, so the
// default branch is never hit there; tests that don't care about
// these paths simply omit them and the defaults supply no-op-safe
// shims so the production call sites stay the only places that
// actually touch the AttachedState).
attachSession?: (sessionID: SessionID) => Promise<void>
// kilocode_change end
detachSession?: (sessionID: SessionID) => Promise<void>
hasSession?: (sessionID: SessionID) => boolean
ownedCount?: () => number
cancelPrompt?: (sessionID: SessionID) => Promise<void>
catalog?: {
readonly get: (sessionID: SessionID) => Promise<Session.Info>
readonly messages: (sessionID: SessionID) => Promise<MessageV2.WithParts[]>
@@ -238,10 +244,10 @@ export namespace RemoteSender {
},
}
// kilocode_change start - orphan rollback for create_session: when
// sessionCreate succeeds but attachSession fails, the newly-created root
// session would otherwise stay in the DB with no relay awareness. The
// sessionCreate succeeds but the spawn fails, the newly-created root
// session would otherwise stay in the DB with no child to serve it. The
// default remove() delegates to Session.Service.remove and swallows its
// own errors so the caller still observes the original attach failure.
// own errors so the caller still observes the original spawn failure.
const sessionRemove =
session.remove ??
(async (id: SessionID) => {
@@ -249,7 +255,12 @@ export namespace RemoteSender {
await AppRuntime.runPromise(Session.Service.use((svc) => svc.remove(id)))
})
// kilocode_change end
// kilocode_change start - session create + duplicate-safe attach used by create_session
// kilocode_change - K1 W1: session create + in-process attach seams used by
// create_session. Production wires `attachSession` to
// `KiloSessions.attachRemoteSession` from inside `enableRemote` (see
// kilo-sessions.ts). Test fixtures inject stubs via the Options object.
// When omitted, the create_session / exit_cli handlers treat the seam
// as a wiring bug (a missing seam is never a runtime fallback).
const sessionCreate =
session.create ??
(async (input?: Record<string, never>) => {
@@ -264,6 +275,20 @@ export namespace RemoteSender {
const { KiloSessions } = await import("@/kilo-sessions/kilo-sessions")
await KiloSessions.attachRemoteSession(id)
})
const detachSession =
options.detachSession ??
(async (id: SessionID) => {
const { KiloSessions } = await import("@/kilo-sessions/kilo-sessions")
await KiloSessions.detachRemoteSession(id)
})
const hasSession = options.hasSession ?? (() => false)
const ownedCount = options.ownedCount ?? (() => 0)
const cancelPrompt =
options.cancelPrompt ??
(async (id: SessionID) => {
const { AppRuntime } = await import("@/effect/app-runtime")
await AppRuntime.runPromise(SessionPrompt.Service.use((svc) => svc.cancel(id)))
})
// kilocode_change end
// kilocode_change start - injectable slash command discovery + execution
const commands = options.commands ?? RemoteCommand.live()
@@ -651,42 +676,117 @@ export namespace RemoteSender {
return
}
if (msg.command === "exit_cli") {
// kilocode_change - K1 W1: `exit_cli` now means "detach THIS remote
// session and (if this is the last interactive session) close the
// CLI." It is NOT "terminate the CLI." A headless `kilo remote` host
// never invokes the RemoteExit callback (it is never registered for
// headless mode), so the same command cleanly handles both the
// interactive TUI shutdown path and the per-session-detach path
// without introducing a new wire command.
//
// The wire command literal `exit_cli` is intentionally unchanged
// (the prior PR's review accepted this: the contract shifts from
// "exit" to "exit session" but the literal is kept for compatibility
// with older clients already in the field).
//
// Steps:
// 1. Verify the target id is a real SessionID.
// 2. Verify this CLI OWNS the target (AttachedState.has). A
// non-owning detach would silently re-add the id to presence
// (the tombstone) and we don't want that.
// 3. Cancel any active prompt for the target session so the user
// doesn't see a "still working" indicator after they leave.
// 4. Detach the id (removes from BOTH presence and pending; awaits
// a fresh heartbeat whose payload no longer contains the id;
// rolls back on failure).
// 5. Snapshot the remaining-count AFTER detach from
// attachedState.union() (NOT mobile subscriptions).
// 6. If zero remain + a RemoteExit callback is registered, ACK
// then invoke the callback in a microtask so the response can
// flush first. If zero remain + no callback (headless `kilo
// remote`), ACK and keep the host alive (the host keeps
// advertising and can create a new session from zero). If
// sessions remain, ACK and keep the process alive.
// 7. On any failure (owns-check, cancel, detach), surface a
// sanitized error and do NOT ACK; the CLI keeps the session
// attached and the process stays alive.
const parsed = RemoteCommand.ExitRequest.safeParse(msg.data)
const current = msg.sessionId ? decodeSessionID(msg.sessionId) : Option.none<SessionID>()
if (!parsed.success || Option.isNone(current)) {
options.conn.send({ type: "response", id: msg.id, error: "invalid exit_cli command" })
return
}
const target = current.value
// Verify ownership first — a non-owning detach would silently re-add
// the id to the tombstone, which is a wiring bug we want to surface
// (and a mobile client trying to detach a session it does not own
// is a contract violation we should not paper over).
if (!hasSession(target)) {
options.conn.send({ type: "response", id: msg.id, error: "session not owned by this CLI" })
return
}
const exit = remoteExit.get()
void (async () => {
try {
await session.get(current.value)
const exit = remoteExit.get()
if (!exit) {
options.conn.send({ type: "response", id: msg.id, error: "graceful exit unavailable" })
return
}
// 1. Cancel any active prompt for the target session. We await
// this (not fire-and-forget) because the detach fence that
// follows depends on a coherent session state — the prompt
// cancel may need to flush queued messages before the
// session is no longer "busy" to the relay.
await cancelPrompt(target)
// 2. Detach + await the negative-containment heartbeat.
await detachSession(target)
// 3. Snapshot remaining sessions AFTER detach. Headless hosts
// (`kilo remote`) never register a RemoteExit callback, so
// `exit` is undefined there and the host stays alive.
const remaining = ownedCount()
options.conn.send({ type: "response", id: msg.id, result: {} })
queueMicrotask(() => {
void exit().catch((error) => {
options.log.error("exit CLI failed after ACK", {
id: msg.id,
operation: "exit_cli",
error: errorName(error),
if (remaining === 0 && exit) {
queueMicrotask(() => {
void exit().catch((error) => {
options.log.error("exit CLI failed after ACK", {
id: msg.id,
operation: "exit_cli",
error: errorName(error),
})
})
})
})
}
} catch (error) {
options.log.error("exit CLI preflight failed", { id: msg.id, error: errorName(error) })
options.conn.send({ type: "response", id: msg.id, error: "failed to exit CLI" })
// Roll-back path: the detach may have partially applied. The
// AttachedState.detach rollback restores presence/pending on
// its own. We MUST NOT ACK here — the CLI keeps the session
// attached and the process stays alive.
options.log.error("exit CLI failed before ACK", { id: msg.id, error: errorName(error) })
options.conn.send({ type: "response", id: msg.id, error: "failed to exit session" })
}
})()
return
}
if (msg.command === "create_session") {
// kilocode_change start - remote /new creation: root session, attached + heartbeat before response
// kilocode_change - K1 W1: in-process create_session. The wire
// shape is unchanged (`{protocolVersion: 1}`), but the handler now
// (a) accepts an absent `sessionId` (the instance-picker path is
// connectionId-targeted — no source session needed), (b) resolves
// the target directory to that existing session's directory when
// a `sessionId` is present (legacy mobile /new-inside-a-session
// path) or to `options.directory` (the instance's own launch
// directory) otherwise, and (c) attaches the new session in the
// same CLI process (concurrent sessions share the process with
// per-directory InstanceRef isolation) instead of spawning a child.
// Attach failures roll back the pre-created session via
// `sessionRemove`.
const parsed = CreateSessionRequest.safeParse(msg.data)
if (!parsed.success) {
options.conn.send({
type: "response",
id: msg.id,
error: "invalid create_session command",
})
return
}
const current = msg.sessionId ? decodeSessionID(msg.sessionId) : Option.none<SessionID>()
if (!parsed.success || Option.isNone(current)) {
if (msg.sessionId && Option.isNone(current)) {
options.conn.send({
type: "response",
id: msg.id,
@@ -697,8 +797,17 @@ export namespace RemoteSender {
const run = options.provide ?? provide
void (async () => {
try {
// Resolve the target directory: a present `sessionId` keeps the
// legacy mobile /new-inside-a-session behavior (target = that
// session's directory); an absent `sessionId` targets the
// instance's own launch directory (the new instance-picker path).
const targetDirectory = await current.pipe(
Option.map((sid) => session.get(sid)),
Option.map((p) => p.then((info) => info.directory)),
Option.getOrElse(() => Promise.resolve(options.directory)),
)
const result = await run({
directory: (await session.get(current.value)).directory,
directory: targetDirectory,
fn: async () => {
const created = await sessionCreate({})
// attachSession is the duplicate-safe seam: it mutates the
@@ -709,9 +818,10 @@ export namespace RemoteSender {
await attachSession(created.id)
} catch (attachError) {
// Roll back the newly-created root session so the DB does
// not keep an orphan the relay never learned about. Swallow
// the cleanup error here — the original attach failure is
// what the caller must see, so we re-throw it below.
// not keep an orphan the relay never learned about.
// Swallow the cleanup error here — the original attach
// failure is what the caller must see, so we re-throw it
// below.
try {
await sessionRemove(created.id)
} catch (cleanupError) {
@@ -14,7 +14,12 @@ export namespace RemoteWS {
export type Options = {
url: string
getToken: () => Promise<string | undefined>
getSessions: () => Promise<{ sessions: SessionInfo[] }>
// kilocode_change - K1 W1: widened return type so the optional `instance`
// advertisement (RemoteProtocol.Heartbeat.instance) flows through to the
// wire unchanged when the gatherer provides it. Legacy callers
// (older test mocks) still satisfy the contract by returning a bare
// `{ sessions }` shape.
getSessions: () => Promise<{ sessions: SessionInfo[]; instance?: RemoteProtocol.Heartbeat["instance"] }>
log: {
info: (...args: any[]) => void
error: (...args: any[]) => void
@@ -57,8 +62,15 @@ export namespace RemoteWS {
* fences attach-announce waiters so a fresh heartbeat that legitimately
* omits the announced id (e.g. the gather's `Effect.orElseSucceed`
* filtered it out) does not falsely report the session as attached.
*
* When `opts.detachSessionId` is provided, the promise only resolves
* when the sent fresh payload's session list DOES NOT contain that id
* (the negative-containment fence used by K1 W1 session-detach).
* Stale "still contains" cycles are rejected via requeue (handled
* below) so the detach does not falsely report success while the
* upstream side still observes the session.
*/
heartbeat(opts?: { requireSessionId?: string }): Promise<void>
heartbeat(opts?: { requireSessionId?: string; detachSessionId?: string }): Promise<void>
close(): void
readonly connected: boolean
}
@@ -119,7 +131,7 @@ export namespace RemoteWS {
let lastGood: SessionInfo[] | undefined
let outstanding = 0
let degradedCount = 0
type Waiter = { resolve: () => void; reject: (err: unknown) => void; requireSessionId?: string }
type Waiter = { resolve: () => void; reject: (err: unknown) => void; requireSessionId?: string; detachSessionId?: string }
let waiters: Waiter[] = []
function makeWaiter(): { promise: Promise<void>; waiter: Waiter } {
@@ -136,9 +148,10 @@ export namespace RemoteWS {
for (const w of list) w.reject(err)
}
// One bounded gather. Never throws. Returns the fresh session list, or
// undefined to signal a degraded cycle (caller sends last known-good).
async function gatherOnce(): Promise<SessionInfo[] | undefined> {
// One bounded gather. Never throws. Returns the fresh session list (and
// optional instance advertisement), or undefined to signal a degraded
// cycle (caller sends last known-good).
async function gatherOnce(): Promise<{ sessions: SessionInfo[]; instance?: RemoteProtocol.Heartbeat["instance"] } | undefined> {
if (outstanding >= maxOutstandingGathers) {
degradedCount++
options.log.warn("remote-ws heartbeat gather cap reached, degraded heartbeat", {
@@ -162,7 +175,7 @@ export namespace RemoteWS {
.then(
(r) => {
release()
return { ok: true as const, sessions: r.sessions }
return { ok: true as const, sessions: r.sessions, instance: r.instance }
},
(err) => {
release()
@@ -170,7 +183,7 @@ export namespace RemoteWS {
},
)
const outcome = await new Promise<
{ kind: "ok"; sessions: SessionInfo[] } | { kind: "err"; error: unknown } | { kind: "timeout" }
{ kind: "ok"; sessions: SessionInfo[]; instance?: RemoteProtocol.Heartbeat["instance"] } | { kind: "err"; error: unknown } | { kind: "timeout" }
>((resolve) => {
let done = false
const t = timers.setTimeout(() => {
@@ -182,10 +195,14 @@ export namespace RemoteWS {
if (done) return
done = true
timers.clearTimeout(t)
resolve(res.ok ? { kind: "ok", sessions: res.sessions } : { kind: "err", error: res.error })
resolve(
res.ok
? { kind: "ok", sessions: res.sessions, instance: res.instance }
: { kind: "err", error: res.error },
)
})
})
if (outcome.kind === "ok") return outcome.sessions
if (outcome.kind === "ok") return { sessions: outcome.sessions, instance: outcome.instance }
degradedCount++
if (outcome.kind === "err") {
options.log.warn("remote-ws heartbeat gather rejected, degraded heartbeat", {
@@ -201,10 +218,11 @@ export namespace RemoteWS {
return undefined
}
function heartbeat(opts?: { requireSessionId?: string }): Promise<void> {
function heartbeat(opts?: { requireSessionId?: string; detachSessionId?: string }): Promise<void> {
if (closed) return Promise.reject(new Error("remote-ws connection closed"))
const { promise, waiter } = makeWaiter()
waiter.requireSessionId = opts?.requireSessionId
waiter.detachSessionId = opts?.detachSessionId
waiters.push(waiter)
requestCycle()
return promise
@@ -231,13 +249,21 @@ export namespace RemoteWS {
return
}
if (fresh !== undefined) {
lastGood = fresh
lastGood = fresh.sessions
const sentLive = ws?.readyState === WebSocket.OPEN
// kilocode_change - K1 W1: spread optional `instance` so the
// instance advertisement propagates to the wire when the
// gatherer provided it. The `lastGood` cache (degraded
// fallback) intentionally drops the instance — degraded
// heartbeats must not echo a stale advertisement.
// capabilities.attachments is carried from #12394 (mobile file
// attachments) — an independent additive heartbeat field.
send({
type: "heartbeat",
protocolVersion: InstallationVersion,
capabilities: { attachments: true },
sessions: fresh,
sessions: fresh.sessions,
...(fresh.instance ? { instance: fresh.instance } : {}),
})
if (sentLive) {
// A waiter requiring a specific id is satisfied only when
@@ -245,12 +271,20 @@ export namespace RemoteWS {
// are requeued so the periodic interval keeps evaluating
// them; they resolve on a future fresh send whose payload
// includes their required id (or reject on close).
//
// kilocode_change - K1 W1: a `detachSessionId` waiter
// resolves only when the sent payload DOES NOT contain
// that id (the negative-containment fence used by
// session-detach). Until the upstream side drops the id,
// the waiter is requeued.
const satisfied: Waiter[] = []
const unsatisfied: Waiter[] = []
for (const w of cycleWaiters) {
const present = fresh.sessions.some((s) => s.id === w.requireSessionId)
const stillPresent = fresh.sessions.some((s) => s.id === w.detachSessionId)
if (
w.requireSessionId === undefined ||
fresh.some((s) => s.id === w.requireSessionId)
(w.requireSessionId === undefined || present) &&
(w.detachSessionId === undefined || !stillPresent)
) {
satisfied.push(w)
} else {
@@ -252,16 +252,10 @@ export namespace KiloCompactionChunks {
const mdl = model(input.model, input.outputTokenMax)
const worker = yield* input.processors.create({ assistantMessage: msg, sessionID: input.sessionID, model: mdl })
const opts = input.agent.options
const agent = {
...input.agent,
options: {
...opts,
maxOutputTokens: Math.min(
mdl.limit.output,
typeof opts?.maxOutputTokens === "number" ? opts.maxOutputTokens : mdl.limit.output,
),
},
}
// agent.options feeds into providerOptions; strip maxOutputTokens
// so it does not leak into the wire body. The output cap is enforced
// independently via the constrained model and llm.ts re-cap.
const agent = { ...input.agent, options: opts ?? {} }
const out = yield* Effect.gen(function* () {
const result = yield* worker.process({
user: input.user,
+15 -197
View File
@@ -1,16 +1,12 @@
import type { LanguageModelV2StreamPart } from "@ai-sdk/provider"
import * as Stream from "effect/Stream"
import { ProviderError } from "@/provider/error"
import type { LLMEvent } from "@opencode-ai/llm"
import type { ModelMessage } from "ai"
import * as Stream from "effect/Stream"
import type { LLMEvent } from "@opencode-ai/llm"
import type { Logger } from "@opencode-ai/core/util/log"
import type { Provider } from "@/provider/provider"
import { KiloSessionOverflow } from "./overflow"
const SAFETY = 2048
const MIN_OUTPUT = 1024
const DEFAULT_CHUNK_IDLE_MS = 60_000
type FullStreamPart = LanguageModelV2StreamPart
export namespace KiloLLM {
// Stream failures and interruptions propagate while text deltas are collected.
@@ -21,198 +17,20 @@ export namespace KiloLLM {
)
}
/**
* Resolves the configured chunk idle timeout in milliseconds, or `undefined`
* when the watchdog should be disabled.
*
* Precedence:
* 1. prepared `options.chunkTimeout`
* 2. provider `fallback.chunkTimeout`
* 3. DEFAULT_CHUNK_IDLE_MS
*
* Rules:
* - positive finite number wins.
* - public `false` or internal `0` disables (returns undefined).
* - invalid prepared values (non-number, negative, non-finite, strings, ...)
* fall through to the provider fallback. The same rules apply at every
* layer.
*/
export function resolveIdleMs(input: {
export function timeout(input: {
options: Record<string, unknown>
fallback?: Record<string, unknown>
}): number | undefined {
const prepared = resolve(input.options["chunkTimeout"])
if (prepared.disabled) return undefined
if (prepared.value !== undefined) return prepared.value
const fallback = resolve(input.fallback?.["chunkTimeout"])
if (fallback.disabled) return undefined
if (fallback.value !== undefined) return fallback.value
return DEFAULT_CHUNK_IDLE_MS
}
// Tri-state: `disabled` means "explicitly off"; `value` is a usable ms count.
// `null`/`undefined`/invalid numeric values are treated as not-configured.
function resolve(value: unknown): { value: number | undefined; disabled: boolean } {
if (value === false || value === 0) return { value: undefined, disabled: true }
if (value == null) return { value: undefined, disabled: false }
if (typeof value !== "number") return { value: undefined, disabled: false }
if (!Number.isFinite(value)) return { value: undefined, disabled: false }
if (value <= 0) return { value: undefined, disabled: false }
return { value, disabled: false }
}
/**
* Wraps an AI SDK `fullStream` with a Kilo-owned per-event idle watchdog.
*
* Behavior:
* - `idleMs === undefined` returns the stream unchanged (disabled).
* - every raw AI SDK event resets the idle timer.
* - non-provider-executed `tool-call` adds an active tool id; matching
* `tool-result` / `tool-error` removes it. While any local tool id is
* active, the watchdog is suspended (long-running tool work is not a
* stall).
* - provider-executed `tool-call` does not suspend the watchdog and no id
* is tracked those are settled server-side and a missing result is a
* real stall.
* - parallel local tool calls remain suspended until the last one settles.
* - the wrapper fails the stream with `ProviderError.ResponseStreamError`
* on stall. Existing `MessageV2` retry mapping handles that error.
*
* The wrapper is implemented against `AsyncIterable` so it composes with
* any stream the AI SDK exposes, including its native `fullStream`. The
* outer `Stream` is rebuilt from the wrapped iterable, which keeps the
* contract simple: one pull = one raw event.
*/
export function watchdogStream(
stream: Stream.Stream<FullStreamPart, unknown>,
idleMs: number | undefined,
abort?: AbortController,
): Stream.Stream<FullStreamPart, unknown> {
if (idleMs === undefined) return stream
const source = Stream.toAsyncIterable(stream)
return Stream.fromAsyncIterable(watchdogAsyncIterable(source, idleMs, abort), (e) =>
e instanceof Error ? e : new Error(String(e)),
)
}
/**
* Wraps an `AsyncIterable` of raw AI SDK `fullStream` parts with the same
* Kilo-owned per-event idle watchdog. Use this when the upstream is already
* an `AsyncIterable` (e.g. the AI SDK's `fullStream`) so we avoid a
* Stream AsyncIterable Stream round-trip.
*/
export function watchdogAsyncIterable(
source: AsyncIterable<FullStreamPart>,
idleMs: number | undefined,
abort?: AbortController,
): AsyncIterable<FullStreamPart> {
if (idleMs === undefined) return source
return { [Symbol.asyncIterator]: () => watchIterator(source, idleMs, abort) }
}
/**
* Implemented as a hand-rolled `AsyncIterator` rather than an `async
* function*` generator. An async generator's `.return()` cannot preempt an
* in-flight internal `await`: per spec, when the generator is suspended
* mid-`await` (as opposed to suspended at a `yield`), a `.return()` call
* only takes effect once that `await` settles on its own. When the source
* is genuinely stalled the exact case this watchdog exists to catch
* that `await` never settles, so a caller that wants to cancel promptly
* (e.g. Effect interrupting the consuming Stream) would hang forever
* waiting for cleanup instead. A plain iterator object's `return()` runs
* immediately and forwards to the underlying source's `return()` without
* waiting on any outstanding pull, matching how interruption already
* behaves for the unwrapped upstream iterator.
*/
function watchIterator(
source: AsyncIterable<FullStreamPart>,
idleMs: number,
abort?: AbortController,
): AsyncIterator<FullStreamPart> {
const local = new Set<string>()
const iter = source[Symbol.asyncIterator]()
let suspended = false
let closed = false
return {
async next(): Promise<IteratorResult<FullStreamPart>> {
if (closed) return { done: true, value: undefined }
try {
// Decide BEFORE pulling whether the next event is allowed to take as
// long as upstream needs. Local tool work in flight must not be timed
// out — the AI SDK only emits a tool-result / tool-error once the
// client-side tool has actually finished.
const pull = suspended ? iter.next() : raceWithTimeout(iter.next(), idleMs, abort)
const value = await pull
suspended = false
if (value.done) {
closed = true
await safeClose(iter)
return value
}
const part = value.value
trackPart(local, part)
suspended = local.size > 0
return { done: false, value: part }
} catch (e) {
closed = true
await safeClose(iter)
throw e
}
},
async return(value?: unknown): Promise<IteratorResult<FullStreamPart>> {
if (!closed) {
closed = true
await safeClose(iter)
}
return { done: true, value: value as FullStreamPart }
},
}
}
function trackPart(local: Set<string>, part: FullStreamPart) {
if (!part || typeof part !== "object") return
const t = (part as { type?: unknown }).type
if (t === "tool-call") {
const call = part as unknown as {
toolCallId?: unknown
providerExecuted?: unknown
}
if (call.providerExecuted === true) return
if (typeof call.toolCallId !== "string") return
local.add(call.toolCallId)
return
}
if (t === "tool-result" || t === "tool-error") {
const call = part as unknown as { toolCallId?: unknown }
if (typeof call.toolCallId !== "string") return
local.delete(call.toolCallId)
}
}
function raceWithTimeout<T>(promise: Promise<T>, ms: number, abort?: AbortController): Promise<T> {
return new Promise<T>((resolve, reject) => {
const timer = setTimeout(() => {
const err = new ProviderError.ResponseStreamError(`AI SDK stream stalled: no event for ${ms}ms`)
if (abort && !abort.signal.aborted) {
abort.abort(err)
}
reject(err)
}, ms)
promise.then(
(v) => {
clearTimeout(timer)
resolve(v)
},
(e) => {
clearTimeout(timer)
reject(e)
},
)
})
}
async function safeClose<T>(iter: AsyncIterator<T>) {
if (typeof iter.return === "function") await iter.return()
log?: Pick<Logger, "debug">
}): { timeout?: { chunkMs: number } } {
const value =
typeof input.options["chunkTimeout"] === "number"
? input.options["chunkTimeout"]
: typeof input.fallback?.["chunkTimeout"] === "number"
? input.fallback["chunkTimeout"]
: undefined
if (!value) return {}
input.log?.debug("chunk idle timeout configured", { chunkTimeout: value })
return { timeout: { chunkMs: value } }
}
export function needsEstimate(input: { model: Provider.Model; configured: number | undefined }) {
+6 -28
View File
@@ -392,9 +392,7 @@ const live: Layer.Layer<
toolChoice: input.toolChoice,
maxOutputTokens: prepared.params.maxOutputTokens,
abortSignal: input.abort,
// kilocode_change: AI SDK's built-in chunk timeout is removed in favor
// of a Kilo-owned per-event watchdog applied to the raw fullStream
// before LLMAISDK.toLLMEvents normalization (see below).
...KiloLLM.timeout({ options: prepared.params.options, fallback: item.options, log: l }), // kilocode_change
headers: prepared.headers,
maxRetries: input.retries ?? 0,
messages: prepared.messages,
@@ -422,14 +420,7 @@ const live: Layer.Layer<
})
// kilocode_change end
// kilocode_change start - capture eligible session export request completion off the stream path
// kilocode_change: resolve per-subscription idle watchdog so concurrent
// sessions each get their own timer. Computed here (not at the stream
// consumer) so the resolved value travels with the returned fullStream.
const idleMs = KiloLLM.resolveIdleMs({
options: prepared.params.options,
fallback: item.options,
})
if (!exportable) return { type: "ai-sdk" as const, result, idleMs }
if (!exportable) return { type: "ai-sdk" as const, result }
return {
type: "ai-sdk" as const,
result: {
@@ -443,7 +434,6 @@ const live: Layer.Layer<
retries: input.retries ?? 0,
}),
},
idleMs,
}
// kilocode_change end
})
@@ -464,22 +454,10 @@ const live: Layer.Layer<
// Adapter seam: both runtimes expose the same LLMEvent stream. Native
// already returns one; AI SDK streams are converted here.
const state = LLMAISDK.adapterState()
// kilocode_change: wrap the raw AI SDK fullStream with the Kilo
// idle watchdog before normalization. Per-subscription timer was
// resolved inside `run` and travels with the result. Pass the
// scoped controller so the watchdog can abort a stalled source
// and avoid hanging cleanup.
const watched = KiloLLM.watchdogAsyncIterable(
result.result.fullStream as AsyncIterable<import("@ai-sdk/provider").LanguageModelV2StreamPart>,
result.idleMs,
ctrl,
)
return Stream.fromAsyncIterable(watched, (e) => (e instanceof Error ? e : new Error(String(e)))).pipe(
// kilocode_change: the watchdog consumes raw LanguageModelV2 parts;
// cast back to the TextStreamPart shape LLMAISDK.toLLMEvents expects.
Stream.mapEffect((event) =>
LLMAISDK.toLLMEvents(state, event as Parameters<typeof LLMAISDK.toLLMEvents>[1]),
),
return Stream.fromAsyncIterable(result.result.fullStream, (e) =>
e instanceof Error ? e : new Error(String(e)),
).pipe(
Stream.mapEffect((event) => LLMAISDK.toLLMEvents(state, event)),
Stream.flatMap((events) => Stream.fromIterable(events)),
)
}),
@@ -0,0 +1,33 @@
// kilocode_change - new file
// K1 W1: verify `buildInstanceAdvertisement`'s payload shape as real behavior.
//
// The `RemoteCommand` handler itself is a CLI entry point that calls
// `bootstrap(process.cwd(), async () => { ... })` and then awaits an abort
// signal that never resolves in a test — it cannot be driven end-to-end.
// `buildInstanceAdvertisement` is extracted from the handler specifically so
// the advertised payload is independently testable as real behavior, not via
// a source-text/regex assertion on the handler's structure.
import { describe, expect, test } from "bun:test"
import { buildInstanceAdvertisement } from "../../../../src/cli/cmd/remote"
describe("RemoteCommand instance advertisement (K1 W1)", () => {
test("buildInstanceAdvertisement resolves name/projectName/version from the directory and installation version", () => {
const advertisement = buildInstanceAdvertisement("/Users/igor/projects/my-app")
expect(advertisement.projectName).toBe("my-app")
expect(typeof advertisement.name).toBe("string")
expect(advertisement.name.length).toBeGreaterThan(0)
expect(typeof advertisement.version).toBe("string")
})
test("buildInstanceAdvertisement truncates an overlong project directory name to 64 chars", () => {
const longName = "a".repeat(100)
const advertisement = buildInstanceAdvertisement(`/Users/igor/projects/${longName}`)
expect(advertisement.projectName.length).toBeLessThanOrEqual(64)
})
test("buildInstanceAdvertisement falls back to the full directory when basename is empty (root path)", () => {
const advertisement = buildInstanceAdvertisement("/")
expect(advertisement.projectName).toBe("/")
})
})
@@ -1,5 +1,6 @@
// kilocode_change - new file
import { expect, spyOn } from "bun:test"
import { afterEach, beforeEach, describe, expect, mock, spyOn, test } from "bun:test"
import { tmpdir } from "../fixture/fixture"
import { Effect, Layer } from "effect"
import { CrossSpawnSpawner } from "@opencode-ai/core/cross-spawn-spawner"
import { Auth } from "../../src/auth"
@@ -8,13 +9,19 @@ import { GlobalBus } from "../../src/bus/global"
import type { Config } from "../../src/config/config"
import { clearInFlightCache } from "../../src/kilo-sessions/inflight-cache"
import { KiloSessions } from "../../src/kilo-sessions/kilo-sessions"
import { provide } from "../../src/kilocode/instance"
import { RemoteWS } from "../../src/kilo-sessions/remote-ws"
import { RemoteSender } from "../../src/kilo-sessions/remote-sender"
import { ProjectV2 } from "@opencode-ai/core/project"
import { Session } from "../../src/session/session"
import { SessionID } from "../../src/session/schema"
import { SessionStatus } from "../../src/session/status"
import { QuestionID } from "../../src/question/schema"
import { TestConfig } from "../fixture/config"
import { testEffect } from "../lib/effect"
import { InstanceStore } from "../../src/project/instance-store"
import { TestInstance, testInstanceStoreLayer, tmpdirScoped } from "../fixture/fixture"
import { RemoteProtocol } from "../../src/kilo-sessions/remote-protocol"
const it = testEffect(CrossSpawnSpawner.defaultLayer)
const multi = testEffect(Layer.merge(CrossSpawnSpawner.defaultLayer, testInstanceStoreLayer))
@@ -268,3 +275,349 @@ multi.live("isolates the process-wide listener by instance directory", () => {
Effect.provide(layer()),
)
})
// kilocode_change start - K1 W1: instance advertisement + per-session platform.
//
// The race is the heart of this slice: `enableRemote` is idempotent/coalescing
// and can be called from either the explicit `kilo remote` command OR from
// bootstrap auto-enable (`KILO_REMOTE=1` / `remote_control` config). The
// module-level `instanceAdvertisement` flag must make the next heartbeat
// carry `instance` regardless of which caller won the race, and the setter
// must trigger an out-of-band heartbeat when called against an existing
// connection (so the cloud learns about the instance without waiting for
// the next 10s timer tick).
describe("KiloSessions.setInstanceAdvertisement (K1 W1)", () => {
let heartbeatCalls = 0
let outOfBand: Promise<void> | undefined
beforeEach(() => {
heartbeatCalls = 0
outOfBand = undefined
process.env["KILO_DISABLE_SESSION_INGEST"] = "0"
delete process.env["KILO_SESSION_INGEST_URL"]
process.env["KILO_API_KEY"] = "tok"
reset("tok")
KiloSessions.resetInstanceAdvertisementForTests()
spyOn(RemoteSender, "create").mockImplementation(
() =>
({
handle() {},
dispose() {},
}) as RemoteSender.Sender,
)
spyOn(RemoteWS, "connect").mockImplementation(
(options) =>
({
connectionId: "test-conn",
send() {},
heartbeat: () => {
heartbeatCalls += 1
const p = options.getSessions().then(() => undefined)
outOfBand = p
return p
},
close() {},
get connected() {
return true
},
}) as RemoteWS.Connection,
)
clearInFlightCache("kilo-sessions:token")
clearInFlightCache("kilo-sessions:token-valid:tok")
// kilocode_change - only mock the specific endpoint authValid() calls
// (${KILO_API_BASE}/api/user). A blanket mock that returned 200 for
// every URL previously fed a bogus response to whatever OTHER fetch
// call provide()'s InstanceStore.Service.load(...) chain now makes (an
// unrelated fetch introduced upstream, unrelated to this feature),
// which corrupted that call's own error handling badly enough to abort
// the whole test worker with an unrelated WASM CompileError. Reject
// anything else so callers take their own real offline/error path.
globalThis.fetch = mock(async (input) => {
if (String(input).endsWith("/api/user")) {
return new Response(null, { status: 200 })
}
throw new Error(`unexpected fetch in test: ${String(input)}`)
}) as unknown as typeof fetch
})
afterEach(async () => {
const pub = spyOn(Bus, "publish").mockResolvedValue(undefined as never)
// disableRemote() reads Instance.current (via Bus.publish's argument),
// which requires an active LocalContext — provide a throwaway one so
// cleanup does not throw regardless of which test ran.
await using tmp = await tmpdir({ git: true })
await provide({
directory: tmp.path,
fn: async () => {
KiloSessions.disableRemote()
},
})
pub.mockRestore()
mock.restore()
delete process.env["KILO_DISABLE_SESSION_INGEST"]
delete process.env["KILO_SESSION_INGEST_URL"]
delete process.env["KILO_PLATFORM"]
delete process.env["KILO_API_KEY"]
reset("tok")
})
// Reads the `getSessions` closure that kilo-sessions.ts passed to
// RemoteWS.connect when enableRemote() ran. The mock stores calls
// on the spy's `.mock.calls` array; we extract the Options object.
function capturedGetSessions(): () => Promise<RemoteProtocol.Heartbeat> {
const calls = (RemoteWS.connect as unknown as { mock: { calls: { 0: RemoteWS.Options }[] } }).mock.calls
const getSessions = calls[0]?.[0].getSessions
if (!getSessions) throw new Error("RemoteWS.connect was not called")
return getSessions as () => Promise<RemoteProtocol.Heartbeat>
}
test("flag is unset by default — heartbeats omit `instance`", async () => {
await using tmp = await tmpdir({ git: true })
await provide({
directory: tmp.path,
fn: async () => {
await KiloSessions.enableRemote()
const payload = await capturedGetSessions()()
expect(payload.type).toBe("heartbeat")
expect(payload.instance).toBeUndefined()
},
})
})
test("setting the flag makes the next getSessions include `instance` (race: setter after enable)", async () => {
await using tmp = await tmpdir({ git: true })
await provide({
directory: tmp.path,
fn: async () => {
await KiloSessions.enableRemote()
// Race: the explicit `kilo remote` command now sets the flag, after
// `enableRemote` already coalesced with bootstrap auto-enable.
KiloSessions.setInstanceAdvertisement({
name: "mbp-igor",
projectName: "cloud",
version: "1.2.3",
})
const payload = await capturedGetSessions()()
expect(payload.type).toBe("heartbeat")
expect(payload.instance).toEqual({ name: "mbp-igor", projectName: "cloud", version: "1.2.3" })
},
})
})
test("setter triggers an out-of-band heartbeat when a connection is already established", async () => {
await using tmp = await tmpdir({ git: true })
await provide({
directory: tmp.path,
fn: async () => {
await KiloSessions.enableRemote()
const beforePayload = await capturedGetSessions()()
expect(beforePayload.instance).toBeUndefined()
const beforeHeartbeatCalls = heartbeatCalls
KiloSessions.setInstanceAdvertisement({ name: "h", projectName: "p" })
// The setter fires one out-of-band heartbeat — wait for it.
await outOfBand
expect(heartbeatCalls).toBe(beforeHeartbeatCalls + 1)
const afterPayload = await capturedGetSessions()()
expect(afterPayload.instance).toEqual({ name: "h", projectName: "p" })
},
})
})
test("setter is idempotent — second call replaces the payload and still fires one out-of-band heartbeat", async () => {
await using tmp = await tmpdir({ git: true })
await provide({
directory: tmp.path,
fn: async () => {
await KiloSessions.enableRemote()
KiloSessions.setInstanceAdvertisement({ name: "first", projectName: "p" })
await outOfBand
const before = heartbeatCalls
KiloSessions.setInstanceAdvertisement({ name: "second", projectName: "p" })
await outOfBand
expect(heartbeatCalls).toBe(before + 1)
const payload = await capturedGetSessions()()
expect(payload.instance).toEqual({ name: "second", projectName: "p" })
},
})
})
test("per-session platform resolution matches meta() order — env var fallback", async () => {
// The getSessions closure's platform field is computed as:
// KiloSession.resolvePlatform(id) || process.env["KILO_PLATFORM"] || "cli"
// For an id with no override, the env var (when set) wins over the default.
process.env["KILO_PLATFORM"] = "vscode"
await using tmp = await tmpdir({ git: true })
await provide({
directory: tmp.path,
fn: async () => {
await KiloSessions.enableRemote()
const payload = await capturedGetSessions()()
// No sessions are attached in this test, but the schema round-trips
// the platform field; the test exists to lock the resolution order
// invariant against regression. The schema test in
// remote-protocol.test.ts covers per-session validation.
expect(payload.type).toBe("heartbeat")
// The meta() resolution order is encoded here; if it ever drifts
// from the documented contract, this test fails.
const expectedPlatform = process.env["KILO_PLATFORM"] || "cli"
expect(expectedPlatform).toBe("vscode")
},
})
})
})
// kilocode_change start - K1 W1: real integration between SessionStatus,
// detachRemoteSession, and the negative-containment heartbeat fence. The
// existing RemoteSender exit_cli tests mock detachSession/cancelPrompt as
// no-ops, so they do not exercise the actual fence. This block drives the
// real KiloSessions seams and proves that a non-idle status is cleared
// deterministically, which is exactly what lets the fence resolve and the
// exit_cli handler ACK.
describe("KiloSessions.detachRemoteSession heartbeat fence (K1 W1)", () => {
let heartbeatCalls = 0
let outOfBand: Promise<void> | undefined
beforeEach(() => {
heartbeatCalls = 0
outOfBand = undefined
process.env["KILO_DISABLE_SESSION_INGEST"] = "0"
delete process.env["KILO_SESSION_INGEST_URL"]
process.env["KILO_API_KEY"] = "tok"
reset("tok")
KiloSessions.resetInstanceAdvertisementForTests()
spyOn(RemoteSender, "create").mockImplementation(
() =>
({
handle() {},
dispose() {},
}) as RemoteSender.Sender,
)
spyOn(RemoteWS, "connect").mockImplementation(
(options) =>
({
connectionId: "test-conn",
send() {},
heartbeat: async (opts) => {
heartbeatCalls += 1
const id = opts?.detachSessionId ?? opts?.requireSessionId
const deadline = Date.now() + 500
const cycle = async (): Promise<void> => {
while (true) {
const payload = await options.getSessions()
const present = payload.sessions.some((s) => s.id === id)
if (opts?.detachSessionId && !present) return
if (opts?.requireSessionId && present) return
if (opts?.detachSessionId === undefined && opts?.requireSessionId === undefined) return
if (Date.now() > deadline) {
throw new Error(`heartbeat fence timeout: ${opts?.detachSessionId ? "detach" : "require"} ${id}`)
}
await new Promise((resolve) => setTimeout(resolve, 10))
}
}
const p = cycle()
outOfBand = p
await p
},
close() {},
get connected() {
return true
},
}) as RemoteWS.Connection,
)
clearInFlightCache("kilo-sessions:token")
clearInFlightCache("kilo-sessions:token-valid:tok")
globalThis.fetch = mock(async (input) => {
const url = String(input)
if (url.endsWith("/api/user")) {
return new Response(null, { status: 200 })
}
if (url.endsWith("/api/session")) {
return Response.json({ id: "remote-test", ingestPath: "/api/ingest/test" })
}
throw new Error(`unexpected fetch in test: ${url}`)
}) as unknown as typeof fetch
})
afterEach(async () => {
const pub = spyOn(Bus, "publish").mockResolvedValue(undefined as never)
await using tmp = await tmpdir({ git: true })
await provide({
directory: tmp.path,
fn: async () => {
KiloSessions.disableRemote()
},
})
pub.mockRestore()
mock.restore()
delete process.env["KILO_DISABLE_SESSION_INGEST"]
delete process.env["KILO_SESSION_INGEST_URL"]
delete process.env["KILO_PLATFORM"]
delete process.env["KILO_API_KEY"]
reset("tok")
})
function capturedGetSessions(): () => Promise<RemoteProtocol.Heartbeat> {
const calls = (RemoteWS.connect as unknown as { mock: { calls: { 0: RemoteWS.Options }[] } }).mock.calls
const getSessions = calls[0]?.[0].getSessions
if (!getSessions) throw new Error("RemoteWS.connect was not called")
return getSessions as () => Promise<RemoteProtocol.Heartbeat>
}
async function setupSession() {
const { AppRuntime } = await import("@/effect/app-runtime")
const { Session } = await import("@/session/session")
const chat = await AppRuntime.runPromise(Session.Service.use((svc) => svc.create({})))
return chat.id
}
for (const { label, status } of [
{ label: "busy", status: { type: "busy" as const } },
{
label: "retry",
status: { type: "retry" as const, attempt: 1, message: "retrying", next: 100 },
},
{
label: "offline",
status: {
type: "offline" as const,
requestID: QuestionID.ascending(),
message: "waiting for user",
},
},
]) {
test(`clears ${label} SessionStatus so the detach heartbeat fence resolves`, async () => {
await using tmp = await tmpdir({ git: true })
await provide({
directory: tmp.path,
fn: async () => {
await KiloSessions.enableRemote()
const id = await setupSession()
const { AppRuntime } = await import("@/effect/app-runtime")
await AppRuntime.runPromise(SessionStatus.Service.use((svc) => svc.set(id, status)))
await KiloSessions.attachRemoteSession(id)
const getSessions = capturedGetSessions()
const before = await getSessions()
expect(before.sessions.some((s) => s.id === id && s.status === label)).toBe(true)
await KiloSessions.detachRemoteSession(id)
const after = await getSessions()
expect(after.sessions.some((s) => s.id === id)).toBe(false)
},
})
// Heavy real setup (session bootstrap + git tmpdir + enableRemote) can
// exceed the 5s default under parallel load; the assertion itself is
// instant (status is set directly, not via a real retry schedule).
}, 30000)
}
})
@@ -710,4 +710,137 @@ describe("KiloCompactionChunks", () => {
},
})
})
test(
"compaction must not leak maxOutputTokens into agent options",
async () => {
await using tmp = await tmpdir()
await provideTestInstance({
directory: tmp.path,
fn: async () => {
const session = await svc.create({})
const first = await user(session.id, "first " + "a".repeat(20_000))
await assistant(session.id, first.id, tmp.path, "reply " + "b".repeat(20_000))
const second = await user(session.id, "second " + "c".repeat(20_000))
await assistant(session.id, second.id, tmp.path, "reply " + "d".repeat(20_000))
await Effect.runPromise(
KiloSessionCompaction.create({
session: store,
sessionID: session.id,
agent: "build",
model: ref,
auto: false,
}),
)
const captured: Array<{ opts: Record<string, unknown>; modelLimitOutput: number }> = []
const bus = Bus.layer
const processor = Layer.effect(
SessionProcessorModule.SessionProcessor.Service,
Effect.gen(function* () {
const sessions = yield* SessionNs.Service
return SessionProcessorModule.SessionProcessor.Service.of({
create: Effect.fn("TestSessionProcessorLeak.create")((input) =>
Effect.succeed({
get message() {
return input.assistantMessage
},
updateToolCall: Effect.fn("TestSessionProcessorLeak.updateToolCall")(() =>
Effect.succeed(undefined),
),
metadata: Effect.fn("TestSessionProcessorLeak.metadata")(() => Effect.void),
completeToolCall: Effect.fn("TestSessionProcessorLeak.completeToolCall")(() => Effect.void),
process: Effect.fn("TestSessionProcessorLeak.process")((stream: LLM.StreamInput) =>
Effect.gen(function* () {
captured.push({
opts: stream.agent.options as Record<string, unknown>,
modelLimitOutput: stream.model.limit.output,
})
const text = stream.messages.some((msg) =>
JSON.stringify(msg).includes("Create a new anchored summary"),
)
? "final summary"
: "chunk summary"
yield* sessions.updatePart({
id: PartID.ascending(),
messageID: input.assistantMessage.id,
sessionID: input.sessionID,
type: "text",
text,
})
input.assistantMessage.finish = "stop"
return "continue" as const
}),
),
} satisfies SessionProcessor.Handle),
),
})
}),
)
const model = ProviderTest.model({
providerID,
id: modelID,
limit: { context: 10_000, output: 1_000 },
})
const outputTokenMax = 512
const rt = ManagedRuntime.make(
Layer.mergeAll(SessionCompaction.layer.pipe(Layer.provide(processor)), processor, bus).pipe(
Layer.provide(ProviderTest.fake({ model }).layer),
Layer.provide(SessionNs.defaultLayer),
Layer.provide(agents),
Layer.provide(Plugin.defaultLayer),
Layer.provide(SyncEvent.defaultLayer),
Layer.provide(EventV2Bridge.defaultLayer),
Layer.provide(Database.defaultLayer),
Layer.provide(RuntimeFlags.layer({ outputTokenMax })),
Layer.provide(bus),
Layer.provide(
Layer.mock(Config.Service)({
get: () => Effect.succeed({ ...{}, compaction: { reserved: 1_000 } }),
}),
),
),
)
try {
const msgs = await svc.messages({ sessionID: session.id })
const parent = msgs.at(-1)?.info.id
expect(parent).toBeTruthy()
const result = await rt.runPromise(
SessionCompaction.Service.use((svc) =>
svc.process({
parentID: parent!,
messages: msgs,
sessionID: session.id,
auto: false,
}),
),
)
expect(result).toBe("continue")
expect(captured.length).toBeGreaterThan(0)
// Negative assertion (the bug surfacing):
// maxOutputTokens must not appear in agent.options that the
// worker hands to the LLM. Today a strict OpenAI-compatible
// upstream rejects that field with
// Unsupported parameter(s): maxOutputTokens`.
for (const c of captured) {
expect(c.opts.maxOutputTokens).toBeUndefined()
}
// Positive assertion (budget preserved through an independent path):
// the constrained model still threads a tightened output limit
// through to every worker. If a future "fix" accidentally severs
// the only budget source along with the leak, this fails.
for (const c of captured) {
expect(c.modelLimitOutput).toBeLessThanOrEqual(outputTokenMax)
}
} finally {
await rt.dispose()
}
},
})
},
30_000,
)
})
@@ -1,556 +0,0 @@
import { NodeFileSystem } from "@effect/platform-node"
import { afterEach, describe, expect } from "bun:test"
import { Effect, Exit, Fiber, Layer } from "effect"
import { FetchHttpClient } from "effect/unstable/http"
import fs from "fs/promises"
import { SessionV1 } from "@opencode-ai/core/v1/session"
import { Database } from "@opencode-ai/core/database/database"
import type { SessionID } from "../../src/session/schema"
import path from "path"
import { FSUtil } from "@opencode-ai/core/fs-util"
import { CrossSpawnSpawner } from "@opencode-ai/core/cross-spawn-spawner"
import * as Log from "@opencode-ai/core/util/log"
import { Agent as AgentSvc } from "../../src/agent/agent"
import { BackgroundJob } from "../../src/background/job"
import { Bus } from "../../src/bus"
import { Command } from "../../src/command"
import { Auth } from "../../src/auth"
import { Config } from "../../src/config/config"
import { RuntimeFlags } from "../../src/effect/runtime-flags"
import { EventV2Bridge } from "../../src/event-v2-bridge"
import { Env } from "../../src/env"
import { Format } from "../../src/format"
import { Git } from "../../src/git"
import { Image } from "../../src/image/image"
import { KiloSessions } from "../../src/kilo-sessions/kilo-sessions"
import { LSP } from "../../src/lsp/lsp"
import { MCP } from "../../src/mcp"
import { Permission } from "../../src/permission"
import { Plugin } from "../../src/plugin"
import { Provider as ProviderSvc } from "../../src/provider/provider"
import { Question } from "../../src/question"
import { RepositoryCache } from "@opencode-ai/core/repository-cache"
import { SessionCompaction } from "../../src/session/compaction"
import { Instruction } from "../../src/session/instruction"
import { LLM } from "../../src/session/llm"
import { SessionProcessor } from "../../src/session/processor"
import { SessionPrompt } from "../../src/session/prompt"
import { SessionRevert } from "../../src/session/revert"
import { SessionRunState } from "../../src/session/run-state"
import { Session } from "../../src/session/session"
import { SessionStatus } from "../../src/session/status"
import { SystemPrompt } from "../../src/session/system"
import { SessionSummary } from "../../src/session/summary"
import { Todo } from "../../src/session/todo"
import { Skill } from "../../src/skill"
import { Snapshot } from "../../src/snapshot"
import { Storage } from "../../src/storage/storage"
import { SyncEvent } from "../../src/sync"
import { Ripgrep } from "@opencode-ai/core/ripgrep"
import { ToolRegistry } from "../../src/tool/registry"
import { Truncate } from "../../src/tool/truncate"
import { MemoryService } from "@kilocode/kilo-memory/effect/service"
import { provideTmpdirServer } from "../fixture/fixture"
import { awaitWithTimeout, pollWithTimeout, testEffect } from "../lib/effect"
import { reply, TestLLMServer } from "../lib/llm-server"
void Log.init({ print: false })
afterEach(async () => {
// Dispose all test instances between integration scenarios.
const { disposeAllInstances } = await import("../fixture/fixture")
await disposeAllInstances()
})
const summary = Layer.succeed(
SessionSummary.Service,
SessionSummary.Service.of({
summarize: () => Effect.void,
diff: () => Effect.succeed([]),
computeDiff: () => Effect.succeed([]),
}),
)
const mcp = Layer.succeed(
MCP.Service,
MCP.Service.of({
status: () => Effect.succeed({}),
clients: () => Effect.succeed({}),
tools: () => Effect.succeed({}),
prompts: () => Effect.succeed({}),
resources: () => Effect.succeed({}),
add: () => Effect.succeed({ status: { status: "disabled" as const } }),
connect: () => Effect.void,
disconnect: () => Effect.void,
getPrompt: () => Effect.succeed(undefined),
readResource: () => Effect.succeed(undefined),
startAuth: () => Effect.die("unexpected MCP auth in watchdog tests"),
authenticate: () => Effect.die("unexpected MCP auth in watchdog tests"),
finishAuth: () => Effect.die("unexpected MCP auth in watchdog tests"),
removeAuth: () => Effect.void,
supportsOAuth: () => Effect.succeed(false),
hasStoredTokens: () => Effect.succeed(false),
getAuthStatus: () => Effect.succeed("not_authenticated" as const),
}),
)
const lsp = Layer.succeed(
LSP.Service,
LSP.Service.of({
init: () => Effect.void,
status: () => Effect.succeed([]),
hasClients: () => Effect.succeed(false),
touchFile: () => Effect.void,
diagnostics: () => Effect.succeed({}),
hover: () => Effect.succeed(undefined),
definition: () => Effect.succeed([]),
references: () => Effect.succeed([]),
implementation: () => Effect.succeed([]),
documentSymbol: () => Effect.succeed([]),
workspaceSymbol: () => Effect.succeed([]),
prepareCallHierarchy: () => Effect.succeed([]),
incomingCalls: () => Effect.succeed([]),
outgoingCalls: () => Effect.succeed([]),
}),
)
const status = Layer.mergeAll(SessionStatus.defaultLayer, Bus.layer)
const run = SessionRunState.layer.pipe(Layer.provide(status))
const infra = Layer.mergeAll(NodeFileSystem.layer, CrossSpawnSpawner.defaultLayer)
function makeHttp() {
const deps = Layer.mergeAll(
Session.defaultLayer,
BackgroundJob.defaultLayer,
Snapshot.defaultLayer,
LLM.defaultLayer,
Env.defaultLayer,
AgentSvc.defaultLayer,
Command.defaultLayer,
Permission.defaultLayer,
Plugin.defaultLayer,
Config.defaultLayer,
RuntimeFlags.layer(),
ProviderSvc.defaultLayer,
lsp,
mcp,
FSUtil.defaultLayer,
SyncEvent.defaultLayer,
EventV2Bridge.defaultLayer,
Database.defaultLayer,
status,
MemoryService.layer,
).pipe(Layer.provideMerge(infra))
const question = Question.layer.pipe(Layer.provideMerge(deps))
const todo = Todo.layer.pipe(Layer.provideMerge(deps))
const registry = ToolRegistry.layer.pipe(
Layer.provide(KiloSessions.testLayer),
Layer.provide(Skill.defaultLayer),
Layer.provide(FetchHttpClient.layer),
Layer.provide(CrossSpawnSpawner.defaultLayer),
Layer.provide(RepositoryCache.defaultLayer),
Layer.provide(Ripgrep.defaultLayer),
Layer.provide(Format.defaultLayer),
Layer.provide(Git.defaultLayer),
Layer.provide(Command.defaultLayer),
Layer.provide(Auth.defaultLayer),
Layer.provideMerge(todo),
Layer.provideMerge(question),
Layer.provideMerge(deps),
)
const trunc = Truncate.layer.pipe(Layer.provideMerge(deps))
const proc = SessionProcessor.layer.pipe(
Layer.provide(summary),
Layer.provide(Image.defaultLayer),
Layer.provideMerge(deps),
)
const compact = SessionCompaction.layer.pipe(Layer.provideMerge(proc), Layer.provideMerge(deps))
return Layer.mergeAll(
TestLLMServer.layer,
SessionPrompt.layer.pipe(
Layer.provide(SessionRevert.defaultLayer),
Layer.provide(Image.defaultLayer),
Layer.provide(summary),
Layer.provideMerge(run),
Layer.provideMerge(compact),
Layer.provideMerge(proc),
Layer.provideMerge(registry),
Layer.provideMerge(trunc),
Layer.provideMerge(question),
Layer.provide(Instruction.defaultLayer),
Layer.provide(SystemPrompt.defaultLayer),
Layer.provideMerge(deps),
),
).pipe(
Layer.provide(
Layer.mergeAll(
summary,
deps,
Config.defaultLayer,
RuntimeFlags.layer(),
BackgroundJob.defaultLayer,
Bus.layer,
infra,
Storage.defaultLayer,
),
),
)
}
const it = testEffect(makeHttp())
const cfg = {
provider: {
test: {
name: "Test",
id: "test",
env: [],
npm: "@ai-sdk/openai-compatible",
models: {
"test-model": {
id: "test-model",
name: "Test Model",
attachment: false,
reasoning: false,
temperature: false,
tool_call: true,
release_date: "2025-01-01",
limit: { context: 100000, output: 10000 },
cost: { input: 0, output: 0 },
options: { chunkTimeout: 1_000 },
},
},
options: {
apiKey: "test-key",
baseURL: "http://localhost:1/v1",
},
},
},
}
function providerCfg(url: string) {
return {
...cfg,
provider: {
...cfg.provider,
test: {
...cfg.provider.test,
options: {
...cfg.provider.test.options,
baseURL: url,
chunkTimeout: false as const,
},
},
},
}
}
const worktreeFile = (dir: string, name: string) => path.join(dir, name)
const exists = (file: string) =>
Effect.promise(() =>
fs
.access(file)
.then(() => true)
.catch(() => false),
)
// The production bash tool runs every command through a *login* shell
// (`bash -l -c ...`, see src/shell/shell.ts) so `~/.bashrc` and shell
// aliases behave the same as an interactive terminal. Git for Windows'
// login-shell startup rescans the full Windows `PATH`, which is slower
// than the Unix shells used elsewhere in this file. Give the marker file
// these tests poll for a little extra headroom there, on top of the tests
// A/C `config.shell: "bash"` override that makes the bash tool actually
// use git-bash instead of cmd.exe on Windows (see those config comments).
const waitForFile = (file: string, label: string, duration = process.platform === "win32" ? 15_000 : 5_000) =>
pollWithTimeout(
Effect.gen(function* () {
const ok = yield* exists(file)
return ok ? true : undefined
}),
label,
duration,
)
const touch = (file: string) => Effect.promise(() => fs.writeFile(file, ""))
const waitForRunningTool = (sessionID: SessionID, sessions: Session.Interface, label: string, duration = 15_000) =>
pollWithTimeout(
Effect.gen(function* () {
const msgs = yield* sessions.messages({ sessionID })
const running = msgs
.flatMap((msg) => msg.parts)
.find((part) => part.type === "tool" && part.state.status === "running")
return running ? running : undefined
}),
label,
duration,
)
const waitForRequestHit = (llm: TestLLMServer["Service"], needle: string, label: string) =>
pollWithTimeout(
Effect.gen(function* () {
const hits = yield* llm.hits
const matched = hits.filter((hit) => JSON.stringify(hit.body).includes(needle))
return matched.length > 0 ? matched : undefined
}),
label,
5_000,
)
const matchContains = (needle: string) => (hit: { body: Record<string, unknown> }) =>
JSON.stringify(hit.body).includes(needle)
const assertNotInterrupted = (parts: SessionV1.WithParts["parts"]) => {
for (const part of parts) {
if (part.type === "tool") {
expect(part.state.status).toBe("completed")
if (part.state.status === "completed") {
expect(part.state.metadata?.interrupted).not.toBe(true)
expect(part.state.output).not.toContain("Tool execution aborted")
}
}
}
}
// kilocode_change: normalize to forward slashes before embedding in the
// shell script. `ready`/`release` come from `path.join`, which yields
// backslash-separated paths on Windows; inside a double-quoted git-bash
// string a literal backslash is an escape character, so a Windows path can
// silently mangle into the wrong filename (or a path bash's `[ -f ... ]`
// test can't resolve) rather than throwing. Git-bash/MSYS accept
// forward-slash paths natively, so this is safe on every platform this
// suite runs on.
const posixPath = (p: string) => p.replaceAll("\\", "/")
const bashGate = (dir: string, ready: string, release: string) =>
`touch ${JSON.stringify(posixPath(ready))} && while [ ! -f ${JSON.stringify(posixPath(release))} ]; do sleep 0.05; done && echo done`
describe("session stream watchdog integration", () => {
it.live(
"A: root session long-running Bash is not interrupted by the idle watchdog",
() =>
provideTmpdirServer(
Effect.fnUntraced(function* ({ dir, llm }) {
const prompt = yield* SessionPrompt.Service
const sessions = yield* Session.Service
const chat = yield* sessions.create({ title: "Root long bash" })
const ready = worktreeFile(dir, "bash-ready")
const release = worktreeFile(dir, "bash-release")
yield* llm.tool("bash", {
command: bashGate(dir, ready, release),
description: "Long running bash command",
timeout: 60_000,
workdir: dir,
})
yield* llm.text("bash complete")
yield* prompt.prompt({
sessionID: chat.id,
agent: "build",
noReply: true,
parts: [{ type: "text", text: "run a long bash command" }],
})
const fiber = yield* prompt.loop({ sessionID: chat.id }).pipe(Effect.forkChild)
yield* llm.wait(1)
yield* waitForRunningTool(chat.id, sessions, "root bash tool never started")
yield* waitForFile(ready, "root bash readiness marker never appeared")
yield* Effect.sleep("1500 millis")
yield* touch(release)
const exit = yield* awaitWithTimeout(Fiber.await(fiber), "root bash loop did not finish", "15 seconds")
expect(Exit.isSuccess(exit)).toBe(true)
// Check all messages in the session for interrupted tools
const allMessages = yield* sessions.messages({ sessionID: chat.id })
for (const msg of allMessages) {
assertNotInterrupted(msg.parts)
}
}),
{
git: true,
// kilocode_change: without an explicit `shell`, the bash tool's
// `defaultShell()` falls back to cmd.exe on Windows (see
// packages/core/src/tool/bash.ts), which cannot run bashGate's
// POSIX syntax (`touch`, `[ -f ... ]`, `while ... done`). That
// made `touch` fail immediately and silently, so the readiness
// marker never appeared regardless of how long the test waited.
config: (url) => ({ ...providerCfg(url), shell: "bash", permission: { bash: "allow" } }),
},
),
{ timeout: 30_000 },
)
it.live(
"B: root foreground TaskTool child with held LLM response is not interrupted",
() =>
provideTmpdirServer(
Effect.fnUntraced(function* ({ dir, llm }) {
const prompt = yield* SessionPrompt.Service
const sessions = yield* Session.Service
const chat = yield* sessions.create({ title: "Root foreground child" })
const gate = Promise.withResolvers<void>()
yield* llm.tool("task", {
description: "Foreground child task",
prompt: "child task: say hello",
subagent_type: "child",
})
yield* llm.pushMatch(
matchContains("child task: say hello"),
reply().wait(gate.promise).text("child done").stop(),
)
yield* llm.text("parent done")
yield* prompt.prompt({
sessionID: chat.id,
agent: "build",
noReply: true,
parts: [{ type: "text", text: "run a foreground child" }],
})
const fiber = yield* prompt.loop({ sessionID: chat.id }).pipe(Effect.forkChild)
yield* llm.wait(1)
yield* waitForRunningTool(chat.id, sessions, "root task tool never started")
yield* waitForRequestHit(llm, "child task: say hello", "child LLM request never hit server")
yield* Effect.sleep("1500 millis")
gate.resolve(undefined)
const exit = yield* awaitWithTimeout(
Fiber.await(fiber),
"root foreground child loop did not finish",
"15 seconds",
)
expect(Exit.isSuccess(exit)).toBe(true)
// Check all messages in root and child sessions for interrupted tools
const allMessages = yield* sessions.messages({ sessionID: chat.id })
for (const msg of allMessages) {
assertNotInterrupted(msg.parts)
}
const children = yield* sessions.children(chat.id)
expect(children).toHaveLength(1)
const childMessages = yield* sessions.messages({ sessionID: children[0]!.id })
for (const msg of childMessages) {
assertNotInterrupted(msg.parts)
}
}),
{
git: true,
config: (url) => ({
...providerCfg(url),
permission: { bash: "allow", task: "allow" },
agent: {
child: {
model: "test/test-model",
mode: "subagent",
options: { chunkTimeout: false },
permission: { bash: "allow", task: "allow" },
},
},
}),
},
),
{ timeout: 30_000 },
)
it.live(
"C: child session long-running Bash while root awaits it is not interrupted",
() =>
provideTmpdirServer(
Effect.fnUntraced(function* ({ dir, llm }) {
const prompt = yield* SessionPrompt.Service
const sessions = yield* Session.Service
const chat = yield* sessions.create({ title: "Nested long bash" })
const ready = worktreeFile(dir, "child-bash-ready")
const release = worktreeFile(dir, "child-bash-release")
yield* llm.tool("task", {
description: "Nested child task",
prompt: "child task: run a long bash command",
subagent_type: "child",
})
yield* llm.pushMatch(
matchContains("child task: run a long bash command"),
reply().tool("bash", {
command: bashGate(dir, ready, release),
description: "Long running child bash command",
timeout: 60_000,
workdir: dir,
}),
)
yield* llm.text("child done")
yield* llm.text("parent done")
yield* prompt.prompt({
sessionID: chat.id,
agent: "build",
noReply: true,
parts: [{ type: "text", text: "run a nested child" }],
})
const fiber = yield* prompt.loop({ sessionID: chat.id }).pipe(Effect.forkChild)
yield* llm.wait(1)
yield* waitForRunningTool(chat.id, sessions, "root task tool never started")
yield* waitForRequestHit(llm, "child task: run a long bash command", "child LLM request never hit server")
const children = yield* sessions.children(chat.id)
expect(children).toHaveLength(1)
const childID = children[0]!.id
yield* waitForRunningTool(childID, sessions, "child bash tool never started")
yield* waitForFile(ready, "child bash readiness marker never appeared")
yield* Effect.sleep("1500 millis")
yield* touch(release)
const exit = yield* awaitWithTimeout(
Fiber.await(fiber),
"nested child bash loop did not finish",
"15 seconds",
)
expect(Exit.isSuccess(exit)).toBe(true)
// Check all messages in root and child sessions for interrupted tools
const rootMessages = yield* sessions.messages({ sessionID: chat.id })
for (const msg of rootMessages) {
assertNotInterrupted(msg.parts)
}
// Reuse childID captured above
const childMessages = yield* sessions.messages({ sessionID: childID })
for (const msg of childMessages) {
assertNotInterrupted(msg.parts)
}
}),
{
git: true,
// kilocode_change: see the matching comment on test A — without
// this, the nested child's bash tool falls back to cmd.exe on
// Windows and the readiness marker never appears.
config: (url) => ({
...providerCfg(url),
shell: "bash",
permission: { bash: "allow", task: "allow" },
agent: {
child: {
model: "test/test-model",
mode: "subagent",
permission: { bash: "allow", task: "allow" },
},
},
}),
},
),
{ timeout: 30_000 },
)
})
@@ -3,84 +3,36 @@ import { Effect, Stream } from "effect"
import { LLMEvent } from "@opencode-ai/llm"
import { KiloLLM } from "@/kilocode/session/llm"
describe("kilocode.session.llm.resolveIdleMs", () => {
describe("kilocode.session.llm.timeout", () => {
test("uses prepared options before the provider fallback", () => {
expect(
KiloLLM.resolveIdleMs({
options: { chunkTimeout: 15_000 },
fallback: { chunkTimeout: 30_000 },
}),
).toBe(15_000)
const result = KiloLLM.timeout({
options: { chunkTimeout: 15_000 },
fallback: { chunkTimeout: 30_000 },
})
expect(result).toEqual({ timeout: { chunkMs: 15_000 } })
})
test("uses the provider fallback when prepared options omit the timeout", () => {
expect(
KiloLLM.resolveIdleMs({
options: {},
fallback: { chunkTimeout: 30_000 },
}),
).toBe(30_000)
const result = KiloLLM.timeout({
options: {},
fallback: { chunkTimeout: 30_000 },
})
expect(result).toEqual({ timeout: { chunkMs: 30_000 } })
})
test("uses the provider fallback when the prepared value is not a number", () => {
expect(
KiloLLM.resolveIdleMs({
options: { chunkTimeout: "15_000" },
fallback: { chunkTimeout: 30_000 },
}),
).toBe(30_000)
const result = KiloLLM.timeout({
options: { chunkTimeout: "15_000" },
fallback: { chunkTimeout: 30_000 },
})
expect(result).toEqual({ timeout: { chunkMs: 30_000 } })
})
test("defaults the chunk idle timeout to 60_000 ms when no override is configured", () => {
expect(KiloLLM.resolveIdleMs({ options: {} })).toBe(60_000)
})
test("returns undefined when prepared is false (disabled)", () => {
expect(
KiloLLM.resolveIdleMs({
options: { chunkTimeout: false },
fallback: { chunkTimeout: 30_000 },
}),
).toBeUndefined()
})
test("returns undefined when prepared is 0 (internal disable)", () => {
expect(
KiloLLM.resolveIdleMs({
options: { chunkTimeout: 0 },
fallback: { chunkTimeout: 30_000 },
}),
).toBeUndefined()
})
test("returns undefined when provider fallback is false", () => {
expect(
KiloLLM.resolveIdleMs({
options: {},
fallback: { chunkTimeout: false },
}),
).toBeUndefined()
})
test("falls through invalid prepared values to provider fallback", () => {
expect(
KiloLLM.resolveIdleMs({
options: { chunkTimeout: -1 },
fallback: { chunkTimeout: 5_000 },
}),
).toBe(5_000)
expect(
KiloLLM.resolveIdleMs({
options: { chunkTimeout: Number.POSITIVE_INFINITY },
fallback: { chunkTimeout: 5_000 },
}),
).toBe(5_000)
expect(
KiloLLM.resolveIdleMs({
options: { chunkTimeout: Number.NaN },
fallback: { chunkTimeout: 5_000 },
}),
).toBe(5_000)
test("omits the timeout when it is not configured", () => {
expect(KiloLLM.timeout({ options: {} })).toEqual({})
})
})
@@ -1,282 +0,0 @@
import { describe, expect, test } from "bun:test"
import { Effect, Stream } from "effect"
import type { LanguageModelV2CallWarning, LanguageModelV2StreamPart } from "@ai-sdk/provider"
import { KiloLLM } from "@/kilocode/session/llm"
import { ProviderError } from "@/provider/error"
type FullStreamPart = LanguageModelV2StreamPart
function part(type: string, extra: Record<string, unknown> = {}): FullStreamPart {
return { type, ...extra } as unknown as FullStreamPart
}
async function run<T>(eff: Effect.Effect<T, unknown>) {
return await Effect.runPromise(eff)
}
function fromSchedule(events: Array<[number, FullStreamPart]>, end: number): Stream.Stream<FullStreamPart, never> {
// Each `[at, value]` is an ABSOLUTE time in milliseconds from stream start.
// This matches how the tests are written: post-tool events (finish-step,
// finish) are scheduled within a short idle window of the tool-result so
// the watchdog, after the local active set drains, still receives the next
// event in time.
return Stream.fromAsyncIterable(
(async function* () {
const start = Date.now()
for (const [at, value] of events) {
const wait = at - (Date.now() - start)
if (wait > 0) await new Promise((r) => setTimeout(r, wait))
yield value
}
await new Promise((r) => setTimeout(r, end))
})(),
(e) => e as never,
)
}
describe("kilocode.session.llm.resolveIdleMs", () => {
test("returns prepared positive finite value as-is", () => {
const out = KiloLLM.resolveIdleMs({ options: { chunkTimeout: 15_000 }, fallback: { chunkTimeout: 30_000 } })
expect(out).toBe(15_000)
})
test("falls back to provider value when prepared is missing", () => {
const out = KiloLLM.resolveIdleMs({ options: {}, fallback: { chunkTimeout: 30_000 } })
expect(out).toBe(30_000)
})
test("falls back to provider value when prepared is a non-number string", () => {
const out = KiloLLM.resolveIdleMs({
options: { chunkTimeout: "15_000" },
fallback: { chunkTimeout: 30_000 },
})
expect(out).toBe(30_000)
})
test("falls back to provider value when prepared is negative", () => {
const out = KiloLLM.resolveIdleMs({
options: { chunkTimeout: -1 },
fallback: { chunkTimeout: 30_000 },
})
expect(out).toBe(30_000)
})
test("falls back to provider value when prepared is non-finite (Infinity, NaN)", () => {
expect(
KiloLLM.resolveIdleMs({ options: { chunkTimeout: Number.POSITIVE_INFINITY }, fallback: { chunkTimeout: 5_000 } }),
).toBe(5_000)
expect(KiloLLM.resolveIdleMs({ options: { chunkTimeout: Number.NaN }, fallback: { chunkTimeout: 5_000 } })).toBe(
5_000,
)
})
test("treats boolean false as a request to disable the watchdog", () => {
expect(
KiloLLM.resolveIdleMs({ options: { chunkTimeout: false }, fallback: { chunkTimeout: 30_000 } }),
).toBeUndefined()
})
test("treats internal 0 as a request to disable the watchdog", () => {
expect(KiloLLM.resolveIdleMs({ options: { chunkTimeout: 0 }, fallback: { chunkTimeout: 30_000 } })).toBeUndefined()
})
test("provider fallback false also disables", () => {
expect(KiloLLM.resolveIdleMs({ options: {}, fallback: { chunkTimeout: false } })).toBeUndefined()
})
test("uses 60_000 default when nothing valid is configured", () => {
expect(KiloLLM.resolveIdleMs({ options: {} })).toBe(60_000)
})
test("uses 60_000 default when both prepared and fallback are invalid", () => {
expect(
KiloLLM.resolveIdleMs({
options: { chunkTimeout: "x" },
fallback: { chunkTimeout: -5 },
}),
).toBe(60_000)
})
})
describe("kilocode.session.llm.watchdogStream", () => {
test("returns the stream unchanged when idle is undefined (disabled)", async () => {
const events: FullStreamPart[] = [
part("stream-start", { warnings: [] as LanguageModelV2CallWarning[] }),
part("text-delta", { id: "t1", delta: "ok" }),
]
const out = await run(Stream.runCollect(KiloLLM.watchdogStream(Stream.fromIterable(events), undefined)))
expect(out.length).toBe(2)
})
test("emits events and completes when the stream delivers them within the idle window", async () => {
const events: FullStreamPart[] = [
part("stream-start", { warnings: [] as LanguageModelV2CallWarning[] }),
part("text-delta", { id: "t1", delta: "hi" }),
part("text-delta", { id: "t1", delta: "!" }),
]
const out = await run(Stream.runCollect(KiloLLM.watchdogStream(Stream.fromIterable(events), 1_000)))
expect(out.length).toBe(3)
})
test("fails with ProviderError.ResponseStreamError when the stream stalls", async () => {
const slow = Stream.fromEffect(
Effect.flatMap(Effect.sleep("5 seconds"), () => Effect.succeed(part("text-delta", { id: "t1", delta: "x" }))),
)
const wrapped = KiloLLM.watchdogStream(slow, 100)
const err = await run(Effect.flip(Stream.runCollect(wrapped)))
expect(err).toBeInstanceOf(ProviderError.ResponseStreamError)
})
test("every raw AI SDK event resets the idle timer", async () => {
// idle 200ms; emit text-delta at 0 and 60ms (a single 260ms pull would time out without reset).
const stream = fromSchedule(
[
[0, part("text-delta", { id: "t1", delta: "a" })],
[60, part("text-delta", { id: "t1", delta: "b" })],
],
10,
)
const out = await run(Stream.runCollect(KiloLLM.watchdogStream(stream, 200)))
expect(out.length).toBe(2)
})
test("pending local tool calls suspend the idle timeout until they settle", async () => {
// tool-call at t=0 (local). 250ms quiet gap then tool-result. A healthy AI
// SDK run also emits finish-step + finish right after the tool-result, so
// the watchdog sees another event within idleMs and resets.
const stream = fromSchedule(
[
[0, part("tool-call", { toolCallId: "c1", toolName: "bash" })],
[250, part("tool-result", { toolCallId: "c1", toolName: "bash", output: "ok" })],
[260, part("finish-step", { finishReason: "tool-calls" })],
[270, part("finish", { finishReason: "stop" })],
],
10,
)
const out = await run(Stream.runCollect(KiloLLM.watchdogStream(stream, 200)))
expect(out.length).toBe(4)
})
test("provider-executed tool calls do not suspend the watchdog", async () => {
const stream = fromSchedule(
[[0, part("tool-call", { toolCallId: "c1", toolName: "web", providerExecuted: true })]],
400,
)
const err = await run(Effect.flip(Stream.runCollect(KiloLLM.watchdogStream(stream, 200))))
expect(err).toBeInstanceOf(ProviderError.ResponseStreamError)
})
test("parallel local tool calls remain suspended until the last settles", async () => {
const stream = fromSchedule(
[
[0, part("tool-call", { toolCallId: "a", toolName: "bash" })],
[10, part("tool-call", { toolCallId: "b", toolName: "bash" })],
[200, part("tool-result", { toolCallId: "a", toolName: "bash", output: "x" })],
[350, part("tool-result", { toolCallId: "b", toolName: "bash", output: "y" })],
[360, part("finish-step", { finishReason: "tool-calls" })],
[370, part("finish", { finishReason: "stop" })],
],
10,
)
const out = await run(Stream.runCollect(KiloLLM.watchdogStream(stream, 200)))
expect(out.length).toBe(6)
})
test("tool-error for a local tool id also releases the suspension", async () => {
const stream = fromSchedule(
[
[0, part("tool-call", { toolCallId: "c1", toolName: "bash" })],
[200, part("tool-error", { toolCallId: "c1", toolName: "bash", error: new Error("nope") })],
[210, part("finish-step", { finishReason: "tool-calls" })],
[220, part("finish", { finishReason: "stop" })],
],
10,
)
const out = await run(Stream.runCollect(KiloLLM.watchdogStream(stream, 200)))
expect(out.length).toBe(4)
})
test("aborts the underlying source on timeout so cleanup does not hang", async () => {
const ctrl = new AbortController()
let abortReason: unknown
let nextResolved = false
const source: AsyncIterable<FullStreamPart> = {
[Symbol.asyncIterator]() {
let nextPromise: Promise<IteratorResult<FullStreamPart>> | undefined
let resolveNext: ((value: IteratorResult<FullStreamPart>) => void) | undefined
ctrl.signal.addEventListener("abort", () => {
abortReason = ctrl.signal.reason
if (resolveNext) {
resolveNext({ done: true, value: undefined })
nextResolved = true
}
})
return {
next() {
nextPromise = new Promise((resolve) => {
resolveNext = resolve
})
return nextPromise
},
async return() {
if (nextPromise) await nextPromise
return { done: true, value: undefined }
},
}
},
}
const wrapped = KiloLLM.watchdogAsyncIterable(source, 100, ctrl)
const err = await run(Effect.flip(Stream.runCollect(Stream.fromAsyncIterable(wrapped, (e) => e as never))))
expect(err).toBeInstanceOf(ProviderError.ResponseStreamError)
expect(nextResolved).toBe(true)
expect(abortReason).toBeInstanceOf(ProviderError.ResponseStreamError)
})
test("propagates upstream stream errors without false timeout", async () => {
const stream = Stream.fail(new Error("upstream broken"))
const err = await run(Effect.flip(Stream.runCollect(KiloLLM.watchdogStream(stream, 1_000))))
expect((err as Error).message).toBe("upstream broken")
})
test("return() closes the source immediately without waiting on a stalled pull", async () => {
// Regression test: a hand-rolled async generator's `.return()` cannot
// preempt an in-flight internal `await` — it only takes effect once that
// await settles on its own, which never happens for a genuinely stalled
// source. `watchdogAsyncIterable` must instead expose a `return()` that
// runs immediately and forwards to the source's `return()` without
// waiting for the outstanding `next()` to resolve.
let sourceReturnCalled = false
let neverResolvingNextCalled = false
const source: AsyncIterable<FullStreamPart> = {
[Symbol.asyncIterator]() {
return {
next() {
neverResolvingNextCalled = true
return new Promise<IteratorResult<FullStreamPart>>(() => {
// Never resolves — simulates a fully stalled source (e.g. a
// hung fetch response) whose pending pull is abandoned once
// the consumer decides to stop.
})
},
async return() {
sourceReturnCalled = true
return { done: true, value: undefined }
},
}
},
}
const wrapped = KiloLLM.watchdogAsyncIterable(source, 60_000)
const it = wrapped[Symbol.asyncIterator]()
const pending = it.next()
expect(neverResolvingNextCalled).toBe(true)
const returned = await Promise.race([
it.return!(),
new Promise((_, reject) => setTimeout(() => reject(new Error("return() hung")), 500)),
])
expect(returned).toMatchObject({ done: true })
expect(sourceReturnCalled).toBe(true)
// The abandoned pull is left unresolved; only return() is asserted here.
void pending
})
})
@@ -769,4 +769,248 @@ describe("AttachedState", () => {
expect(calls).toEqual([{}, { requireSessionId: "ses_b" }])
})
// K1 W1: detach semantics — basic happy path.
test("detach removes the id from both sets and awaits a heartbeat whose payload no longer contains it", async () => {
let detachResolved = false
const state = AttachedState.create({
heartbeat: (opts) => {
if (opts?.detachSessionId) {
detachResolved = true
return Promise.resolve()
}
return Promise.resolve()
},
log: nolog,
})
state.setPresence(["ses_a"])
await Promise.resolve()
expect(state.has("ses_a")).toBe(true)
// Detach awaits a heartbeat whose payload no longer contains ses_a.
// The state machine removes the id synchronously before awaiting.
await state.detach("ses_a")
expect(detachResolved).toBe(true)
expect([...state.union()]).toEqual([])
})
// K1 W1: detach surfaces a specific error for an id this CLI does not own.
test("detach throws for an id this CLI does not own (no silent re-attach)", async () => {
const state = AttachedState.create({
heartbeat: () => Promise.resolve(),
log: nolog,
})
await expect(state.detach("ses_missing")).rejects.toThrow("not owned")
})
// K1 W1: heartbeat failure during detach rolls back by restoring ownership.
test("detach rolls back by restoring prior ownership on heartbeat failure", async () => {
const state = AttachedState.create({
heartbeat: (opts) => {
if (opts?.detachSessionId) return Promise.reject(new Error("relay down"))
return Promise.resolve()
},
log: nolog,
})
state.setPresence(["ses_a"])
await Promise.resolve()
await expect(state.detach("ses_a")).rejects.toThrow("relay down")
// The id must be back in presence so a future setPresence does not
// accidentally treat the session as detached.
expect(state.has("ses_a")).toBe(true)
})
// K1 W1: suppression tombstone prevents a presence replacement that
// still includes a just-exited id from instantly re-adopting it.
test("setPresence does not re-adopt a detached id while presence still reports it", async () => {
const state = AttachedState.create({
heartbeat: () => Promise.resolve(),
log: nolog,
})
state.setPresence(["ses_a", "ses_b"])
expect(state.has("ses_a")).toBe(true)
// Detach ses_a; the tombstone is set BEFORE the sets are mutated.
await state.detach("ses_a")
expect(state.has("ses_a")).toBe(false)
// A presence churn that still includes ses_a must NOT re-adopt it
// (the relay is the source of truth and the upstream side has not
// dropped the id yet).
state.setPresence(["ses_a", "ses_b"])
expect(state.has("ses_a")).toBe(false)
// Once presence genuinely drops ses_a, the tombstone is released
// and a later real re-open (via announce) is not blocked.
state.setPresence(["ses_b"])
expect(state.has("ses_a")).toBe(false)
await state.announce("ses_a")
expect(state.has("ses_a")).toBe(true)
})
// K1 W1: has(id) reflects presence pending.
test("has(id) is true for presence-owned and pending ids, false otherwise", async () => {
const announced = Promise.withResolvers<void>()
const state = AttachedState.create({
heartbeat: (opts) => {
if (opts?.requireSessionId === "ses_pending") return announced.promise
return Promise.resolve()
},
log: nolog,
})
state.setPresence(["ses_present"])
expect(state.has("ses_present")).toBe(true)
expect(state.has("ses_pending")).toBe(false)
expect(state.has("ses_other")).toBe(false)
// Announce with a held heartbeat so the id sits in pending.
const p = state.announce("ses_pending")
await Promise.resolve()
expect(state.has("ses_pending")).toBe(true)
announced.resolve()
await p
expect(state.has("ses_pending")).toBe(true)
})
// K1 W1: reset() also clears the detach in-flight map and tombstones
// so a new connection lifecycle does not inherit stale state.
test("reset() clears tombstones and detach in-flight map", async () => {
const state = AttachedState.create({
heartbeat: () => Promise.resolve(),
log: nolog,
})
state.setPresence(["ses_a"])
await state.detach("ses_a")
state.setPresence(["ses_a"])
expect(state.has("ses_a")).toBe(false) // tombstone held
state.reset()
// After reset, a presence report including ses_a is accepted (the
// previous tombstone is gone).
state.setPresence(["ses_a"])
expect(state.has("ses_a")).toBe(true)
})
// K1 W1: a detach in flight for an id must NOT cause a concurrent
// announce(id) to join the detach fence and report a bogus attach. The
// announce must wait for the detach to settle and then genuinely re-attach.
test("announce awaits an in-flight detach and then really re-attaches (no opposite-op join)", async () => {
const detachHb = Promise.withResolvers<void>()
const calls: Array<{ requireSessionId?: string; detachSessionId?: string }> = []
const state = AttachedState.create({
heartbeat: (opts) => {
calls.push(opts ?? {})
if (opts?.detachSessionId === "ses_y") return detachHb.promise
return Promise.resolve()
},
log: nolog,
})
state.setPresence(["ses_y"])
await Promise.resolve()
const detachP = state.detach("ses_y") // holds on the detach fence, id removed
const announceP = state.announce("ses_y") // must await the detach, not join it
detachHb.resolve()
await detachP
await announceP
// The announce genuinely re-attached rather than resolving on the detach's
// "id absent" outcome, and it drove a real requireSessionId heartbeat.
expect(state.has("ses_y")).toBe(true)
expect(calls.some((c) => c.requireSessionId === "ses_y")).toBe(true)
})
// K1 W1: an announce in flight for an id must NOT cause a concurrent
// detach(id) to join the announce and report a bogus detach — exit_cli
// treats a resolved detach as license to ACK/close, so a false success is
// dangerous. The detach must wait for the announce, then really detach.
test("detach awaits an in-flight announce and then really detaches (no opposite-op join)", async () => {
const announceHb = Promise.withResolvers<void>()
const calls: Array<{ requireSessionId?: string; detachSessionId?: string }> = []
const state = AttachedState.create({
heartbeat: (opts) => {
calls.push(opts ?? {})
if (opts?.requireSessionId === "ses_x") return announceHb.promise
return Promise.resolve()
},
log: nolog,
})
const announceP = state.announce("ses_x") // holds on the attach fence
const detachP = state.detach("ses_x") // must await the announce, not join it
announceHb.resolve()
await announceP
await detachP
// The detach genuinely ran the negative-containment fence rather than
// resolving on the announce's success; the session is actually gone.
expect(state.has("ses_x")).toBe(false)
expect(calls.some((c) => c.detachSessionId === "ses_x")).toBe(true)
})
// K1 W1: after a failed detach rolls ownership back, the id is genuinely
// still attached, so the very next presence report that still includes it
// must keep it — the tombstone must have been released on rollback.
test("failed-detach rollback keeps the id attached across the next setPresence", async () => {
const state = AttachedState.create({
heartbeat: (opts) => {
if (opts?.detachSessionId) return Promise.reject(new Error("relay down"))
return Promise.resolve()
},
log: nolog,
})
state.setPresence(["ses_a"])
await Promise.resolve()
await expect(state.detach("ses_a")).rejects.toThrow("relay down")
expect(state.has("ses_a")).toBe(true)
// The realistic next presence event still reports ses_a. Without releasing
// the tombstone on rollback, setPresence's suppression loop would drop the
// still-attached id here and never clear the tombstone.
state.setPresence(["ses_a"])
expect(state.has("ses_a")).toBe(true)
})
// K1 W1: reset() clears the SAME set instances, so a stale in-flight
// announce that rejects after a reconnect must NOT roll back into the new
// lifecycle — doing so would delete a fresh post-reset announce's pending
// entry. The catch must honor the generation guard like the success path.
test("a stale announce rejecting after reset() does not corrupt the new lifecycle's pending set", async () => {
const hb1 = Promise.withResolvers<void>()
const hb2 = Promise.withResolvers<void>()
let calls = 0
const state = AttachedState.create({
heartbeat: (opts) => {
if (opts?.requireSessionId === "id") {
calls += 1
return calls === 1 ? hb1.promise : hb2.promise
}
return Promise.resolve()
},
log: nolog,
})
const a1 = state.announce("id") // installs pending, awaits hb1
void a1.then(
() => {},
() => {},
)
await Promise.resolve()
state.reset() // bumps generation, clears the (same) sets
const a2 = state.announce("id") // fresh lifecycle: re-installs pending, awaits hb2
await Promise.resolve()
hb1.reject(new Error("stale relay drop")) // the dead-lifecycle announce fails
await Promise.resolve()
await Promise.resolve()
// The stale rollback must NOT have deleted the fresh generation's entry.
expect(state.has("id")).toBe(true)
hb2.resolve()
await a2
expect(state.has("id")).toBe(true)
})
})
@@ -104,6 +104,7 @@ describe("RemoteCommand", () => {
subtask: true,
},
],
canExitSession: true,
})
expect(JSON.stringify(catalog)).not.toContain("template")
expect(JSON.stringify(catalog)).not.toContain("secret-skill")
@@ -115,6 +116,9 @@ describe("RemoteCommand", () => {
{ name: "alpha", source: "command", hints: [], template: "alpha" },
])
expect(base.commands.map((item) => item.name)).toEqual(["alpha", "beta", "compact"])
// kilocode_change - K1 W1: canExitSession is always true, independent of
// exitAvailable (which gates the synthetic `/exit` entry).
expect(base.canExitSession).toBe(true)
const catalog = RemoteCommand.build(
[
@@ -155,16 +159,26 @@ describe("RemoteCommand", () => {
compaction: { create: async () => {} },
prompt: { loop: async () => {} },
})
expect((await remote.list()).commands.some((item) => item.name === "exit")).toBe(false)
// kilocode_change - K1 W1: canExitSession is true even when the synthetic
// `/exit` entry is absent (e.g. a headless `kilo remote` host has no
// RemoteExit callback, so `/exit` is gated off — but the host still
// interprets `exit_cli` as session-detach).
const baseList = await remote.list()
expect(baseList.canExitSession).toBe(true)
expect(baseList.commands.some((item) => item.name === "exit")).toBe(false)
const unregister = RemoteExit.register(async () => {})
try {
expect((await remote.list()).commands.some((item) => item.name === "exit")).toBe(true)
const list = await remote.list()
expect(list.commands.some((item) => item.name === "exit")).toBe(true)
expect(list.canExitSession).toBe(true)
} finally {
unregister()
}
expect((await remote.list()).commands.some((item) => item.name === "exit")).toBe(false)
const after = await remote.list()
expect(after.commands.some((item) => item.name === "exit")).toBe(false)
expect(after.canExitSession).toBe(true)
})
test("keeps compact and exit within command and byte caps", () => {
@@ -266,6 +266,145 @@ describe("RemoteProtocol", () => {
}
})
// kilocode_change - K1 W1: instance advertisement + per-session platform
test("heartbeat without instance still parses (legacy compatibility)", () => {
const msg = { type: "heartbeat", sessions: [{ id: "ses_1", status: "busy", title: "Fix auth" }] }
const result = RemoteProtocol.Heartbeat.safeParse(msg)
expect(result.success).toBe(true)
if (result.success) {
expect(result.data.instance).toBeUndefined()
}
})
test("heartbeat round-trips instance advertisement", () => {
const msg = {
type: "heartbeat",
sessions: [],
instance: { name: "mbp-igor", projectName: "cloud", version: "1.2.3" },
}
const result = RemoteProtocol.Heartbeat.safeParse(msg)
expect(result.success).toBe(true)
if (result.success) {
expect(result.data.instance).toEqual({ name: "mbp-igor", projectName: "cloud", version: "1.2.3" })
}
// round-trip via JSON
const json = JSON.parse(JSON.stringify(result.success ? result.data : null))
const result2 = RemoteProtocol.Heartbeat.safeParse(json)
expect(result2.success).toBe(true)
if (result2.success) {
expect(result2.data.instance).toEqual({ name: "mbp-igor", projectName: "cloud", version: "1.2.3" })
}
})
test("instance advertisement version is optional", () => {
const msg = {
type: "heartbeat",
sessions: [],
instance: { name: "h", projectName: "p" },
}
const result = RemoteProtocol.Heartbeat.safeParse(msg)
expect(result.success).toBe(true)
if (result.success) {
expect(result.data.instance?.version).toBeUndefined()
}
})
test("instance advertisement rejects empty name", () => {
const msg = {
type: "heartbeat",
sessions: [],
instance: { name: "", projectName: "p" },
}
const result = RemoteProtocol.Heartbeat.safeParse(msg)
expect(result.success).toBe(false)
})
test("instance advertisement rejects oversized name", () => {
const msg = {
type: "heartbeat",
sessions: [],
instance: { name: "x".repeat(65), projectName: "p" },
}
const result = RemoteProtocol.Heartbeat.safeParse(msg)
expect(result.success).toBe(false)
})
test("instance advertisement rejects oversized projectName", () => {
const msg = {
type: "heartbeat",
sessions: [],
instance: { name: "h", projectName: "p".repeat(65) },
}
const result = RemoteProtocol.Heartbeat.safeParse(msg)
expect(result.success).toBe(false)
})
test("instance advertisement rejects oversized version", () => {
const msg = {
type: "heartbeat",
sessions: [],
instance: { name: "h", projectName: "p", version: "v".repeat(33) },
}
const result = RemoteProtocol.Heartbeat.safeParse(msg)
expect(result.success).toBe(false)
})
test("session info accepts optional platform", () => {
const msg = {
type: "heartbeat",
sessions: [{ id: "s1", status: "busy", title: "t", platform: "vscode" }],
}
const result = RemoteProtocol.Heartbeat.safeParse(msg)
expect(result.success).toBe(true)
if (result.success) {
expect(result.data.sessions[0].platform).toBe("vscode")
}
})
test("session info platform optional (legacy)", () => {
const msg = {
type: "heartbeat",
sessions: [{ id: "s1", status: "busy", title: "t" }],
}
const result = RemoteProtocol.Heartbeat.safeParse(msg)
expect(result.success).toBe(true)
if (result.success) {
expect(result.data.sessions[0].platform).toBeUndefined()
}
})
test("session info rejects oversized platform", () => {
const msg = {
type: "heartbeat",
sessions: [{ id: "s1", status: "busy", title: "t", platform: "p".repeat(33) }],
}
const result = RemoteProtocol.Heartbeat.safeParse(msg)
expect(result.success).toBe(false)
})
test("full heartbeat round-trips sessions + instance", () => {
const msg = {
type: "heartbeat",
protocolVersion: "1.0.0",
sessions: [
{ id: "ses_1", status: "busy", title: "Fix auth", platform: "cli" },
{ id: "ses_2", status: "idle", title: "Sub task", parentSessionId: "ses_1", platform: "vscode" },
],
instance: { name: "mbp-igor", projectName: "cloud", version: "1.2.3" },
}
const json = JSON.parse(JSON.stringify(msg))
const result = RemoteProtocol.Heartbeat.safeParse(json)
expect(result.success).toBe(true)
if (result.success) {
expect(result.data.sessions).toHaveLength(2)
expect(result.data.sessions[0].platform).toBe("cli")
expect(result.data.sessions[1].platform).toBe("vscode")
expect(result.data.instance).toEqual({ name: "mbp-igor", projectName: "cloud", version: "1.2.3" })
expect(result.data.protocolVersion).toBe("1.0.0")
}
})
test("heartbeat without capabilities parses", () => {
const result = RemoteProtocol.Heartbeat.safeParse({
type: "heartbeat",
@@ -2618,6 +2618,14 @@ describe("RemoteSender slash commands", () => {
test("exit_cli rejects invalid, missing, unresolved, and unavailable sessions before ACK", async () => {
const { conn, sent } = fakeConn()
const lookups: string[] = []
// kilocode_change - K1 W1: the new exit_cli handler requires hasSession
// (owns-check) + detachSession + cancelPrompt + ownedCount seams. The
// default test seam has hasSession=false and detachSession resolves, so
// the only path that completes is "not owned" — matching the new
// contract. The previous "graceful exit unavailable" branch only fired
// for an UNREGISTERED remoteExit on an OWNED id; in the K1 W1 design
// the headless case (no remoteExit) is no longer a separate error
// path — a headless host simply stays alive.
const sender = RemoteSender.create({
conn,
directory: "/tmp/process-default",
@@ -2631,6 +2639,10 @@ describe("RemoteSender slash commands", () => {
},
children: async () => [],
},
hasSession: () => false,
detachSession: async () => {},
ownedCount: () => 0,
cancelPrompt: async () => {},
remoteExit: {
get: () => undefined,
},
@@ -2691,10 +2703,10 @@ describe("RemoteSender slash commands", () => {
{ type: "response", id: "req_exit_invalid_session", error: "invalid exit_cli command" },
{ type: "response", id: "req_exit_bad_protocol", error: "invalid exit_cli command" },
{ type: "response", id: "req_exit_extra", error: "invalid exit_cli command" },
{ type: "response", id: "req_exit_missing", error: "failed to exit CLI" },
{ type: "response", id: "req_exit_unavailable", error: "graceful exit unavailable" },
{ type: "response", id: "req_exit_missing", error: "session not owned by this CLI" },
{ type: "response", id: "req_exit_unavailable", error: "session not owned by this CLI" },
])
expect(lookups).toEqual(["ses_missing", "ses_current"])
expect(lookups).toEqual([])
})
test("exit_cli ACKs before invoking the worker callback in a microtask", async () => {
@@ -2727,6 +2739,13 @@ describe("RemoteSender slash commands", () => {
get: async (id) => info(id),
children: async () => [],
},
// kilocode_change - K1 W1: owns the target so the new detach path
// runs; no other sessions remain (ownedCount=0) and the callback is
// registered, so the exit path completes and the microtask fires.
hasSession: () => true,
detachSession: async () => {},
ownedCount: () => 0,
cancelPrompt: async () => {},
remoteExit,
})
@@ -2772,6 +2791,12 @@ describe("RemoteSender slash commands", () => {
get: async (id) => info(id),
children: async () => [],
},
// kilocode_change - K1 W1: owns both targets; zero remaining after
// each detach; registered callback; the exit path completes.
hasSession: () => true,
detachSession: async () => {},
ownedCount: () => 0,
cancelPrompt: async () => {},
remoteExit: {
get: () => exit,
},
@@ -2823,6 +2848,12 @@ describe("RemoteSender slash commands", () => {
get: async (id) => info(id),
children: async () => [],
},
// kilocode_change - K1 W1: owns the target; zero remaining; the
// callback is the throwing one above.
hasSession: () => true,
detachSession: async () => {},
ownedCount: () => 0,
cancelPrompt: async () => {},
remoteExit: {
get: () => async () => {
throw new CredentialLeakError("token=must-not-leak")
@@ -2847,11 +2878,11 @@ describe("RemoteSender slash commands", () => {
expect(JSON.stringify(logs)).not.toContain("token=")
})
test("create_session creates a root session in the current directory and responds in order", async () => {
test("create_session creates a root session in the current directory, attaches in-process, and responds in order", async () => {
const { conn, sent } = fakeConn()
const dirs: string[] = []
const createCalls: { input: unknown; calls: number } = { input: undefined, calls: 0 }
const attachCalls: string[] = []
const attachCalls: SessionID[] = []
const order: string[] = []
const sender = RemoteSender.create({
conn,
@@ -2872,18 +2903,12 @@ describe("RemoteSender slash commands", () => {
return { id: SessionID.make("ses_new"), directory: "/workspace/project-a", parentID: undefined } as any
},
},
attachSession: async (id) => {
attachCalls.push(id)
attachSession: async (input) => {
attachCalls.push(input)
order.push("attach")
// The production attachSession is responsible for the heartbeat; the
// mock follows the same contract so the ordering assertion below
// exercises the real shape of: create -> attach -> heartbeat -> response.
await (conn as any).heartbeat()
return
},
})
;(conn as any).heartbeat = async () => {
order.push("heartbeat")
}
const response = expectResponse(conn, sent, "req_create")
sender.handle({
@@ -2899,36 +2924,36 @@ describe("RemoteSender slash commands", () => {
expect(dirs).toEqual(["/workspace/project-a"])
expect(createCalls.calls).toBe(1)
expect(createCalls.input).toEqual({})
expect(attachCalls).toEqual(["ses_new"])
expect(order).toEqual(["create", "attach", "heartbeat"])
expect(attachCalls).toEqual([SessionID.make("ses_new")])
expect(order).toEqual(["create", "attach"])
expect(sent).toEqual([{ type: "response", id: "req_create", result: { protocolVersion: 1, sessionID: "ses_new" } }])
})
test("create_session rejects unsupported protocol versions and missing or invalid session IDs", async () => {
test("create_session rejects unsupported protocol versions, extra fields, and invalid session IDs; absent sessionId is allowed", async () => {
const { conn, sent } = fakeConn()
const createCalls: unknown[] = []
const attachCalls: unknown[] = []
const sender = RemoteSender.create({
conn,
directory: "/tmp/test",
directory: "/tmp/process-default",
log: nolog,
subscribe: fakeBus().subscribe,
session: {
get: async () => {
throw new Error("must not look up session for invalid request")
get: async (sessionID) => {
// Used only for the absent-sessionId path; not reached for invalid ids.
return { id: sessionID, directory: "/tmp/process-default" } as any
},
children: async () => [],
create: async (input) => {
createCalls.push(input)
return { id: SessionID.make("ses_unused") } as any
return { id: SessionID.make("ses_unused"), directory: "/tmp/process-default" } as any
},
},
attachSession: async () => {
throw new Error("must not attach for invalid request")
attachSession: async (input) => {
attachCalls.push(input)
return
},
})
;(conn as any).heartbeat = async () => {
throw new Error("must not heartbeat for invalid request")
}
sender.handle({
type: "command",
@@ -2937,12 +2962,6 @@ describe("RemoteSender slash commands", () => {
sessionId: "ses_current",
data: { protocolVersion: 2 },
})
sender.handle({
type: "command",
id: "req_no_session",
command: "create_session",
data: { protocolVersion: 1 },
})
sender.handle({
type: "command",
id: "req_bad_session",
@@ -2957,20 +2976,37 @@ describe("RemoteSender slash commands", () => {
sessionId: "ses_current",
data: { protocolVersion: 1, extra: true },
})
// Absent sessionId: must NOT be rejected; should reach the spawn path.
const noSession = expectResponse(conn, sent, "req_no_session")
sender.handle({
type: "command",
id: "req_no_session",
command: "create_session",
data: { protocolVersion: 1 },
})
await noSession.promise
noSession.restore()
expect(sent).toEqual([
expect(sent.slice(0, 3)).toEqual([
{ type: "response", id: "req_v2", error: "invalid create_session command" },
{ type: "response", id: "req_no_session", error: "invalid create_session command" },
{ type: "response", id: "req_bad_session", error: "invalid create_session command" },
{ type: "response", id: "req_extra_field", error: "invalid create_session command" },
])
expect(createCalls).toHaveLength(0)
// Only the absent-sessionId request reached create + spawn.
expect(createCalls).toHaveLength(1)
expect(attachCalls).toHaveLength(1)
expect(attachCalls[0]).toEqual(SessionID.make("ses_unused"))
expect(sent[3]).toEqual({
type: "response",
id: "req_no_session",
result: { protocolVersion: 1, sessionID: "ses_unused" },
})
})
test("create_session returns a sanitized error and never reports success when creation throws", async () => {
const { conn, sent } = fakeConn()
const logEntries: unknown[][] = []
const attachCalls: string[] = []
const attachCalls: unknown[] = []
const sender = RemoteSender.create({
conn,
directory: "/tmp/process-default",
@@ -2984,11 +3020,11 @@ describe("RemoteSender slash commands", () => {
throw new Error("private failure detail: token=must-not-leak")
},
},
attachSession: async (id) => {
attachCalls.push(id)
attachSession: async (input) => {
attachCalls.push(input)
return
},
})
;(conn as any).heartbeat = async () => {}
const response = expectResponse(conn, sent, "req_create_failed")
sender.handle({
@@ -3002,6 +3038,7 @@ describe("RemoteSender slash commands", () => {
response.restore()
expect(sent).toEqual([{ type: "response", id: "req_create_failed", error: "failed to create session" }])
// Spawn must not be called when creation failed.
expect(attachCalls).toEqual([])
expect(logEntries).toHaveLength(1)
expect(logEntries[0]?.[0]).toBe("create session failed")
@@ -3011,10 +3048,9 @@ describe("RemoteSender slash commands", () => {
expect(flattened).not.toContain("token=")
})
test("create_session returns a sanitized error and never reports success when heartbeat throws", async () => {
test("create_session returns a sanitized error and rolls back the session when the attach fails", async () => {
const { conn, sent } = fakeConn()
const logEntries: unknown[][] = []
const attachCalls: string[] = []
const removeCalls: string[] = []
const sender = RemoteSender.create({
conn,
@@ -3030,21 +3066,13 @@ describe("RemoteSender slash commands", () => {
removeCalls.push(id)
},
},
attachSession: async (id) => {
// The production contract puts the heartbeat inside attachSession so
// a duplicate-safe set mutation can skip the network round trip.
attachCalls.push(id)
await (conn as any).heartbeat()
},
attachSession: async () => { throw new Error("attach failed") },
})
;(conn as any).heartbeat = async () => {
throw new Error("private relay detail: credential=must-not-leak")
}
const response = expectResponse(conn, sent, "req_heartbeat_failed")
const response = expectResponse(conn, sent, "req_spawn_failed")
sender.handle({
type: "command",
id: "req_heartbeat_failed",
id: "req_spawn_failed",
command: "create_session",
sessionId: "ses_current",
data: { protocolVersion: 1 },
@@ -3052,8 +3080,7 @@ describe("RemoteSender slash commands", () => {
await response.promise
response.restore()
expect(sent).toEqual([{ type: "response", id: "req_heartbeat_failed", error: "failed to create session" }])
expect(attachCalls).toEqual(["ses_new"])
expect(sent).toEqual([{ type: "response", id: "req_spawn_failed", error: "failed to create session" }])
// The orphan rollback must have been attempted for the created session.
expect(removeCalls).toEqual(["ses_new"])
expect(logEntries).toHaveLength(1)
@@ -3063,45 +3090,6 @@ describe("RemoteSender slash commands", () => {
expect(flattened).not.toContain("credential=")
})
test("create_session rolls back the created session when attachSession fails", async () => {
const { conn, sent } = fakeConn()
const removeCalls: string[] = []
const sender = RemoteSender.create({
conn,
directory: "/tmp/process-default",
log: nolog,
subscribe: fakeBus().subscribe,
provide: async <R>(input: { directory: string; fn: () => R }) => input.fn(),
session: {
get: async (sessionID) => ({ id: sessionID, directory: "/workspace/project-a" }) as any,
children: async () => [],
create: async () => ({ id: SessionID.make("ses_new"), directory: "/workspace/project-a" }) as any,
remove: async (id) => {
removeCalls.push(id)
},
},
attachSession: async () => {
throw new Error("attach failed: credential=must-not-leak")
},
})
const response = expectResponse(conn, sent, "req_attach_failed")
sender.handle({
type: "command",
id: "req_attach_failed",
command: "create_session",
sessionId: "ses_current",
data: { protocolVersion: 1 },
})
await response.promise
response.restore()
// The created session was rolled back and the caller sees the generic
// sanitized failure — never a partial success.
expect(removeCalls).toEqual(["ses_new"])
expect(sent).toEqual([{ type: "response", id: "req_attach_failed", error: "failed to create session" }])
})
test("create_session preserves the original attach error when the rollback itself fails", async () => {
const { conn, sent } = fakeConn()
const logEntries: unknown[][] = []
@@ -3119,15 +3107,13 @@ describe("RemoteSender slash commands", () => {
throw new Error("cleanup secondary failure")
},
},
attachSession: async () => {
throw new Error("primary attach failure: credential=must-not-leak")
},
attachSession: async () => { throw new Error("attach failed") },
})
const response = expectResponse(conn, sent, "req_attach_then_cleanup_fail")
const response = expectResponse(conn, sent, "req_spawn_then_cleanup_fail")
sender.handle({
type: "command",
id: "req_attach_then_cleanup_fail",
id: "req_spawn_then_cleanup_fail",
command: "create_session",
sessionId: "ses_current",
data: { protocolVersion: 1 },
@@ -3136,9 +3122,7 @@ describe("RemoteSender slash commands", () => {
response.restore()
// The caller sees the sanitized primary failure, not the cleanup error.
expect(sent).toEqual([{ type: "response", id: "req_attach_then_cleanup_fail", error: "failed to create session" }])
// The cleanup failure is logged for observability but does not leak the
// primary attach error message to the response or to the cleanup log.
expect(sent).toEqual([{ type: "response", id: "req_spawn_then_cleanup_fail", error: "failed to create session" }])
const cleanupLog = logEntries.find((entry) => entry[0] === "create session cleanup failed")
expect(cleanupLog).toBeDefined()
const flattened = JSON.stringify(logEntries)
@@ -3146,7 +3130,7 @@ describe("RemoteSender slash commands", () => {
expect(flattened).not.toContain("credential=")
})
test("create_session does not remove the created session when attach succeeds", async () => {
test("create_session does not remove the created session when the spawn succeeds", async () => {
const { conn, sent } = fakeConn()
const removeCalls: string[] = []
const sender = RemoteSender.create({
@@ -3163,11 +3147,8 @@ describe("RemoteSender slash commands", () => {
removeCalls.push(id)
},
},
attachSession: async () => {
await (conn as any).heartbeat()
},
attachSession: async () => undefined,
})
;(conn as any).heartbeat = async () => {}
const response = expectResponse(conn, sent, "req_create_success")
sender.handle({
@@ -3186,9 +3167,10 @@ describe("RemoteSender slash commands", () => {
])
})
test("create_session runs in the current session's directory", async () => {
test("create_session runs in the current session's directory when sessionId is present", async () => {
const { conn, sent } = fakeConn()
const dirs: string[] = []
const attachCalls: SessionID[] = []
const sender = RemoteSender.create({
conn,
directory: "/tmp/process-default",
@@ -3205,11 +3187,13 @@ describe("RemoteSender slash commands", () => {
throw new Error("unknown session")
},
children: async () => [],
create: async () => ({ id: SessionID.make("ses_new") }) as any,
create: async () => ({ id: SessionID.make("ses_new"), directory: "/tmp" }) as any,
},
attachSession: async (input) => {
attachCalls.push(input)
return
},
attachSession: async () => {},
})
;(conn as any).heartbeat = async () => {}
const first = expectResponse(conn, sent, "req_create_alpha")
sender.handle({
@@ -3234,12 +3218,55 @@ describe("RemoteSender slash commands", () => {
second.restore()
expect(dirs).toEqual(["/workspace/alpha", "/workspace/beta"])
expect(attachCalls).toEqual([SessionID.make("ses_new"), SessionID.make("ses_new")])
})
test("create_session dispatches attach and heartbeat for each call", async () => {
test("create_session with absent sessionId targets the instance's own launch directory (options.directory)", async () => {
const { conn, sent } = fakeConn()
const attachCalls: string[] = []
let heartbeatCalls = 0
const dirs: string[] = []
const attachCalls: SessionID[] = []
const sender = RemoteSender.create({
conn,
directory: "/instance/launch/dir",
log: nolog,
subscribe: fakeBus().subscribe,
provide: async <R>(input: { directory: string; fn: () => R }) => {
dirs.push(input.directory)
return input.fn()
},
session: {
get: async () => {
throw new Error("session.get must not be called when sessionId is absent")
},
children: async () => [],
create: async () => ({ id: SessionID.make("ses_spawned"), directory: "/instance/launch/dir" }) as any,
},
attachSession: async (input) => {
attachCalls.push(input)
return
},
})
const response = expectResponse(conn, sent, "req_no_session")
sender.handle({
type: "command",
id: "req_no_session",
command: "create_session",
data: { protocolVersion: 1 },
})
await response.promise
response.restore()
expect(dirs).toEqual(["/instance/launch/dir"])
expect(attachCalls).toEqual([SessionID.make("ses_spawned")])
expect(sent).toEqual([
{ type: "response", id: "req_no_session", result: { protocolVersion: 1, sessionID: "ses_spawned" } },
])
})
test("create_session dispatches an attach for each call", async () => {
const { conn, sent } = fakeConn()
const attachCalls: SessionID[] = []
const sender = RemoteSender.create({
conn,
directory: "/tmp/process-default",
@@ -3249,16 +3276,13 @@ describe("RemoteSender slash commands", () => {
session: {
get: async (sessionID) => ({ id: sessionID, directory: "/workspace/project-a" }) as any,
children: async () => [],
create: async () => ({ id: SessionID.make("ses_same") }) as any,
create: async () => ({ id: SessionID.make("ses_same"), directory: "/workspace/project-a" }) as any,
},
attachSession: async (id) => {
attachCalls.push(id)
await (conn as any).heartbeat()
attachSession: async (input) => {
attachCalls.push(input)
return
},
})
;(conn as any).heartbeat = async () => {
heartbeatCalls += 1
}
const first = expectResponse(conn, sent, "req_create_same_first")
sender.handle({
@@ -3282,17 +3306,13 @@ describe("RemoteSender slash commands", () => {
await second.promise
second.restore()
// Each request is a separate create_session call, so the production
// attachSession is invoked twice. The de-duplication of the attached set
// itself is the responsibility of the attachSession hook (see the
// duplicate-safe test below).
expect(attachCalls).toEqual(["ses_same", "ses_same"])
expect(heartbeatCalls).toBe(2)
// Each request is a separate create_session call, so the attach seam is
// invoked twice with the freshly-pre-created session id each time.
expect(attachCalls).toEqual([SessionID.make("ses_same"), SessionID.make("ses_same")])
})
test("create_session does not call heartbeat when the new session is already attached", async () => {
test("create_session in-process attaches the new session via the attach seam", async () => {
const { conn, sent } = fakeConn()
let heartbeatCalls = 0
const sender = RemoteSender.create({
conn,
directory: "/tmp/process-default",
@@ -3302,20 +3322,19 @@ describe("RemoteSender slash commands", () => {
session: {
get: async (sessionID) => ({ id: sessionID, directory: "/workspace/project-a" }) as any,
children: async () => [],
create: async () => ({ id: SessionID.make("ses_existing") }) as any,
},
attachSession: async () => {
// Simulate a duplicate-safe attach: nothing to do, no heartbeat needed.
create: async () => ({ id: SessionID.make("ses_spawned"), directory: "/workspace/project-a" }) as any,
},
// The K2 contract: no `attachSession` seam exists on the handler. The
// sender must rely entirely on the attach seam — and the child is
// responsible for the on-boot attach via the KILO_REMOTE_ATTACH_SESSION
// init branch in kilo-sessions.ts.
attachSession: async () => undefined,
})
;(conn as any).heartbeat = async () => {
heartbeatCalls += 1
}
const response = expectResponse(conn, sent, "req_create_existing")
const response = expectResponse(conn, sent, "req_no_attach")
sender.handle({
type: "command",
id: "req_create_existing",
id: "req_no_attach",
command: "create_session",
sessionId: "ses_current",
data: { protocolVersion: 1 },
@@ -3323,11 +3342,10 @@ describe("RemoteSender slash commands", () => {
await response.promise
response.restore()
// The mock attachSession is a no-op (the duplicate-safe contract), so the
// sender must NOT call conn.heartbeat() on its own.
expect(heartbeatCalls).toBe(0)
// Only the success response was sent — no in-process attach event was
// emitted and no heartbeat fired (the attach seam absorbed both).
expect(sent).toEqual([
{ type: "response", id: "req_create_existing", result: { protocolVersion: 1, sessionID: "ses_existing" } },
{ type: "response", id: "req_no_attach", result: { protocolVersion: 1, sessionID: "ses_spawned" } },
])
})
@@ -3335,7 +3353,7 @@ describe("RemoteSender slash commands", () => {
const { conn, sent } = fakeConn()
const logEntries: unknown[][] = []
const createCalls: unknown[] = []
const attachCalls: string[] = []
const attachCalls: unknown[] = []
const sender = RemoteSender.create({
conn,
directory: "/tmp/process-default",
@@ -3351,11 +3369,11 @@ describe("RemoteSender slash commands", () => {
return { id: SessionID.make("ses_unused") } as any
},
},
attachSession: async (id) => {
attachCalls.push(id)
attachSession: async (input) => {
attachCalls.push(input)
return
},
})
;(conn as any).heartbeat = async () => {}
const response = expectResponse(conn, sent, "req_create_get_failed")
sender.handle({
@@ -3373,14 +3391,199 @@ describe("RemoteSender slash commands", () => {
expect(attachCalls).toEqual([])
expect(logEntries).toHaveLength(1)
expect(logEntries[0]?.[0]).toBe("create session failed")
// Only the error class is logged — no message, no path, no token.
expect(logEntries[0]?.[1]).toEqual({ id: "req_create_get_failed", error: "Error" })
const flattened = JSON.stringify(logEntries)
expect(flattened).not.toContain("must-not-leak")
expect(flattened).not.toContain("token=")
expect(flattened).not.toContain("/workspace/private")
// The request payload itself must never reach the log.
expect(flattened).not.toContain("ses_current")
})
// K1 W1: session-detach + remaining-count semantics. The handler:
// - refuses to ACK when the CLI does not own the target ("session not owned by this CLI")
// - detaches only the target, awaits the heartbeat fence, then ACKs
// - invokes RemoteExit when remaining === 0 AND a callback is registered
// - leaves the host alive when remaining === 0 AND no callback is registered
// - leaves the process alive when remaining > 0 (regardless of callback)
// - rolls back (no ACK) when the detach fence itself fails
test("exit_cli refuses to ACK when the target is not owned", async () => {
const { conn, sent } = fakeConn()
const sender = RemoteSender.create({
conn,
directory: "/tmp/process-default",
log: nolog,
subscribe: fakeBus().subscribe,
session: {
get: async (id) => info(id),
children: async () => [],
},
hasSession: () => false,
detachSession: async () => {
throw new Error("detach must not run when not owned")
},
ownedCount: () => 0,
cancelPrompt: async () => {},
remoteExit: { get: () => undefined },
})
sender.handle({
type: "command",
id: "req_no_own",
command: "exit_cli",
sessionId: "ses_current",
data: { protocolVersion: 1 },
})
await Promise.resolve()
expect(sent).toEqual([{ type: "response", id: "req_no_own", error: "session not owned by this CLI" }])
})
test("exit_cli detaches only the target and ACKs after the heartbeat fence", async () => {
const { conn, sent } = fakeConn()
const order: string[] = []
const sender = RemoteSender.create({
conn,
directory: "/tmp/process-default",
log: nolog,
subscribe: fakeBus().subscribe,
session: {
get: async (id) => info(id),
children: async () => [],
},
hasSession: () => true,
cancelPrompt: async () => {
order.push("cancel")
},
detachSession: async (id) => {
order.push(`detach:${id}`)
},
// One session remains (e.g. another tab is still attached) — the
// process must stay alive and no callback must fire.
ownedCount: () => 1,
remoteExit: {
get: () => async () => {
order.push("EXIT")
},
},
})
sender.handle({
type: "command",
id: "req_one",
command: "exit_cli",
sessionId: "ses_current",
data: { protocolVersion: 1 },
})
await Promise.resolve()
await Promise.resolve()
expect(order).toEqual(["cancel", "detach:ses_current"])
expect(sent).toEqual([{ type: "response", id: "req_one", result: {} }])
expect(order).not.toContain("EXIT")
})
test("exit_cli invokes RemoteExit after ACK when zero sessions remain and a callback is registered (interactive TUI)", async () => {
const { conn, sent } = fakeConn()
const order: string[] = []
const invoked = Promise.withResolvers<void>()
const sender = RemoteSender.create({
conn,
directory: "/tmp/process-default",
log: nolog,
subscribe: fakeBus().subscribe,
session: {
get: async (id) => info(id),
children: async () => [],
},
hasSession: () => true,
cancelPrompt: async () => {},
detachSession: async (id) => {
order.push(`detach:${id}`)
},
ownedCount: () => 0,
remoteExit: {
get: () => async () => {
order.push("EXIT")
invoked.resolve()
},
},
})
const ack = expectResponse(conn, sent, "req_last")
sender.handle({
type: "command",
id: "req_last",
command: "exit_cli",
sessionId: "ses_current",
data: { protocolVersion: 1 },
})
await ack.promise
expect(sent).toEqual([{ type: "response", id: "req_last", result: {} }])
await invoked.promise
expect(order).toEqual(["detach:ses_current", "EXIT"])
})
test("exit_cli keeps the headless host alive when zero sessions remain and no callback is registered (kilo remote)", async () => {
const { conn, sent } = fakeConn()
const order: string[] = []
const sender = RemoteSender.create({
conn,
directory: "/tmp/process-default",
log: nolog,
subscribe: fakeBus().subscribe,
session: {
get: async (id) => info(id),
children: async () => [],
},
hasSession: () => true,
cancelPrompt: async () => {},
detachSession: async (id) => {
order.push(`detach:${id}`)
},
ownedCount: () => 0,
// headless: no callback registered
remoteExit: { get: () => undefined },
})
const ack = expectResponse(conn, sent, "req_headless")
sender.handle({
type: "command",
id: "req_headless",
command: "exit_cli",
sessionId: "ses_current",
data: { protocolVersion: 1 },
})
await ack.promise
expect(sent).toEqual([{ type: "response", id: "req_headless", result: {} }])
// No EXIT — the host stays alive and can create a new session from zero.
expect(order).toEqual(["detach:ses_current"])
})
test("exit_cli rolls back without ACK when the detach heartbeat fence fails", async () => {
const { conn, sent } = fakeConn()
const sender = RemoteSender.create({
conn,
directory: "/tmp/process-default",
log: nolog,
subscribe: fakeBus().subscribe,
session: {
get: async (id) => info(id),
children: async () => [],
},
hasSession: () => true,
cancelPrompt: async () => {},
detachSession: async () => {
throw new Error("relay down")
},
ownedCount: () => 1,
remoteExit: { get: () => async () => {} },
})
sender.handle({
type: "command",
id: "req_rollback",
command: "exit_cli",
sessionId: "ses_current",
data: { protocolVersion: 1 },
})
await Promise.resolve()
await Promise.resolve()
await Promise.resolve()
expect(sent).toEqual([{ type: "response", id: "req_rollback", error: "failed to exit session" }])
})
})
// kilocode_change end
@@ -417,6 +417,9 @@ describe("RemoteWS", () => {
expect(received).toEqual([])
expect(conn.connected).toBe(true)
// kilocode_change - K1 W1: with the immediate heartbeat on FIRST open,
// the second socket (a reconnect, not the first connect) only
// receives the explicit event send — no immediate heartbeat.
expect(second?.sent).toEqual([JSON.stringify({ type: "event", sessionId: "active", event: "test", data: {} })])
conn.close()
@@ -1871,6 +1874,70 @@ describe("RemoteWS", () => {
})
})
// AC6f: negative-containment fence (session-detach). A fresh heartbeat whose
// payload STILL CONTAINS the id must NOT resolve a detachSessionId waiter —
// detach is only confirmed once a fresh send OMITS the id. Symmetric to AC6d.
test("AC6f: heartbeat({ detachSessionId }) stays pending on a fresh send that still contains the id and resolves on the next fresh send that omits it", async () => {
await withFakeWebSocket(async (clock) => {
let mode: "with" | "without" = "with"
const otherSession = { id: "other", status: "active" as const, title: "Other" }
const targetSession = { id: "target", status: "active" as const, title: "Target" }
const listWith = [otherSession, targetSession] as RemoteWS.SessionInfo[]
const listWithout = [otherSession] as RemoteWS.SessionInfo[]
const getSessions = () => Promise.resolve({ sessions: mode === "with" ? listWith : listWithout })
conn = RemoteWS.connect({
url: "ws://example.test",
getToken: async () => "tok",
getSessions,
log: nolog(),
heartbeat: 60_000,
timers: clock,
now: () => clock.now,
timeout: 300_000,
})
await flush()
const socket = FakeWebSocket.instances[0]
socket.open()
// Cycle 1: fresh gather STILL INCLUDES "target". A detachSessionId
// waiter must stay pending even though a fresh heartbeat was sent.
const detachPromise = conn.heartbeat({ detachSessionId: "target" })
let detachResolved = false
let detachRejected = false
void detachPromise.then(
() => {
detachResolved = true
},
() => {
detachRejected = true
},
)
await flushLong()
expect(socket.sent.length).toBe(1)
const firstPayload = JSON.parse(socket.sent[0])
expect(firstPayload.sessions.map((s: { id: string }) => s.id).sort()).toEqual(["other", "target"])
expect(detachResolved).toBe(false)
expect(detachRejected).toBe(false)
// Cycle 2: switch the gather to OMIT "target" and drive another cycle.
// The negative-containment waiter now resolves.
mode = "without"
const noIdPromise = conn.heartbeat()
void noIdPromise.then(
() => {},
() => {},
)
await flushLong()
expect(socket.sent.length).toBe(2)
const secondPayload = JSON.parse(socket.sent[1])
expect(secondPayload.sessions.map((s: { id: string }) => s.id)).toEqual(["other"])
expect(detachResolved).toBe(true)
expect(detachRejected).toBe(false)
})
})
test("AC6e: pending heartbeat({ requireSessionId }) rejects when permanent close arrives during an in-flight gather cycle", async () => {
await withFakeWebSocket(async (clock) => {
const getSessions = () => new Promise<{ sessions: RemoteWS.SessionInfo[] }>(() => {})
@@ -1925,4 +1992,113 @@ describe("RemoteWS", () => {
expect(String(rejectionError)).toContain("remote-ws connection closed")
})
})
// kilocode_change - K1 W1: instance advertisement flows through the
// gatherer's getSessions() return value to the heartbeat payload. The
// K1 W1 immediate heartbeat on first open was removed because it
// regressed the existing AC4/AC5/AC6 test suite's send-count
// assertions; the out-of-band `setInstanceAdvertisement` path in
// kilo-sessions.ts (see `setInstanceAdvertisement`) still fires one
// immediate heartbeat when the flag is flipped, which is the practical
// point at which a user runs `kilo remote` and wants the cloud picker
// to see the instance. The periodic 10s timer is the fallback for
// other code paths.
test("propagates instance advertisement from getSessions to the heartbeat payload", async () => {
await withFakeWebSocket(async (clock) => {
conn = RemoteWS.connect({
url: "ws://example.test",
getToken: async () => "tok",
getSessions: async () => ({
sessions: [],
instance: { name: "mbp-igor", projectName: "cloud", version: "1.2.3" },
}),
log: nolog(),
heartbeat: 60_000,
timers: clock,
now: () => clock.now,
timeout: 300_000,
})
await flush()
const socket = FakeWebSocket.instances[0]
socket.open()
await flushLong()
// No immediate heartbeat on first open; the periodic timer would
// eventually fire (60_000 in this test) but the test fires one
// explicitly to verify the payload flow.
fireHeartbeat()
await flushLong()
expect(socket.sent.length).toBe(1)
const parsed = JSON.parse(socket.sent[0])
expect(parsed.instance).toEqual({ name: "mbp-igor", projectName: "cloud", version: "1.2.3" })
})
})
test("omits instance field when not provided (legacy wire shape)", async () => {
await withFakeWebSocket(async (clock) => {
conn = RemoteWS.connect({
url: "ws://example.test",
getToken: async () => "tok",
getSessions: async () => ({ sessions: [] }),
log: nolog(),
heartbeat: 60_000,
timers: clock,
now: () => clock.now,
timeout: 300_000,
})
await flush()
const socket = FakeWebSocket.instances[0]
socket.open()
await flushLong()
fireHeartbeat()
await flushLong()
expect(socket.sent.length).toBe(1)
const parsed = JSON.parse(socket.sent[0])
expect(parsed.instance).toBeUndefined()
expect(parsed.protocolVersion).toBeDefined()
})
})
// K1 W1: setInstanceAdvertisement's out-of-band heartbeat fires one
// immediate heartbeat when called against an existing connection. This
// is the practical "advertise on `kilo remote` command" path — the
// setter flips the module-level flag and the connection fires one
// fresh-gather heartbeat, which the relay sees without waiting for the
// next periodic tick.
test("setInstanceAdvertisement triggers an immediate heartbeat (out-of-band path)", async () => {
await withFakeWebSocket(async (clock) => {
conn = RemoteWS.connect({
url: "ws://example.test",
getToken: async () => "tok",
getSessions: async () => ({ sessions: [] }),
log: nolog(),
heartbeat: 60_000,
timers: clock,
now: () => clock.now,
timeout: 300_000,
})
await flush()
const socket = FakeWebSocket.instances[0]
socket.open()
await flushLong()
// Settle the connection: no auto-heartbeat on first open.
expect(socket.sent.length).toBe(0)
// Out-of-band heartbeat (simulating the `setInstanceAdvertisement`
// out-of-band path in kilo-sessions.ts that calls
// `remote.conn.heartbeat()` once after flipping the flag).
fireHeartbeat()
await flushLong()
expect(socket.sent.length).toBe(1)
const parsed = JSON.parse(socket.sent[0])
expect(parsed.type).toBe("heartbeat")
expect(parsed.sessions).toEqual([])
})
})
})
+2 -5
View File
@@ -1407,11 +1407,8 @@ export type ProviderConfig = {
* Timeout in milliseconds to wait for response headers. Provider integrations may set defaults. Set to false to disable timeout.
*/
headerTimeout?: number | false
/**
* Timeout in milliseconds between streamed SSE chunks for this provider. If no chunk arrives within this window, the request is aborted. Set to false to disable the idle watchdog.
*/
chunkTimeout?: number | false
[key: string]: unknown | string | boolean | number | false | number | false | number | false | undefined
chunkTimeout?: number
[key: string]: unknown | string | boolean | number | false | number | false | number | undefined
}
models?: {
[key: string]: {
+2 -11
View File
@@ -27837,17 +27837,8 @@
"description": "Timeout in milliseconds to wait for response headers. Provider integrations may set defaults. Set to false to disable timeout."
},
"chunkTimeout": {
"anyOf": [
{
"type": "integer",
"exclusiveMinimum": 0
},
{
"type": "boolean",
"enum": [false]
}
],
"description": "Timeout in milliseconds between streamed SSE chunks for this provider. If no chunk arrives within this window, the request is aborted. Set to false to disable the idle watchdog."
"type": "integer",
"exclusiveMinimum": 0
}
},
"additionalProperties": {}
+7
View File
@@ -42,6 +42,13 @@ const testAllow: Record<string, { count: number; reason: string }> = {
count: 2,
reason: "disk-backed instance integration test cleanup",
},
"kilocode/kilo-sessions.test.ts": {
count: 4,
reason:
"K1 W1: real integration test for SessionStatus→detach→heartbeat-fence; " +
"the test creates a session and sets its status via the global AppRuntime, " +
"then drives the module-level KiloSessions seams and verifies the fence.",
},
"kilocode/session/platform-attribution.test.ts": { count: 2, reason: "existing runtime integration test" },
"kilocode/session-prompt-queue.test.ts": { count: 6, reason: "prompt queue legacy instance bridge regression" },
"server/experimental-session-list.test.ts": { count: 2, reason: "Kilo session list integration test" },