fix: tooltip agent total tokens and other minor adjustments

src/main/kotlin/ee/carlrobert/codegpt/agent/AgentService.kt

@@ -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() }
     }
-}
+}

src/main/kotlin/ee/carlrobert/codegpt/agent/strategy/SingleRunStrategyProvider.kt

@@ -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)

src/main/kotlin/ee/carlrobert/codegpt/agent/tools/IntelliJSearchTool.kt

@@ -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,

src/main/kotlin/ee/carlrobert/codegpt/toolwindow/agent/ui/RollbackPanel.kt

@@ -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
             }
         }

src/main/kotlin/ee/carlrobert/codegpt/ui/components/LeftEllipsisLabel.kt

@@ -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 ""
+    }
+}

src/main/kotlin/ee/carlrobert/codegpt/ui/components/TokenUsageCounterPanel.kt

@@ -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()
-    }
-
 }