opensource/kilocode
3
0
Fork 0
mirror of https://github.com/Kilo-Org/kilocode.git synced 2026-08-29 03:44:06 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix(jetbrains): hide file mention payloads

Browse Source
This commit is contained in:
kirillk
2026-06-16 18:33:18 -04:00
parent 54ff8c0df7
commit c23c3e300d
14 changed files with 598 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.changeset/jetbrains-file-mention-parity.md
+5
View File
@@ -0,0 +1,5 @@
---
"@kilocode/kilo-jetbrains": patch
---
Hide raw file contents from mentioned files in JetBrains chat messages.
packages/kilo-jetbrains/backend/src/main/kotlin/ai/kilocode/backend/cli/KiloCliDataParser.kt
+41
View File
@@ -18,6 +18,8 @@ import ai.kilocode.rpc.dto.MessageWithPartsDto
import ai.kilocode.rpc.dto.ModelSelectionDto
import ai.kilocode.rpc.dto.ModelStateDto
import ai.kilocode.rpc.dto.PartDto
import ai.kilocode.rpc.dto.PartSourceDto
import ai.kilocode.rpc.dto.PartSourceTextDto
import ai.kilocode.rpc.dto.PermissionAlwaysRulesDto
import ai.kilocode.rpc.dto.PermissionFileDiffDto
import ai.kilocode.rpc.dto.PermissionReplyDto
@@ -519,6 +521,7 @@ object KiloCliDataParser {
part.mime?.let { fields += "\"mime\":${escape(it)}" }
part.url?.let { fields += "\"url\":${escape(it)}" }
part.filename?.let { fields += "\"filename\":${escape(it)}" }
part.source?.let { fields += "\"source\":${sourceJson(it)}" }
return "{${fields.joinToString(",")}}"
}
fields += "\"text\":${escape(part.text.orEmpty())}"
@@ -649,6 +652,8 @@ object KiloCliDataParser {
mime = obj.str("mime"),
url = obj.str("url"),
filename = obj.str("filename"),
synthetic = obj.flagOrNull("synthetic"),
source = parseSource(obj["source"]),
tool = obj.str("tool"),
callID = obj.str("callID"),
state = state?.str("status"),
@@ -676,6 +681,37 @@ object KiloCliDataParser {
return READ_TOOL_PATH.containsMatchIn(line)
}
private fun parseSource(raw: JsonElement?): PartSourceDto? {
val obj = raw.obj() ?: return null
val type = obj.str("type") ?: return null
val text = obj["text"].obj() ?: return null
val value = text.str("value") ?: return null
val start = text.num("start") ?: return null
val end = text.num("end") ?: return null
return PartSourceDto(
type = type,
text = PartSourceTextDto(value = value, start = start, end = end),
path = obj.str("path"),
clientName = obj.str("clientName"),
uri = obj.str("uri"),
name = obj.str("name"),
kind = obj.long("kind")?.safeInt(),
)
}
private fun sourceJson(source: PartSourceDto): String {
val fields = mutableListOf(
"\"type\":${escape(source.type)}",
"\"text\":{\"value\":${escape(source.text.value)},\"start\":${source.text.start},\"end\":${source.text.end}}",
)
source.path?.let { fields += "\"path\":${escape(it)}" }
source.clientName?.let { fields += "\"clientName\":${escape(it)}" }
source.uri?.let { fields += "\"uri\":${escape(it)}" }
source.name?.let { fields += "\"name\":${escape(it)}" }
source.kind?.let { fields += "\"kind\":$it" }
return "{${fields.joinToString(",")}}"
}
internal fun parseTodos(raw: JsonElement?): List<TodoDto> {
return parseTodosOrNull(raw) ?: emptyList()
}
@@ -1109,6 +1145,11 @@ private fun JsonObject.flag(key: String, default: Boolean): Boolean {
return prim.booleanOrNull ?: prim.contentOrNull?.toBooleanStrictOrNull() ?: default
}
private fun JsonObject.flagOrNull(key: String): Boolean? {
val prim = this[key]?.jsonPrimitive ?: return null
return prim.booleanOrNull ?: prim.contentOrNull?.toBooleanStrictOrNull()
}
private fun Long.safeInt() = coerceIn(Int.MIN_VALUE.toLong(), Int.MAX_VALUE.toLong()).toInt()
private fun JsonObject?.map(key: String): Map<String, String> {
packages/kilo-jetbrains/backend/src/test/kotlin/ai/kilocode/backend/cli/KiloCliDataParserTest.kt
+97
View File
@@ -10,6 +10,8 @@ import ai.kilocode.rpc.dto.PermissionAlwaysRulesDto
import ai.kilocode.rpc.dto.PermissionReplyDto
import ai.kilocode.rpc.dto.ModelSelectionDto
import ai.kilocode.rpc.dto.ModelStateDto
import ai.kilocode.rpc.dto.PartSourceDto
import ai.kilocode.rpc.dto.PartSourceTextDto
import ai.kilocode.rpc.dto.PromptDto
import ai.kilocode.rpc.dto.PromptPartDto
import ai.kilocode.rpc.dto.QuestionReplyDto
@@ -190,6 +192,42 @@ class KiloCliDataParserTest {
assertEquals("a.png", result.part.filename)
}
@Test
fun `parseChatEvent - part preserves synthetic flag and source metadata`() {
val data = globalEvent("""
"type": "message.part.updated",
"properties": {
"sessionID": "ses_1",
"part": {
"id": "file_1",
"sessionID": "ses_1",
"messageID": "msg_1",
"type": "file",
"mime": "text/plain",
"url": "file:///tmp/a.kt",
"filename": "a.kt",
"synthetic": true,
"source": {
"type": "file",
"path": "src/a.kt",
"text": { "value": "@src/a.kt", "start": 4, "end": 13 }
}
}
}
""")
val result = KiloCliDataParser.parseChatEvent("message.part.updated", data)
assertNotNull(result)
assertTrue(result is ChatEventDto.PartUpdated)
assertEquals(true, result.part.synthetic)
assertEquals("file", result.part.source?.type)
assertEquals("src/a.kt", result.part.source?.path)
assertEquals("@src/a.kt", result.part.source?.text?.value)
assertEquals(4.0, result.part.source?.text?.start)
assertEquals(13.0, result.part.source?.text?.end)
}
@Test
fun `ChatEventNormalizer - user part updated sanitizes text`() {
val norm = KiloCliDataParser.ChatEventNormalizer()
@@ -1112,6 +1150,25 @@ class KiloCliDataParserTest {
)
}
@Test
fun `parseMessages - preserves synthetic and source metadata`() {
val raw = """[
{
"info": { "id": "m1", "sessionID": "s1", "role": "user", "time": { "created": 1.0 } },
"parts": [
{ "id": "p1", "sessionID": "s1", "messageID": "m1", "type": "text", "text": "hidden", "synthetic": true },
{ "id": "f1", "sessionID": "s1", "messageID": "m1", "type": "file", "mime": "text/plain", "url": "file:///tmp/a.kt", "source": { "type": "file", "path": "src/a.kt", "text": { "value": "@src/a.kt", "start": 0, "end": 9 } } }
]
}
]"""
val result = KiloCliDataParser.parseMessages(raw).single()
assertEquals(true, result.parts[0].synthetic)
assertEquals("src/a.kt", result.parts[1].source?.path)
assertEquals("@src/a.kt", result.parts[1].source?.text?.value)
}
@Test
fun `parseMessages - message with tool parts`() {
val raw = """[{
@@ -1542,6 +1599,46 @@ class KiloCliDataParserTest {
assertTrue(result.contains(""""filename":"a \"b\".txt""""), result)
}
@Test
fun `buildPromptJson - file part includes source metadata`() {
val prompt = PromptDto(parts = listOf(PromptPartDto(
type = "file",
mime = "text/plain",
url = "file:///tmp/a.kt",
filename = "a.kt",
source = PartSourceDto(
type = "file",
path = "src/a.kt",
text = PartSourceTextDto("@src/a.kt", 4.0, 13.0),
),
)))
val result = KiloCliDataParser.buildPromptJson(prompt)
assertEquals(
"""{"parts":[{"type":"file","mime":"text/plain","url":"file:///tmp/a.kt","filename":"a.kt","source":{"type":"file","text":{"value":"@src/a.kt","start":4.0,"end":13.0},"path":"src/a.kt"}}]}""",
result,
)
}
@Test
fun `buildCommandJson - file part includes source metadata`() {
val prompt = PromptDto(parts = listOf(PromptPartDto(
type = "file",
mime = "text/plain",
url = "file:///tmp/a.kt",
source = PartSourceDto(
type = "file",
path = "src/a.kt",
text = PartSourceTextDto("@src/a.kt", 0.0, 9.0),
),
)))
val result = KiloCliDataParser.buildCommandJson("review", "", prompt)
assertTrue(result.contains(""""source":{"type":"file","text":{"value":"@src/a.kt","start":0.0,"end":9.0},"path":"src/a.kt"}"""), result)
}
// ---- buildSummarizeJson ----
@Test
packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/session/SessionUi.kt
+28 -7
View File
@@ -47,6 +47,8 @@ import ai.kilocode.client.vfs.KiloVfsManager
import ai.kilocode.log.ChatLogSummary
import ai.kilocode.rpc.dto.PromptDto
import ai.kilocode.rpc.dto.PromptPartDto
import ai.kilocode.rpc.dto.PartSourceDto
import ai.kilocode.rpc.dto.PartSourceTextDto
import com.intellij.util.ui.JBUI
import ai.kilocode.log.KiloLog
import com.intellij.ide.BrowserUtil
@@ -415,6 +417,7 @@ class SessionUi(
prompt.reasoning.setItems(m.variants.map { ReasoningPicker.Item(it, variantTitle(it)) }, m.variant)
prompt.setResetVisible(m.modelOverride)
prompt.setReady(m.isReady())
prompt.refreshHighlights()
}
is SessionControllerEvent.ViewChanged.ShowProgress -> {
@@ -655,29 +658,47 @@ class SessionUi(
private fun mentionParts(text: String, paths: Set<String>): List<PromptPartDto> = buildList {
paths.filter { text.contains("@$it") }.forEach { path ->
val token = "@$path"
val start = text.indexOf(token).takeIf { it >= 0 } ?: return@forEach
val target = runCatching {
val item = Path.of(path)
if (item.isAbsolute) item else Path.of(workspace.directory).resolve(item).normalize()
}.getOrNull() ?: return@forEach
add(PromptPartDto(type = "file", mime = "text/plain", url = target.toUri().toString(), filename = target.fileName?.toString()))
add(PromptPartDto(
type = "file",
mime = "text/plain",
url = target.toUri().toString(),
filename = target.fileName?.toString(),
source = source("file", token, start, path = path),
))
}
if (text.contains("@terminal")) {
val terminal = text.indexOf("@terminal")
if (terminal >= 0) {
runBlocking { workspaces.terminalOutput(workspace.directory) }
?.takeIf { it.isNotBlank() }
?.let { add(dataPart("terminal-output.txt", it)) }
?.let { add(dataPart("terminal-output.txt", it, source("resource", "@terminal", terminal, uri = "terminal"))) }
}
if (text.contains("@git-changes")) {
val git = text.indexOf("@git-changes")
if (git >= 0) {
runBlocking { workspaces.gitChanges(workspace.directory) }
?.takeIf { it.isNotBlank() }
?.let { add(dataPart("git-changes.txt", it)) }
?.let { add(dataPart("git-changes.txt", it, source("resource", "@git-changes", git, uri = "git-changes"))) }
}
}
private fun dataPart(name: String, text: String): PromptPartDto {
private fun dataPart(name: String, text: String, source: PartSourceDto? = null): PromptPartDto {
val data = URLEncoder.encode(text, StandardCharsets.UTF_8).replace("+", "%20")
return PromptPartDto(type = "file", mime = "text/plain", url = "data:text/plain;charset=utf-8,$data", filename = name)
return PromptPartDto(type = "file", mime = "text/plain", url = "data:text/plain;charset=utf-8,$data", filename = name, source = source)
}
private fun source(type: String, token: String, start: Int, path: String? = null, uri: String? = null) = PartSourceDto(
type = type,
text = PartSourceTextDto(value = token, start = start.toDouble(), end = (start + token.length).toDouble()),
path = path,
uri = uri,
clientName = if (type == "resource") "jetbrains" else null,
)
private fun openFile(path: String) {
cs.launch {
workspaces.openPath(workspace.directory, path)
packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/session/model/Message.kt
+2
View File
@@ -1,6 +1,7 @@
package ai.kilocode.client.session.model
import ai.kilocode.rpc.dto.MessageDto
import ai.kilocode.rpc.dto.PartSourceDto
import ai.kilocode.rpc.dto.PartTimeDto
import ai.kilocode.rpc.dto.TodoDto
import ai.kilocode.rpc.dto.TodoViewDto
@@ -71,6 +72,7 @@ class FileAttachment(id: String) : Content(id) {
var mime: String = "application/octet-stream"
var url: String = ""
var filename: String? = null
var source: PartSourceDto? = null
}
/** Tool invocation with lifecycle state. */
packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/session/model/SessionModel.kt
+25
View File
@@ -41,6 +41,7 @@ class SessionModel {
private val entries = LinkedHashMap<String, Message>()
private val turnEntries = LinkedHashMap<String, Turn>()
private val hiddenText = mutableSetOf<Pair<String, String>>()
var app: KiloAppStateDto = KiloAppStateDto(KiloAppStatusDto.DISCONNECTED)
var version: String? = null
@@ -140,6 +141,7 @@ class SessionModel {
@RequiresEdt
fun removeMessage(id: String) {
if (entries.remove(id) == null) return
hiddenText.removeAll { it.first == id }
fire(SessionModelEvent.MessageRemoved(id))
regroup()
updateHeader()
@@ -147,6 +149,7 @@ class SessionModel {
@RequiresEdt
fun removeContent(messageId: String, contentId: String) {
hiddenText.remove(messageId to contentId)
val msg = entries[messageId] ?: return
if (msg.parts.remove(contentId) == null) return
fire(SessionModelEvent.ContentRemoved(messageId, contentId))
@@ -157,6 +160,16 @@ class SessionModel {
fun updateContent(messageId: String, dto: PartDto) {
if (dto.type in SILENT_PART_TYPES) return
val msg = entries[messageId] ?: return
val key = messageId to dto.id
if (hiddenSynthetic(msg, dto)) {
hiddenText.add(key)
if (msg.parts.remove(dto.id) != null) {
fire(SessionModelEvent.ContentRemoved(messageId, dto.id))
updateHeader()
}
return
}
hiddenText.remove(key)
val existing = msg.parts[dto.id]
if (empty(dto)) {
if (existing is Text) removeContent(messageId, dto.id)
@@ -175,6 +188,7 @@ class SessionModel {
@RequiresEdt
fun appendDelta(messageId: String, contentId: String, delta: String) {
val msg = entries[messageId] ?: return
if (hiddenText.contains(messageId to contentId)) return
val existing = msg.parts[contentId]
val created = existing == null
if (existing != null) {
@@ -238,6 +252,7 @@ class SessionModel {
@RequiresEdt
fun loadHistory(history: List<MessageWithPartsDto>) {
entries.clear()
hiddenText.clear()
session = null
state = SessionState.Idle
diff = emptyList()
@@ -247,6 +262,10 @@ class SessionModel {
val item = Message(msg.info)
for (part in msg.parts) {
if (part.type in SILENT_PART_TYPES) continue
if (hiddenSynthetic(item, part)) {
hiddenText.add(msg.info.id to part.id)
continue
}
if (empty(part)) continue
val content = fromDto(part, part.text)
item.parts[content.id] = content
@@ -262,6 +281,7 @@ class SessionModel {
fun clear() {
entries.clear()
turnEntries.clear()
hiddenText.clear()
session = null
state = SessionState.Idle
diff = emptyList()
@@ -382,6 +402,7 @@ class SessionModel {
existing.mime = dto.mime ?: "application/octet-stream"
existing.url = dto.url ?: ""
existing.filename = dto.filename
existing.source = dto.source
}
is Tool -> {
existing.kind = toolKind(dto.tool)
@@ -410,6 +431,9 @@ class SessionModel {
private fun empty(dto: PartDto) = dto.type == "text" && dto.text?.isNotBlank() != true
private fun hiddenSynthetic(msg: Message, dto: PartDto) =
msg.info.role == "user" && dto.type == "text" && dto.synthetic == true
private fun fromDto(dto: PartDto, text: CharSequence? = null): Content {
val content = text ?: dto.text
return when (dto.type) {
@@ -424,6 +448,7 @@ class SessionModel {
mime = dto.mime ?: "application/octet-stream"
url = dto.url ?: ""
filename = dto.filename
source = dto.source
}
"tool" -> Tool(dto.id, dto.tool ?: "unknown", toolKind(dto.tool)).apply {
state = parseToolState(dto.state)
packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/session/ui/prompt/KiloPromptCompletionProvider.kt
+34
View File
@@ -38,6 +38,10 @@ class KiloPromptCompletionProvider(
val action: () -> Unit,
)
data class Highlight(val start: Int, val end: Int, val kind: HighlightKind)
enum class HighlightKind { MENTION, COMMAND }
fun mentionPaths(): Set<String> = paths.toSet()
fun clearMentions() {
@@ -47,6 +51,36 @@ class KiloPromptCompletionProvider(
fun clientNames(): Set<String> = actions.mapTo(mutableSetOf()) { it.name }
fun highlights(text: String): List<Highlight> = buildList {
val command = text.takeIf { it.startsWith('/') }
?.drop(1)
?.takeWhile { !it.isWhitespace() }
?.takeIf { it.isNotBlank() }
val commands = workspace.state.value.commands.mapTo(mutableSetOf()) { it.name }
if (command != null && (command in clientNames() || command in commands)) {
add(Highlight(0, command.length + 1, HighlightKind.COMMAND))
}
val ranges = mutableListOf<IntRange>()
val values = (mentionPaths() + setOf("terminal", "git-changes"))
.filter { it.isNotBlank() }
.sortedByDescending { it.length }
values.forEach { value ->
val raw = "@$value"
var idx = text.indexOf(raw)
while (idx >= 0) {
val end = idx + raw.length
val valid = end == text.length || text[end].isWhitespace()
val range = idx until end
if (valid && ranges.none { it.first < end && idx < it.last + 1 }) {
ranges += range
add(Highlight(idx, end, HighlightKind.MENTION))
}
idx = text.indexOf(raw, idx + 1)
}
}
}
override fun getAdvertisement(): String? = null
override fun getPrefix(text: String, offset: Int): String? = token(text, offset)?.prefix
packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/session/ui/prompt/PromptPanel.kt
+43
View File
@@ -39,8 +39,13 @@ import com.intellij.openapi.actionSystem.UiDataProvider
import com.intellij.openapi.actionSystem.ex.ActionUtil
import com.intellij.openapi.actionSystem.IdeActions
import com.intellij.openapi.application.ApplicationManager
import com.intellij.openapi.editor.DefaultLanguageHighlighterColors
import com.intellij.openapi.editor.colors.TextAttributesKey
import com.intellij.openapi.editor.event.DocumentEvent
import com.intellij.openapi.editor.event.DocumentListener
import com.intellij.openapi.editor.markup.HighlighterLayer
import com.intellij.openapi.editor.markup.HighlighterTargetArea
import com.intellij.openapi.editor.markup.RangeHighlighter
import com.intellij.openapi.keymap.Keymap
import com.intellij.openapi.keymap.KeymapManagerListener
import com.intellij.openapi.keymap.KeymapUtil
@@ -98,6 +103,8 @@ class PromptPanel(
private val SHIELD_ICON: Icon = IconLoader.getIcon("/icons/shield.svg", PromptPanel::class.java)
private val SHIELD_FILLED_ICON: Icon = IconLoader.getIcon("/icons/shield-filled.svg", PromptPanel::class.java)
private val WAND_ICON: Icon = IconLoader.getIcon("/icons/wand-sparkles.svg", PromptPanel::class.java)
private val MENTION_KEY = DefaultLanguageHighlighterColors.METADATA
private val COMMAND_KEY = DefaultLanguageHighlighterColors.KEYWORD
}
val mode = ModePicker()
@@ -120,6 +127,7 @@ class PromptPanel(
)
}
private val attachments = mutableListOf<PromptAttachment>()
private val highlighters = mutableListOf<RangeHighlighter>()
private val strip = PromptAttachmentStrip(project) { removeAttachment(it) }
private var bus: MessageBusConnection? = null
private var autoApprove = false
@@ -156,6 +164,7 @@ class PromptPanel(
ScrollPaneConstants.HORIZONTAL_SCROLLBAR_NEVER
installFileDrop(ed.contentComponent, "editor")
installFileDrop(ed.scrollPane, "scroll")
syncHighlights()
ed.contentComponent.addFocusListener(object : FocusAdapter() {
override fun focusGained(e: FocusEvent) {
repaint()
@@ -224,6 +233,7 @@ class PromptPanel(
invalidateEnhancement()
syncEditorHeight()
triggerCompletion(e)
syncHighlights()
onChange()
}
})
@@ -336,6 +346,12 @@ class PromptPanel(
editor.background = style.editorScheme.defaultBackground
syncEditorHeight()
syncAutoApprove()
syncHighlights()
}
@RequiresEdt
fun refreshHighlights() {
syncHighlights()
}
@RequiresEdt
@@ -345,6 +361,33 @@ class PromptPanel(
completion?.clearMentions()
strip.clear()
syncEditorHeight()
syncHighlights()
}
@RequiresEdt
private fun syncHighlights() {
val provider = completion ?: return
val ed = editor.getEditor(false) ?: return
highlighters.forEach(ed.markupModel::removeHighlighter)
highlighters.clear()
val length = ed.document.textLength
provider.highlights(ed.document.text).forEach { item ->
val start = item.start.coerceIn(0, length)
val end = item.end.coerceIn(start, length)
if (start == end) return@forEach
highlighters += ed.markupModel.addRangeHighlighter(
key(item.kind),
start,
end,
HighlighterLayer.SYNTAX + 1,
HighlighterTargetArea.EXACT_RANGE,
)
}
}
private fun key(kind: KiloPromptCompletionProvider.HighlightKind): TextAttributesKey = when (kind) {
KiloPromptCompletionProvider.HighlightKind.MENTION -> MENTION_KEY
KiloPromptCompletionProvider.HighlightKind.COMMAND -> COMMAND_KEY
}
@RequiresEdt
packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/session/views/MessageView.kt
+11
View File
@@ -96,6 +96,14 @@ class MessageView(
sources.remove(content.id)
val stale = if (id == null) parts.remove(content.id) else null
if (stale != null) {
if (stale is PromptAttachmentView) {
stale.remove(content.id)
if (!stale.isEmpty()) {
refresh()
return
}
attachments = null
}
detach(stale)
remove(stale)
Disposer.dispose(stale)
@@ -232,6 +240,9 @@ class MessageView(
* pending/running question tool part linked to the active question.
*/
private fun isHidden(content: Content): Boolean {
if (role == SessionUiStyle.View.Message.USER_ROLE && content is FileAttachment) {
return content.source != null && content.mime.lowercase().startsWith("text/plain")
}
if (content !is Tool) return false
if (role == SessionUiStyle.View.Message.USER_ROLE && content.name == "read") return true
if (content.name == "todoread") return true
packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/session/model/SessionModelTest.kt
+76 -1
View File
@@ -9,6 +9,8 @@ import ai.kilocode.rpc.dto.MessageDto
import ai.kilocode.rpc.dto.MessageTimeDto
import ai.kilocode.rpc.dto.MessageWithPartsDto
import ai.kilocode.rpc.dto.PartDto
import ai.kilocode.rpc.dto.PartSourceDto
import ai.kilocode.rpc.dto.PartSourceTextDto
import ai.kilocode.rpc.dto.PartTimeDto
import ai.kilocode.rpc.dto.SessionDto
import ai.kilocode.rpc.dto.SessionTimeDto
@@ -154,6 +156,37 @@ class SessionModelTest : BasePlatformTestCase() {
assertEquals("ContentRemoved m1/p1", events.single().toString())
}
fun `test updateContent skips user synthetic text`() {
model.addMessage(msg("m1", "user"))
events.clear()
model.updateContent("m1", part("p1", "m1", "text", text = "raw content", synthetic = true))
assertNull(model.message("m1")!!.parts["p1"])
assertTrue(events.isEmpty())
}
fun `test updateContent removes existing user text when marked synthetic`() {
model.addMessage(msg("m1", "user"))
model.updateContent("m1", part("p1", "m1", "text", text = "visible"))
events.clear()
model.updateContent("m1", part("p1", "m1", "text", text = "hidden", synthetic = true))
assertNull(model.message("m1")!!.parts["p1"])
assertEquals("ContentRemoved m1/p1", events.single().toString())
}
fun `test assistant synthetic text remains visible`() {
model.addMessage(msg("m1", "assistant"))
events.clear()
model.updateContent("m1", part("p1", "m1", "text", text = "visible", synthetic = true))
assertEquals("visible", (model.message("m1")!!.parts["p1"] as Text).content.toString())
assertEquals("ContentAdded m1/p1", events.single().toString())
}
fun `test updateContent reasoning creates Reasoning content`() {
model.addMessage(msg("m1", "assistant"))
@@ -387,6 +420,17 @@ class SessionModelTest : BasePlatformTestCase() {
assertTrue(events[1] is SessionModelEvent.ContentDelta)
}
fun `test appendDelta ignores hidden synthetic user text part`() {
model.addMessage(msg("m1", "user"))
model.updateContent("m1", part("p1", "m1", "text", text = "hidden", synthetic = true))
events.clear()
model.appendDelta("m1", "p1", "still hidden")
assertNull(model.message("m1")!!.parts["p1"])
assertTrue(events.isEmpty())
}
fun `test appendDelta on tool content is noop`() {
model.addMessage(msg("m1", "assistant"))
model.updateContent("m1", part("p1", "m1", "tool", tool = "bash", state = "running"))
@@ -579,6 +623,32 @@ class SessionModelTest : BasePlatformTestCase() {
assertTrue(entry.parts["p3"] is StepFinish)
}
fun `test loadHistory skips user synthetic text`() {
model.loadHistory(listOf(MessageWithPartsDto(
msg("m1", "user"),
listOf(
part("p1", "m1", "text", text = "visible"),
part("p2", "m1", "text", text = "hidden", synthetic = true),
),
)))
assertModel("""
user#m1
text#p1:
visible
""")
}
fun `test file attachment preserves source metadata`() {
val source = PartSourceDto("file", PartSourceTextDto("@src/a.kt", 0.0, 9.0), path = "src/a.kt")
model.addMessage(msg("m1", "user"))
model.updateContent("m1", filePart("f1", "m1", "text/plain", "file:///tmp/a.kt", "a.kt", source))
val file = model.message("m1")!!.parts["f1"] as FileAttachment
assertEquals(source, file.source)
}
fun `test updateContent drops patch parts`() {
model.addMessage(msg("m1", "assistant"))
events.clear()
@@ -897,6 +967,8 @@ class SessionModelTest : BasePlatformTestCase() {
mime: String? = null,
url: String? = null,
filename: String? = null,
synthetic: Boolean? = null,
source: PartSourceDto? = null,
) = PartDto(
id = id,
sessionID = "ses",
@@ -906,6 +978,8 @@ class SessionModelTest : BasePlatformTestCase() {
mime = mime,
url = url,
filename = filename,
synthetic = synthetic,
source = source,
tool = tool,
state = state,
title = title,
@@ -921,13 +995,14 @@ class SessionModelTest : BasePlatformTestCase() {
tokens = tokens,
)
private fun filePart(id: String, mid: String, mime: String, url: String, filename: String) = part(
private fun filePart(id: String, mid: String, mime: String, url: String, filename: String, source: PartSourceDto? = null) = part(
id = id,
mid = mid,
type = "file",
mime = mime,
url = url,
filename = filename,
source = source,
)
private fun question(id: String) = Question(
packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/session/ui/PromptPanelTest.kt
+73
View File
@@ -1,16 +1,19 @@
package ai.kilocode.client.session.ui
import ai.kilocode.client.session.ui.style.SessionEditorStyle
import ai.kilocode.client.app.KiloWorkspaceService
import ai.kilocode.client.plugin.KiloBundle
import ai.kilocode.client.session.model.PromptAttachment
import ai.kilocode.client.session.ui.attachment.AttachmentCard
import ai.kilocode.client.session.ui.attachment.AttachmentCardItem
import ai.kilocode.client.session.ui.style.SessionUiStyle
import ai.kilocode.client.session.ui.prompt.KiloPromptCompletionProvider
import ai.kilocode.client.session.ui.prompt.PROMPT_ATTACHMENT_PASTE_HANDLER_KEY
import ai.kilocode.client.session.ui.prompt.PromptAttachmentPasteHandler
import ai.kilocode.client.session.ui.prompt.PromptAttachmentPasteProvider
import ai.kilocode.client.session.ui.prompt.PromptDataKeys
import ai.kilocode.client.session.ui.prompt.PromptPanel
import ai.kilocode.client.testing.FakeWorkspaceRpcApi
import com.intellij.icons.AllIcons
import com.intellij.notification.Notification
import com.intellij.notification.Notifications
@@ -19,6 +22,7 @@ import com.intellij.openapi.actionSystem.DataContext
import com.intellij.openapi.actionSystem.DataSink
import com.intellij.openapi.actionSystem.UiDataProvider
import com.intellij.openapi.application.ApplicationManager
import com.intellij.openapi.editor.DefaultLanguageHighlighterColors
import com.intellij.openapi.editor.Editor
import com.intellij.openapi.editor.EditorFactory
import com.intellij.openapi.editor.actions.PasteAction
@@ -32,7 +36,11 @@ import com.intellij.util.Producer
import com.intellij.util.ui.EmptyIcon
import com.intellij.util.ui.JBUI
import com.intellij.util.ui.UIUtil
import kotlinx.coroutines.CoroutineScope
import kotlinx.coroutines.CancellationException
import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.SupervisorJob
import kotlinx.coroutines.cancel
import java.awt.Container
import java.awt.BorderLayout
import java.awt.datatransfer.DataFlavor
@@ -53,11 +61,22 @@ import javax.swing.SwingUtilities
@Suppress("UnstableApiUsage")
class PromptPanelTest : BasePlatformTestCase() {
private val roots = mutableListOf<SessionRootPanel>()
private lateinit var scope: CoroutineScope
private lateinit var rpc: FakeWorkspaceRpcApi
private lateinit var workspaces: KiloWorkspaceService
override fun setUp() {
super.setUp()
scope = CoroutineScope(SupervisorJob() + Dispatchers.Default)
rpc = FakeWorkspaceRpcApi()
workspaces = KiloWorkspaceService(scope, rpc)
}
override fun tearDown() {
try {
roots.asReversed().forEach { it.removeNotify() }
roots.clear()
scope.cancel()
} finally {
super.tearDown()
}
@@ -197,6 +216,47 @@ class PromptPanelTest : BasePlatformTestCase() {
assertFalse(editor.settings.isPaintSoftWraps)
}
fun `test prompt editor highlights validated commands and mentions`() {
val panel = PromptPanel(project = project, onSend = { _, _ -> }, onAbort = {}, onEnhance = { _, _ -> }, completion = completion())
val field = panel.defaultFocusedComponent as EditorTextField
realize(panel, 260, 400)
field.text = "/new use @terminal and @unknown"
UIUtil.dispatchAllInvocationEvents()
val spans = spans(field)
assertTrue(spans.contains("/new" to DefaultLanguageHighlighterColors.KEYWORD))
assertTrue(spans.contains("@terminal" to DefaultLanguageHighlighterColors.METADATA))
assertFalse(spans.any { it.first == "@unknown" })
}
fun `test prompt clear removes prompt highlighters`() {
val panel = PromptPanel(project = project, onSend = { _, _ -> }, onAbort = {}, onEnhance = { _, _ -> }, completion = completion())
val field = panel.defaultFocusedComponent as EditorTextField
realize(panel, 260, 400)
field.text = "use @terminal"
UIUtil.dispatchAllInvocationEvents()
assertEquals(1, field.getEditor(false)!!.markupModel.allHighlighters.size)
panel.clear()
UIUtil.dispatchAllInvocationEvents()
assertEquals(0, field.getEditor(false)!!.markupModel.allHighlighters.size)
}
fun `test prompt highlighters stay bounded across edits`() {
val panel = PromptPanel(project = project, onSend = { _, _ -> }, onAbort = {}, onEnhance = { _, _ -> }, completion = completion())
val field = panel.defaultFocusedComponent as EditorTextField
realize(panel, 260, 400)
repeat(50) {
field.text = if (it % 2 == 0) "/new @terminal" else "/new @git-changes"
UIUtil.dispatchAllInvocationEvents()
assertTrue(field.getEditor(false)!!.markupModel.allHighlighters.size <= 2)
}
}
fun `test prompt editor shrinks when lines are removed`() {
val panel = PromptPanel(project = project, onSend = { _, _ -> }, onAbort = {}, onEnhance = { _, _ -> })
val editor = panel.defaultFocusedComponent as EditorTextField
@@ -732,6 +792,19 @@ class PromptPanelTest : BasePlatformTestCase() {
return root
}
private fun completion() = KiloPromptCompletionProvider(
workspace = workspaces.workspace("/test"),
service = workspaces,
actions = listOf(KiloPromptCompletionProvider.SlashAction("new", "New") {}),
)
private fun spans(field: EditorTextField): List<Pair<String, com.intellij.openapi.editor.colors.TextAttributesKey?>> {
val editor = field.getEditor(false)!!
return editor.markupModel.allHighlighters.map {
field.text.substring(it.startOffset, it.endOffset) to it.textAttributesKey
}
}
private fun createEditor(): Editor {
val factory = EditorFactory.getInstance()
return factory.createEditor(factory.createDocument(""), project)
packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/session/ui/SessionUiUpdateTest.kt
+76
View File
@@ -15,6 +15,8 @@ import ai.kilocode.rpc.dto.MessageDto
import ai.kilocode.rpc.dto.MessageTimeDto
import ai.kilocode.rpc.dto.MessageWithPartsDto
import ai.kilocode.rpc.dto.PartDto
import ai.kilocode.rpc.dto.PartSourceDto
import ai.kilocode.rpc.dto.PartSourceTextDto
import com.intellij.openapi.Disposable
import com.intellij.openapi.util.Disposer
import com.intellij.testFramework.fixtures.BasePlatformTestCase
@@ -264,6 +266,74 @@ class SessionUiUpdateTest : BasePlatformTestCase() {
assertEquals(listOf("data:image/png;base64,aGVsbG8=", "data:text/plain;base64,aGVsbG8="), opened)
}
fun `test user file mention hides synthetic read payload and text attachment card`() {
model.upsertMessage(msg("u1", "user"))
model.updateContent("u1", part("p1", "u1", "text", text = "read @src/a.kt"))
model.updateContent("u1", PartDto(
id = "p2",
sessionID = "ses",
messageID = "u1",
type = "text",
text = "<path>/tmp/a.kt</path>\n<content>raw</content>",
synthetic = true,
))
model.updateContent("u1", PartDto(
id = "f1",
sessionID = "ses",
messageID = "u1",
type = "file",
mime = "text/plain",
url = "file:///tmp/a.kt",
filename = "a.kt",
source = source("src/a.kt"),
))
val msg = panel.findMessage("u1")!!
assertEquals(listOf("p1"), msg.partIds())
assertNull(msg.part("p2"))
assertNull(msg.part("f1"))
assertEquals("read @src/a.kt", (msg.part("p1") as TextView).markdown())
assertEquals(0, msg.components.filterIsInstance<PromptAttachmentView>().size)
}
fun `test source less text attachment still renders in prompt strip`() {
model.upsertMessage(msg("u1", "user"))
model.updateContent("u1", PartDto(
id = "f1",
sessionID = "ses",
messageID = "u1",
type = "file",
mime = "text/plain",
url = "data:text/plain;base64,aGVsbG8=",
filename = "note.txt",
))
val view = panel.findMessage("u1")!!.part("f1")
assertTrue(view is PromptAttachmentView)
assertNotNull(find(view!!, AttachmentCard::class.java))
}
fun `test source backed image attachment still renders in prompt strip`() {
model.upsertMessage(msg("u1", "user"))
model.updateContent("u1", PartDto(
id = "f1",
sessionID = "ses",
messageID = "u1",
type = "file",
mime = "image/png",
url = "file:///tmp/a.png",
filename = "a.png",
source = source("src/a.png"),
))
val view = panel.findMessage("u1")!!.part("f1")
assertTrue(view is PromptAttachmentView)
assertNotNull(find(view!!, AttachmentCard::class.java))
}
fun `test empty sanitized user text does not create prompt panel`() {
model.upsertMessage(msg("u1", "user"))
model.updateContent("u1", part("p1", "u1", "text", text = "read these screenshots"))
@@ -461,6 +531,12 @@ class SessionUiUpdateTest : BasePlatformTestCase() {
id = id, sessionID = "ses", messageID = mid, type = "tool", tool = tool, state = state,
)
private fun source(path: String) = PartSourceDto(
type = "file",
path = path,
text = PartSourceTextDto("@$path", 5.0, (6 + path.length).toDouble()),
)
private fun <T : Any> find(root: java.awt.Component, type: Class<T>): T? {
if (type.isInstance(root)) return type.cast(root)
if (root is Container) {
packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/session/ui/prompt/KiloPromptCompletionProviderTest.kt
+66 -1
View File
@@ -2,7 +2,10 @@ package ai.kilocode.client.session.ui.prompt
import ai.kilocode.client.app.KiloWorkspaceService
import ai.kilocode.client.testing.FakeWorkspaceRpcApi
import ai.kilocode.rpc.dto.CommandDto
import ai.kilocode.rpc.dto.FileSearchResultDto
import ai.kilocode.rpc.dto.KiloWorkspaceStateDto
import ai.kilocode.rpc.dto.KiloWorkspaceStatusDto
import ai.kilocode.rpc.dto.WorkspaceFileDto
import com.intellij.testFramework.fixtures.BasePlatformTestCase
import com.intellij.util.textCompletion.TextCompletionUtil
@@ -25,7 +28,7 @@ class KiloPromptCompletionProviderTest : BasePlatformTestCase() {
provider = KiloPromptCompletionProvider(
workspace = workspaces.workspace("/test"),
service = workspaces,
actions = emptyList(),
actions = listOf(KiloPromptCompletionProvider.SlashAction("new", "New") {}),
)
}
@@ -74,11 +77,73 @@ class KiloPromptCompletionProviderTest : BasePlatformTestCase() {
assertEquals(listOf("git"), rpc.searchQueries)
}
fun `test highlights known slash command at start`() {
assertEquals(
listOf(KiloPromptCompletionProvider.Highlight(0, 4, KiloPromptCompletionProvider.HighlightKind.COMMAND)),
provider.highlights("/new start fresh"),
)
}
fun `test highlights server slash command at start`() {
rpc.state.value = KiloWorkspaceStateDto(KiloWorkspaceStatusDto.READY, commands = listOf(CommandDto("deploy")))
waitFor { provider.highlights("/deploy prod").isNotEmpty() }
assertEquals(
listOf(KiloPromptCompletionProvider.Highlight(0, 7, KiloPromptCompletionProvider.HighlightKind.COMMAND)),
provider.highlights("/deploy prod"),
)
}
fun `test highlights ignore unknown and non-leading slash commands`() {
assertTrue(provider.highlights("/bogus now").isEmpty())
assertTrue(provider.highlights("hi /new").isEmpty())
}
fun `test highlights special mentions without tracked paths`() {
assertEquals(
listOf(
KiloPromptCompletionProvider.Highlight(4, 13, KiloPromptCompletionProvider.HighlightKind.MENTION),
KiloPromptCompletionProvider.Highlight(18, 30, KiloPromptCompletionProvider.HighlightKind.MENTION),
),
provider.highlights("use @terminal and @git-changes").sortedBy { it.start },
)
}
fun `test highlights tracked mentions longest first`() {
addMention("src/a.ts", "@ts")
addMention("src/a.tsx", "@tsx")
assertEquals(
listOf(KiloPromptCompletionProvider.Highlight(4, 14, KiloPromptCompletionProvider.HighlightKind.MENTION)),
provider.highlights("see @src/a.tsx"),
)
}
fun `test highlights ignore untracked mentions`() {
assertTrue(provider.highlights("see @unknownPath").isEmpty())
}
private fun complete(text: String) {
val file = myFixture.configureByText("prompt.txt", text)
TextCompletionUtil.installProvider(file, provider, true)
myFixture.completeBasic()
}
private fun addMention(path: String, query: String) {
rpc.searchResult = FileSearchResultDto(files = listOf(file(path)))
complete("$query<caret>")
myFixture.type('\n')
assertTrue(provider.mentionPaths().contains(path))
}
private fun waitFor(done: () -> Boolean) {
repeat(50) {
com.intellij.util.ui.UIUtil.dispatchAllInvocationEvents()
if (done()) return
Thread.sleep(20)
}
}
private fun file(path: String) = WorkspaceFileDto(path = path, name = path.substringAfterLast('/'))
}
packages/kilo-jetbrains/shared/src/main/kotlin/ai/kilocode/rpc/dto/ChatDto.kt
+21
View File
@@ -75,6 +75,26 @@ data class PartDto(
val mime: String? = null,
val url: String? = null,
val filename: String? = null,
val synthetic: Boolean? = null,
val source: PartSourceDto? = null,
)
@Serializable
data class PartSourceDto(
val type: String,
val text: PartSourceTextDto,
val path: String? = null,
val clientName: String? = null,
val uri: String? = null,
val name: String? = null,
val kind: Int? = null,
)
@Serializable
data class PartSourceTextDto(
val value: String,
val start: Double,
val end: Double,
)
@Serializable
@@ -103,6 +123,7 @@ data class PromptPartDto(
val mime: String? = null,
val url: String? = null,
val filename: String? = null,
val source: PartSourceDto? = null,
)
// --- Streaming Events ---