Skip to content

Commit 6ae3fff

Browse files
carlrobertohcarlrobertoh
committed
feat: redesign write and edit tool calls with diff previews
1 parent d25e5a7 commit 6ae3fff

17 files changed

Lines changed: 1498 additions & 1682 deletions

‎src/main/kotlin/ee/carlrobert/codegpt/agent/SubagentTools.kt‎

Lines changed: 0 additions & 48 deletions
This file was deleted.

‎src/main/kotlin/ee/carlrobert/codegpt/toolwindow/agent/AgentEventHandler.kt‎

Lines changed: 91 additions & 40 deletions
Original file line numberDiff line numberDiff line change
@@ -19,18 +19,16 @@ import ee.carlrobert.codegpt.agent.*
1919
import ee.carlrobert.codegpt.agent.history.CheckpointRef
2020
import ee.carlrobert.codegpt.agent.rollback.RollbackService
2121
import ee.carlrobert.codegpt.agent.tools.*
22-
import ee.carlrobert.codegpt.agent.tools.ide.BreakpointTool
23-
import ee.carlrobert.codegpt.agent.tools.ide.ExecuteRunConfigurationTool
24-
import ee.carlrobert.codegpt.agent.tools.ide.GetBreakpointsTool
25-
import ee.carlrobert.codegpt.agent.tools.ide.GetDebugSessionsTool
26-
import ee.carlrobert.codegpt.agent.tools.ide.GetRunOutputTool
22+
import ee.carlrobert.codegpt.agent.tools.ide.*
2723
import ee.carlrobert.codegpt.completions.ToolApprovalMode
2824
import ee.carlrobert.codegpt.settings.agents.SubagentRuntimeResolver
2925
import ee.carlrobert.codegpt.settings.service.ServiceType
3026
import ee.carlrobert.codegpt.settings.service.codegpt.CodeGPTApiException
3127
import ee.carlrobert.codegpt.toolwindow.agent.ui.*
3228
import ee.carlrobert.codegpt.toolwindow.agent.ui.approval.*
3329
import ee.carlrobert.codegpt.toolwindow.agent.ui.descriptor.Badge
30+
import ee.carlrobert.codegpt.toolwindow.agent.ui.descriptor.FileChangeSnapshot
31+
import ee.carlrobert.codegpt.toolwindow.agent.ui.descriptor.ToolCallDiffPreview
3432
import ee.carlrobert.codegpt.toolwindow.agent.ui.descriptor.ToolKind
3533
import ee.carlrobert.codegpt.toolwindow.agent.ui.renderer.*
3634
import ee.carlrobert.codegpt.toolwindow.chat.ui.ChatMessageResponseBody
@@ -67,6 +65,7 @@ class AgentEventHandler(
6765
}
6866

