fix: tooltip agent total tokens and other minor adjustments

This commit is contained in:
Carl-Robert Linnupuu 2026-01-25 17:55:49 +00:00
parent ebea17ff6b
commit 1a9bd1143f
6 changed files with 179 additions and 50 deletions

View file

@ -16,7 +16,6 @@ import kotlinx.coroutines.flow.MutableSharedFlow
import kotlinx.coroutines.flow.asSharedFlow
import java.util.*
import java.util.concurrent.ConcurrentHashMap
import kotlin.coroutines.cancellation.CancellationException
import kotlin.io.path.Path
@Service(Service.Level.PROJECT)
@ -121,4 +120,4 @@ class AgentService(private val project: Project) {
fun getTokenTrackerForSession(sessionId: String): TokenUsageTracker {
return sessionTokenTrackers.getOrPut(sessionId) { TokenUsageTracker() }
}
}
}

View file

@ -205,8 +205,7 @@ private suspend fun AIAgentLLMWriteSession.requestResponses(
): List<Message.Response> {
val responses = if (stream) {
val streamFrames = mutableListOf<StreamFrame>()
val frames = requestLLMStreaming()
frames.collect { streamFrames.add(it) }
requestLLMStreaming().collect { streamFrames.add(it) }
streamFrames.toMessageResponses()
} else {
val preparedPrompt = config.missingToolsConversionStrategy.convertPrompt(prompt, tools)

View file

@ -24,6 +24,8 @@ import kotlinx.coroutines.Dispatchers
import kotlinx.serialization.Serializable
import java.nio.file.Paths
import com.intellij.openapi.util.TextRange
import kotlinx.coroutines.TimeoutCancellationException
import kotlinx.coroutines.withTimeout
/**
* Enhanced search tool using IntelliJ's native SearchService and FindModel.
@ -127,12 +129,13 @@ class IntelliJSearchTool(
override suspend fun execute(args: Args): Result {
try {
val (searchScope, matches) = withContext(Dispatchers.Default) {
runReadAction {
val scope = createSearchScope(args, project)
val effectiveLimit = (args.limit ?: 10).coerceAtLeast(1)
val results = searchEverywhere(args.pattern, scope, effectiveLimit)
scope to results
val maxResults = (args.limit ?: 10).coerceIn(1, 50)
val matches = withTimeout(5000) {
withContext(Dispatchers.Default) {
runReadAction {
val scope = createSearchScope(args, project)
searchEverywhere(args.pattern, scope, maxResults)
}
}
}
val output = formatOutput(matches, args)
@ -144,6 +147,14 @@ class IntelliJSearchTool(
matches = matches,
output = output
)
} catch (_: TimeoutCancellationException) {
return Result(
pattern = args.pattern,
scope = args.scope ?: "project",
totalMatches = 0,
matches = emptyList(),
output = "Search timed out. Try a more specific pattern or lower scope."
)
} catch (e: Exception) {
return Result(
pattern = args.pattern,

View file

@ -20,9 +20,13 @@ import ee.carlrobert.codegpt.agent.rollback.*
import ee.carlrobert.codegpt.toolwindow.agent.ui.renderer.ChangeColors
import ee.carlrobert.codegpt.toolwindow.agent.ui.renderer.lineDiffStats
import ee.carlrobert.codegpt.ui.IconActionButton
import ee.carlrobert.codegpt.ui.components.LeftEllipsisLabel
import kotlinx.coroutines.*
import java.awt.BorderLayout
import java.awt.Dimension
import java.awt.FlowLayout
import java.awt.event.ComponentAdapter
import java.awt.event.ComponentEvent
import java.time.Instant
import java.time.ZoneId
import java.time.format.DateTimeFormatter
@ -79,21 +83,26 @@ class RollbackPanel(
changesPanel.apply {
layout = BoxLayout(this, BoxLayout.Y_AXIS)
isOpaque = false
alignmentX = LEFT_ALIGNMENT
addComponentListener(object : ComponentAdapter() {
override fun componentResized(e: ComponentEvent?) {
revalidate()
repaint()
}
})
}
scrollPane.apply {
horizontalScrollBarPolicy = JScrollPane.HORIZONTAL_SCROLLBAR_NEVER
verticalScrollBarPolicy = JScrollPane.VERTICAL_SCROLLBAR_AS_NEEDED
border = null
preferredSize = Dimension(0, 240)
viewportBorder = null
viewport.isOpaque = false
isOpaque = false
}
addToTop(topPanel)
addToCenter(
BorderLayoutPanel().apply {
addToCenter(scrollPane)
}
)
addToCenter(scrollPane)
border = JBUI.Borders.compound(
JBUI.Borders.customLine(JBColor.border(), 0, 0, 1, 0),
JBUI.Borders.empty(0, 0, 8, 0)
@ -171,22 +180,11 @@ class RollbackPanel(
}
private fun createChangeRow(change: FileChange): JComponent {
val row = BorderLayoutPanel().apply {
isOpaque = true
background = JBUI.CurrentTheme.List.background(false, false)
border = JBUI.Borders.compound(
JBUI.Borders.customLine(JBColor.border(), 1),
JBUI.Borders.empty(4, 8)
)
}
val left = JPanel(FlowLayout(FlowLayout.LEFT, 6, 2)).apply {
val leftFixed = JPanel(FlowLayout(FlowLayout.LEFT, 6, 0)).apply {
isOpaque = false
add(changeLabel(change))
add(fileNameComponent(change))
}
left.add(changeLabel(change))
left.add(fileNameComponent(change))
left.add(filePathLabel(change))
addDiffStats(change, left)
val actions = JPanel(FlowLayout(FlowLayout.RIGHT, 6, 0)).apply {
isOpaque = false
@ -194,8 +192,45 @@ class RollbackPanel(
openDiffAction(change).let { actions.add(it) }
actions.add(rollbackAction(change))
row.addToLeft(left)
row.addToRight(actions)
val description = JPanel(BorderLayout()).apply {
isOpaque = false
add(filePathLabel(change), BorderLayout.CENTER)
}
val stats = JPanel(FlowLayout(FlowLayout.LEFT, 6, 0)).apply {
isOpaque = false
}
addDiffStats(change, stats)
val rightFixed = JPanel(FlowLayout(FlowLayout.RIGHT, 6, 0)).apply {
isOpaque = false
add(stats)
add(actions)
}
val row = JPanel(BorderLayout(8, 0)).apply {
isOpaque = true
background = JBUI.CurrentTheme.List.background(false, false)
border = JBUI.Borders.compound(
JBUI.Borders.customLine(JBColor.border(), 1),
JBUI.Borders.empty(4, 8)
)
add(leftFixed, BorderLayout.WEST)
add(description, BorderLayout.CENTER)
add(rightFixed, BorderLayout.EAST)
alignmentX = LEFT_ALIGNMENT
}
description.minimumSize = Dimension(0, description.minimumSize.height)
description.maximumSize = Dimension(Int.MAX_VALUE, description.maximumSize.height)
description.preferredSize = Dimension(0, description.preferredSize.height)
row.minimumSize = Dimension(0, row.minimumSize.height)
row.maximumSize = Dimension(Int.MAX_VALUE, row.preferredSize.height)
row.preferredSize = Dimension(0, row.preferredSize.height)
return row
}
@ -230,10 +265,9 @@ class RollbackPanel(
private fun filePathLabel(change: FileChange): JBLabel {
val display = displayPath(change.path, change)
return JBLabel(display).apply {
return LeftEllipsisLabel(display).apply {
foreground = JBUI.CurrentTheme.Label.disabledForeground()
font = JBUI.Fonts.smallFont()
toolTipText = display
}
}
@ -322,7 +356,7 @@ class RollbackPanel(
componentList.forEach { component ->
if (visibleRows >= maxVisibleItems) return@forEach
height += component.preferredSize.height
if (component is BorderLayoutPanel) {
if (component is JComponent && component !is Box.Filler) {
visibleRows += 1
}
}

View file

@ -0,0 +1,97 @@
package ee.carlrobert.codegpt.ui.components
import com.intellij.ui.components.JBLabel
import java.awt.Dimension
import java.awt.FontMetrics
import java.awt.event.ComponentAdapter
import java.awt.event.ComponentEvent
/**
* A label that keeps the *end* of the text visible by eliding from the left.
*/
class LeftEllipsisLabel(text: String = "") : JBLabel() {
var fullText: String = text
set(value) {
field = value
updateDisplayedText()
}
init {
super.setText(text)
fullText = text
addComponentListener(object : ComponentAdapter() {
override fun componentResized(e: ComponentEvent?) {
updateDisplayedText()
}
})
}
override fun setText(text: String?) {
fullText = text.orEmpty()
}
override fun doLayout() {
super.doLayout()
updateDisplayedText()
}
override fun getMinimumSize(): Dimension {
val size = super.getMinimumSize()
return Dimension(0, size.height)
}
override fun getPreferredSize(): Dimension {
val size = super.getPreferredSize()
val fm = getFontMetrics(font) ?: return size
val ellipsisWidth = fm.stringWidth("...")
val maxWidth = fm.stringWidth(fullText)
return Dimension(minOf(size.width, maxWidth + ellipsisWidth), size.height)
}
private fun updateDisplayedText() {
val availableWidth = width
if (availableWidth <= 0) {
super.setText(fullText)
return
}
val fm = getFontMetrics(font) ?: return
super.setText(leftEllipsize(fullText, fm, availableWidth))
toolTipText = fullText
}
private fun leftEllipsize(text: String, fm: FontMetrics, maxWidth: Int): String {
if (text.isEmpty()) return text
if (fm.stringWidth(text) <= maxWidth) return text
val ellipsis = "..."
val ellipsisWidth = fm.stringWidth(ellipsis)
if (ellipsisWidth >= maxWidth) {
return ""
}
var lo = 0
var hi = text.length
while (lo < hi) {
val mid = (lo + hi) / 2
val candidate = ellipsis + text.substring(mid)
if (fm.stringWidth(candidate) <= maxWidth) {
hi = mid
} else {
lo = mid + 1
}
}
val startIndex = lo.coerceIn(0, text.length)
val result = ellipsis + text.substring(startIndex)
if (fm.stringWidth(result) <= maxWidth) return result
for (i in startIndex + 1..text.length) {
val r = ellipsis + text.substring(i)
if (fm.stringWidth(r) <= maxWidth) return r
}
return ""
}
}

View file

@ -59,7 +59,7 @@ class TokenUsageCounterPanel(
currentSessionId = event.sessionId
val model = getAgentModelForSession(event.sessionId) ?: return
updateDisplay(event, model)
updateTooltipText(model)
updateTooltipText(event, model)
}
}
})
@ -69,8 +69,7 @@ class TokenUsageCounterPanel(
private fun updateDisplay(event: TokenUsageEvent, model: LLModel) {
scope.launch {
withEdt {
val usedPromptTokens = getPromptTokensForSession(event.sessionId) ?: 0
updateColorAndText(usedPromptTokens, model)
updateColorAndText(event.totalTokens, model)
revalidate()
repaint()
}
@ -94,13 +93,12 @@ class TokenUsageCounterPanel(
text = "${percentageLeft.toInt()}% context left"
}
fun updateTooltipText(model: LLModel) {
val usedPrompt = getPromptTokensForSession(currentSessionId) ?: 0
fun updateTooltipText(event: TokenUsageEvent, model: LLModel) {
val budget = computeBudget(model)
toolTipText = buildString {
append("<html><body>")
append("<b>Usage Details</b><br>")
append("Input size: ${numberFormat.format(usedPrompt)} tokens<br>")
append("Input size: ${numberFormat.format(event.totalTokens)} tokens<br>")
append("Max output size: ${numberFormat.format(budget.reservedOutput)} tokens<br>")
append("Max context size: ${numberFormat.format(budget.contextLength)} tokens<br>")
append("</body></html>")
@ -138,13 +136,4 @@ class TokenUsageCounterPanel(
val inputBudget = (contextLength - reserved).coerceAtLeast(1L)
return Budget(contextLength, reserved, inputBudget)
}
private fun getPromptTokensForSession(sessionId: String?): Long? {
if (sessionId == null) return null
return project
?.service<AgentService>()
?.getTokenTrackerForSession(sessionId)
?.getPromptTokens()
}
}