feat(llm): add token usage service with session tracking

Create TokenUsageService to centrally manage token statistics across sessions with availability calculations, and refactor TokenUsagePanel to use cleaner UI layout and improved progress visualization.

Author: phodal

core/src/main/kotlin/cc/unitmesh/devti/gui/chat/ui/TokenUsagePanel.kt

@@ -13,7 +13,6 @@ import com.intellij.util.ui.JBUI
 import com.intellij.util.ui.UIUtil
 import com.intellij.util.ui.components.BorderLayoutPanel
 import java.awt.BorderLayout
-import java.awt.Color
 import java.awt.Font
 import java.awt.GridBagConstraints
 import java.awt.GridBagLayout
@@ -25,10 +24,6 @@ import javax.swing.SwingConstants
  * Panel that displays token usage statistics for the current session
  */
 class TokenUsagePanel(private val project: Project) : BorderLayoutPanel() {
-    
-    private val promptTokensLabel = JBLabel("0", SwingConstants.RIGHT)
-    private val completionTokensLabel = JBLabel("0", SwingConstants.RIGHT)
-    private val totalTokensLabel = JBLabel("0", SwingConstants.RIGHT)
     private val modelLabel = JBLabel("", SwingConstants.LEFT)
     private val progressBar = JProgressBar(0, 100)
     private val usageRatioLabel = JBLabel("", SwingConstants.CENTER)
@@ -46,7 +41,6 @@ class TokenUsagePanel(private val project: Project) : BorderLayoutPanel() {
         isOpaque = false
         border = JBUI.Borders.empty(4, 8)
 
-        // Setup progress bar
         progressBar.apply {
             isStringPainted = false
             preferredSize = java.awt.Dimension(150, 8)
@@ -55,19 +49,16 @@ class TokenUsagePanel(private val project: Project) : BorderLayoutPanel() {
             isOpaque = false
         }
 
-        // Setup usage ratio label
         usageRatioLabel.apply {
             font = font.deriveFont(Font.PLAIN, 10f)
             foreground = UIUtil.getContextHelpForeground()
         }
 
-        // Create main layout
         val mainPanel = JPanel(GridBagLayout())
         mainPanel.isOpaque = false
 
         val gbc = GridBagConstraints()
 
-        // Top row: Model info and progress bar
         gbc.gridx = 0
         gbc.gridy = 0
         gbc.anchor = GridBagConstraints.WEST
@@ -129,18 +120,13 @@ class TokenUsagePanel(private val project: Project) : BorderLayoutPanel() {
             currentUsage = event.usage
             currentModel = event.model
 
-            // Get max tokens for current model
             updateMaxTokens()
-            
-            
-            // Update progress bar
             updateProgressBar(event.usage.totalTokens ?: 0)
-            
+
             if (!event.model.isNullOrBlank()) {
                 modelLabel.text = "Model: ${event.model}"
             }
 
-            // Show the panel when we have data
             isVisible = true
             revalidate()
             repaint()
@@ -154,7 +140,6 @@ class TokenUsagePanel(private val project: Project) : BorderLayoutPanel() {
             val limits = modelManager.getUsedMaxToken()
             maxContextWindowTokens = limits.maxContextWindowTokens?.toLong() ?: 0
         } catch (e: Exception) {
-            // Fallback to default if unable to get limits
             maxContextWindowTokens = 4096
         }
     }
@@ -169,7 +154,6 @@ class TokenUsagePanel(private val project: Project) : BorderLayoutPanel() {
         val usageRatio = (totalTokens.toDouble() / maxContextWindowTokens * 100).toInt()
         progressBar.value = usageRatio.coerceIn(0, 100)
 
-        // Update color based on usage ratio
         progressBar.foreground = when {
             usageRatio >= 90 -> JBColor.RED
             usageRatio >= 75 -> JBColor.ORANGE
@@ -178,12 +162,9 @@ class TokenUsagePanel(private val project: Project) : BorderLayoutPanel() {
             else -> UIUtil.getPanelBackground().brighter()
         }
 
-        // Update ratio label
         usageRatioLabel.text = "${formatTokenCount(totalTokens)}/${formatTokenCount(maxContextWindowTokens)} (${usageRatio}%)"
-        
         progressBar.isVisible = true
         usageRatioLabel.isVisible = true
-        
         progressBar.toolTipText = "Token usage: $usageRatio% of context window"
     }
 
@@ -194,10 +175,7 @@ class TokenUsagePanel(private val project: Project) : BorderLayoutPanel() {
             else -> count.toString()
         }
     }
-    
-    /**
-     * Reset the token usage display
-     */
+
     fun reset() {
         ApplicationManager.getApplication().invokeLater {
             currentUsage = Usage()
@@ -213,14 +191,8 @@ class TokenUsagePanel(private val project: Project) : BorderLayoutPanel() {
             repaint()
         }
     }
-    
-    /**
-     * Get current usage for external access
-     */
+
     fun getCurrentUsage(): Usage = currentUsage
-    
-    /**
-     * Get current model for external access
-     */
+
     fun getCurrentModel(): String? = currentModel
 }

core/src/main/kotlin/cc/unitmesh/devti/llm2/TokenUsageService.kt

@@ -0,0 +1,131 @@
+package cc.unitmesh.devti.llm2
+
+import cc.unitmesh.devti.llms.custom.Usage
+import cc.unitmesh.devti.settings.model.LLMModelManager
+import com.intellij.openapi.application.ApplicationManager
+import com.intellij.openapi.components.Service
+import com.intellij.openapi.project.Project
+import com.intellij.util.messages.MessageBusConnection
+import java.util.concurrent.ConcurrentHashMap
+import java.util.concurrent.atomic.AtomicLong
+
+/**
+ * Service for managing token usage statistics and calculations
+ */
+@Service(Service.Level.PROJECT)
+class TokenUsageService(private val project: Project) : TokenUsageListener {
+    private val sessionTokenUsage = ConcurrentHashMap<String, Usage>()
+    private val totalTokenUsage = AtomicLong(0)
+    private val totalPromptTokens = AtomicLong(0)
+    private val totalCompletionTokens = AtomicLong(0)
+    
+    private var currentSessionId: String? = null
+    private var currentModel: String? = null
+    private var messageBusConnection: MessageBusConnection? = null
+    
+    init {
+        setupTokenUsageListener()
+    }
+    
+    private fun setupTokenUsageListener() {
+        messageBusConnection = ApplicationManager.getApplication().messageBus.connect()
+        messageBusConnection?.subscribe(TokenUsageListener.TOPIC, this)
+    }
+    
+    override fun onTokenUsage(event: TokenUsageEvent) {
+        event.sessionId?.let { sessionId ->
+            sessionTokenUsage[sessionId] = event.usage
+            currentSessionId = sessionId
+        }
+        
+        // Update current model
+        currentModel = event.model
+        
+        // Update total usage statistics
+        totalTokenUsage.addAndGet(event.usage.totalTokens)
+        totalPromptTokens.addAndGet(event.usage.promptTokens)
+        totalCompletionTokens.addAndGet(event.usage.completionTokens)
+    }
+    
+    /**
+     * Get the current consumed tokens for the active session
+     */
+    fun getCurrentConsumedTokens(): Usage {
+        return currentSessionId?.let { sessionId ->
+            sessionTokenUsage[sessionId]
+        } ?: Usage()
+    }
+    
+    /**
+     * Get total consumed tokens across all sessions
+     */
+    fun getTotalConsumedTokens(): Usage {
+        return Usage(
+            promptTokens = totalPromptTokens.get(),
+            completionTokens = totalCompletionTokens.get(),
+            totalTokens = totalTokenUsage.get()
+        )
+    }
+    
+    /**
+     * Get the maximum tokens allowed for the current model
+     */
+    fun getUsedMaxToken(): Long {
+        return LLMModelManager.getInstance().getUsedMaxToken().maxContextWindowTokens?.toLong() ?: 0L
+    }
+    
+    /**
+     * Calculate token availability ratio (0.0 to 1.0)
+     * Returns the percentage of tokens still available
+     */
+    fun calculateTokenAvailability(): Double {
+        val maxTokens = getUsedMaxToken()
+        if (maxTokens <= 0) return 1.0
+        
+        val currentUsage = getCurrentConsumedTokens()
+        val usedTokens = currentUsage.totalTokens
+        
+        return if (usedTokens >= maxTokens) {
+            0.0
+        } else {
+            1.0 - (usedTokens.toDouble() / maxTokens.toDouble())
+        }
+    }
+    
+    /**
+     * Calculate token usage ratio (0.0 to 1.0)
+     * Returns the percentage of tokens already used
+     */
+    fun calculateTokenUsageRatio(): Double {
+        return 1.0 - calculateTokenAvailability()
+    }
+    
+    /**
+     * Check if token usage is approaching the limit
+     * @param threshold Warning threshold (default 0.8 = 80%)
+     */
+    fun isApproachingTokenLimit(threshold: Double = 0.8): Boolean {
+        return calculateTokenUsageRatio() >= threshold
+    }
+    
+    /**
+     * Get remaining tokens for current session
+     */
+    fun getRemainingTokens(): Long {
+        val maxTokens = getUsedMaxToken()
+        val currentUsage = getCurrentConsumedTokens()
+        val remaining = maxTokens - currentUsage.totalTokens
+        return maxOf(0L, remaining)
+    }
+
+    fun dispose() {
+        messageBusConnection?.disconnect()
+        messageBusConnection = null
+    }
+    
+    companion object {
+        fun getInstance(project: Project): TokenUsageService {
+            return project.getService(TokenUsageService::class.java)
+        }
+    }
+}