6967
private val mainToolCards = ConcurrentHashMap<String, ToolCallCard>()
68+
private val fileChangeSnapshots = ConcurrentHashMap<String, FileChangeSnapshot>()
7069
private val pendingToolOutput = ConcurrentHashMap<String, MutableList<ToolOutputLine>>()
7170
private val scheduledToolOutputFlushes =
7271
Collections.newSetFromMap(ConcurrentHashMap<String, Boolean>())
@@ -90,9 +89,15 @@ class AgentEventHandler(
9089
@Volatile
9190
private var lastWriteArgs: WriteTool.Args? = null
9291

92+
@Volatile
93+
private var lastWriteSnapshot: FileChangeSnapshot? = null
94+
9395
@Volatile
9496
private var lastEditArgs: EditTool.Args? = null
9597

98+
@Volatile
99+
private var lastEditSnapshot: FileChangeSnapshot? = null
100+
96101
@Volatile
97102
private var currentRollbackRunId: String? = null
98103

@@ -133,11 +138,14 @@ class AgentEventHandler(
133138

134139
fun resetForNewSubmission() {
135140
mainToolCards.clear()
141+
fileChangeSnapshots.clear()
136142
pendingToolOutput.clear()
137143
scheduledToolOutputFlushes.clear()
138144
currentResponseBody = null
139145
lastWriteArgs = null
146+
lastWriteSnapshot = null
140147
lastEditArgs = null
148+
lastEditSnapshot = null
141149
currentApproval = null
142150
approvalQueue.clear()
143151
currentQuestion = null
@@ -417,74 +425,78 @@ class AgentEventHandler(
417425
}
418426

419427
override fun onToolStarting(id: String, toolName: String, args: Any?) {
420-
val uiArgs = ToolSpecs.coerceArgsForUi(toolName, args)
421-
if (toolName == "Tool" && uiArgs == null) {
428+
if (toolName == "Tool" && args == null) {
422429
logger.debug("Deferring placeholder tool card session=$sessionId toolId=$id toolName=$toolName")
423430
return
424431
}
425432

426-
when (uiArgs) {
433+
when (args) {
427434
is TodoWriteTool.Args -> {
428435
runInEdt {
429436
val inProgressTask =
430-
uiArgs.todos.find { it.status == TodoWriteTool.TodoStatus.IN_PROGRESS }
437+
args.todos.find { it.status == TodoWriteTool.TodoStatus.IN_PROGRESS }
431438
if (inProgressTask != null) {
432439
showLoading(inProgressTask.activeForm)
433440
}
434-
todoListPanel.updateTodos(uiArgs.todos)
441+
todoListPanel.updateTodos(args.todos)
435442
todoListPanel.isVisible = true
436-
showMainToolCard(id, toolName, uiArgs)
443+
showMainToolCard(id, toolName, args)
437444
}
438445
}
439446

440447
is TaskTool.Args -> {
441448
runInEdt {
442449
val host = ensureRunViewForSubagent(id)
443-
host.addEntry(createTaskEntry(id, null, uiArgs))
450+
host.addEntry(createTaskEntry(id, null, args))
444451
host.refresh()
445452
requestUiRefresh()
446453
}
447454
}
448455

449456
else -> {
450-
when (uiArgs) {
451-
is EditTool.Args -> {
452-
trackEditOperation(uiArgs)
457+
val fileChangeSnapshot = when (args) {
458+
is EditTool.Args -> captureEditSnapshot(args).also {
459+
fileChangeSnapshots[keyFor(id)] = it
460+
lastEditSnapshot = it
461+
trackEditOperation(args)
453462
}
454463

455-
is WriteTool.Args -> {
456-
trackWriteOperation(uiArgs)
464+
is WriteTool.Args -> captureWriteSnapshot(args).also {
465+
fileChangeSnapshots[keyFor(id)] = it
466+
lastWriteSnapshot = it
467+
trackWriteOperation(args)
457468
}
458-
}
459469

470+
else -> null
471+
}
460472
runInEdt {
461-
showMainToolCard(id, toolName, uiArgs)
473+
showMainToolCard(id, toolName, args, fileChangeSnapshot)
462474
}
463475
}
464476
}
465477
}
466478

467479
override fun onToolCompleted(id: String?, toolName: String, result: Any?) {
468-
val uiResult = ToolSpecs.coerceResultForUi(toolName, result)
469480
runInEdt {
470-
if (id != null && (toolName == "Task" || uiResult is TaskTool.Result)) {
481+
if (id != null && (toolName == "Task" || result is TaskTool.Result)) {
471482
val holder = runViewHolder ?: subagentViewHolders.values.firstOrNull { viewHolder ->
472483
viewHolder.getItems().any { entry -> entry.id == id }
473484
}
474-
holder?.completeEntry(id, uiResult)
485+
holder?.completeEntry(id, result)
475486
holder?.refresh()
476487
} else if (id != null && mainToolCards.containsKey(keyFor(id))) {
477-
val success = uiResult !is ToolError && uiResult != null
478-
mainToolCards[keyFor(id)]?.complete(success, uiResult)
488+
val success = result !is ToolError && result != null
489+
mainToolCards[keyFor(id)]?.complete(success, result)
490+
fileChangeSnapshots.remove(keyFor(id))
479491
mainToolCards[keyFor(id)]?.getDescriptor()?.let { descriptor ->
480492
if (descriptor.kind == ToolKind.OTHER || descriptor.titleMain.isBlank()) {
481493
logger.warn(
482-
"Completed generic tool card session=$sessionId toolId=$id toolName=$toolName resultType=${uiResult?.javaClass?.name ?: "null"} title=${descriptor.titleMain}"
494+
"Completed generic tool card session=$sessionId toolId=$id toolName=$toolName resultType=${result?.javaClass?.name ?: "null"} title=${descriptor.titleMain}"
483495
)
484496
}
485497
}
486498

487-
val bgId = (uiResult as? BashTool.Result)?.bashId
499+
val bgId = (result as? BashTool.Result)?.bashId
488500
if (bgId == null) {
489501
mainToolCards.remove(keyFor(id))
490502
} else {
@@ -543,12 +555,16 @@ class AgentEventHandler(
543555

544556
is WriteTool.Args -> {
545557
lastWriteArgs = args
546-
RunEntry.WriteEntry(cid, parentId, args, null)
558+
val snapshot = captureWriteSnapshot(args)
559+
lastWriteSnapshot = snapshot
560+
RunEntry.WriteEntry(cid, parentId, args, null, snapshot)
547561
}
548562

549563
is EditTool.Args -> {
550564
lastEditArgs = args
551-
RunEntry.EditEntry(cid, parentId, args, null)
565+
val snapshot = captureEditSnapshot(args)
566+
lastEditSnapshot = snapshot
567+
RunEntry.EditEntry(cid, parentId, args, null, snapshot)
552568
}
553569

554570
is TaskTool.Args -> createTaskEntry(cid, parentId, args)
@@ -611,11 +627,17 @@ class AgentEventHandler(
611627
}
612628
}
613629

614-
private fun showMainToolCard(id: String, toolName: String, uiArgs: Any?) {
630+
private fun showMainToolCard(
631+
id: String,
632+
toolName: String,
633+
uiArgs: Any?,
634+
fileChangeSnapshot: FileChangeSnapshot? = null
635+
) {
615636
val key = keyFor(id)
616637
val existingCard = mainToolCards[key]
617638
if (existingCard == null) {
618-
val card = ToolCallCard(project, toolName, uiArgs)
639+
val card =
640+
ToolCallCard(project, toolName, uiArgs, fileChangeSnapshot = fileChangeSnapshot)
619641
mainToolCards[key] = card
620642
currentResponseBody?.addToolStatusPanel(card)
621643
requestUiRefresh()
@@ -838,14 +860,19 @@ class AgentEventHandler(
838860
val payload = request.payload ?: return
839861
val (path, before, after) = when (payload) {
840862
is EditPayload -> {
841-
val currentContent = getFileContentWithFallback(payload.filePath)
842-
val proposed = payload.proposedContent ?: applyStringReplacement(
843-
currentContent,
844-
payload.oldString,
845-
payload.newString,
846-
payload.replaceAll
847-
)
848-
Triple(payload.filePath, currentContent, proposed)
863+
val snapshot = lastEditSnapshot
864+
if (snapshot != null) {
865+
Triple(payload.filePath, snapshot.beforeText, snapshot.afterText)
866+
} else {
867+
val currentContent = getFileContentWithFallback(payload.filePath)
868+
val proposed = payload.proposedContent ?: applyStringReplacement(
869+
currentContent,
870+
payload.oldString,
871+
payload.newString,
872+
payload.replaceAll
873+
)
874+
Triple(payload.filePath, currentContent, proposed)
875+
}
849876
}
850877

851878
else -> return
@@ -868,14 +895,15 @@ class AgentEventHandler(
868895
val nonDiffBadges = descriptor.secondaryBadges.filterNot { isDiffBadge(it) }
869896
descriptor.copy(
870897
secondaryBadges = nonDiffBadges + diffBadges,
871-
summary = null
898+
summary = null,
899+
diffPreview = ToolCallDiffPreview(path, FileChangeSnapshot(before, after))
872900
)
873901
}
874902
}
875903

876904
private fun isDiffBadge(badge: Badge): Boolean {
877905
val text = badge.text
878-
return text.startsWith("[+") || text.startsWith("[-") || text.startsWith("[~")
906+
return text.startsWith("+") || text.startsWith("-") || text.startsWith("~")
879907
}
880908

881909
private fun maybeShowNextQuestion() {
@@ -993,6 +1021,7 @@ class AgentEventHandler(
9931021
serviceScope.dispose()
9941022
agentApprovalManager.dispose()
9951023
mainToolCards.clear()
1024+
fileChangeSnapshots.clear()
9961025
synchronized(toolOutputLock) {
9971026
pendingToolOutput.clear()
9981027
scheduledToolOutputFlushes.clear()
@@ -1089,4 +1118,26 @@ class AgentEventHandler(
10891118
rollbackService.trackWrite(sessionId, normalizedPath)
10901119
}
10911120
}
1121+
1122+
private fun captureEditSnapshot(args: EditTool.Args): FileChangeSnapshot {
1123+
val normalizedPath = args.filePath.replace("\\", "/")
1124+
val currentContent = getFileContentWithFallback(normalizedPath)
1125+
val proposedContent = applyStringReplacement(
1126+
currentContent,
1127+
args.oldString,
1128+
args.newString,
1129+
args.replaceAll
1130+
)
1131+
return FileChangeSnapshot(currentContent, proposedContent)
1132+
}
1133+
1134+
private fun captureWriteSnapshot(args: WriteTool.Args): FileChangeSnapshot {
1135+
val normalizedPath = args.filePath.replace("\\", "/")
1136+
val beforeContent = getFileContentWithFallback(normalizedPath)
1137+
return FileChangeSnapshot(
1138+
beforeText = beforeContent,
1139+
afterText = args.content,
1140+
isNewFile = beforeContent.isEmpty() && !java.io.File(normalizedPath).exists()
1141+
)
1142+
}
10921143
}

0 commit comments

Comments
 (0)