refactor: derive compress graph from token buckets

Author: Tarquinen

lib/hooks.ts

@@ -13,7 +13,7 @@ import { handleHelpCommand } from "./commands/help"
 import { handleSweepCommand } from "./commands/sweep"
 import { handleManualToggleCommand, handleManualTriggerCommand } from "./commands/manual"
 import { ensureSessionInitialized } from "./state/state"
-import { getCurrentParams } from "./strategies/utils"
+import { cacheSystemPromptTokens } from "./ui/utils"
 
 const INTERNAL_AGENT_SIGNATURES = [
     "You are a title generator",
@@ -107,6 +107,8 @@ export function createChatMessageTransformHandler(
             return
         }
 
+        cacheSystemPromptTokens(state, output.messages)
+
         syncToolCache(state, config, logger, output.messages)
         buildToolIdList(state, output.messages)
 

lib/state/state.ts

@@ -82,6 +82,7 @@ export function createSessionState(): SessionState {
         currentTurn: 0,
         variant: undefined,
         modelContextLimit: undefined,
+        systemPromptTokens: undefined,
     }
 }
 
@@ -108,6 +109,7 @@ export function resetSessionState(state: SessionState): void {
     state.currentTurn = 0
     state.variant = undefined
     state.modelContextLimit = undefined
+    state.systemPromptTokens = undefined
 }
 
 export async function ensureSessionInitialized(

lib/state/types.ts

@@ -53,4 +53,5 @@ export interface SessionState {
     currentTurn: number
     variant: string | undefined
     modelContextLimit: number | undefined
+    systemPromptTokens: number | undefined
 }

lib/tools/compress.ts

@@ -8,6 +8,7 @@ import { getCurrentParams, countAllMessageTokens, countTokens } from "../strateg
 import type { AssistantMessage } from "@opencode-ai/sdk/v2"
 import { findStringInMessages, collectToolIdsInRange, collectMessageIdsInRange } from "./utils"
 import { sendCompressNotification } from "../ui/notification"
+import { cacheSystemPromptTokens } from "../ui/utils"
 import { prune as applyPruneTransforms } from "../messages/prune"
 import { clog, C } from "../compress-logger"
 
@@ -109,6 +110,8 @@ export function createCompressTool(ctx: ToolContext): ReturnType<typeof tool> {
                     ctx.config.manualMode.enabled,
                 )
 
+                cacheSystemPromptTokens(state, messages)
+
                 clog.info(C.STATE, `State Snapshot (before boundary matching)`, {
                     sessionId: state.sessionId,
                     isSubAgent: state.isSubAgent,
@@ -406,20 +409,6 @@ export function createCompressTool(ctx: ToolContext): ReturnType<typeof tool> {
                     compressedToolIds: compressedToolIds.length,
                 })
 
-                // Build token weight map for all messages (for proportional bar graph)
-                const weights = new Map<string, number>()
-                let contentWeight = 0
-                for (const msg of messages) {
-                    const w = countAllMessageTokens(msg)
-                    weights.set(msg.info.id, w)
-                    contentWeight += w
-                }
-
-                // System prompt = API total minus message content
-                const systemWeight = Math.max(0, totalSessionTokens - contentWeight)
-                const sessionIds = ["__system__", ...messages.map((m) => m.info.id)]
-                weights.set("__system__", systemWeight)
-
                 await sendCompressNotification(
                     client,
                     logger,
@@ -431,11 +420,7 @@ export function createCompressTool(ctx: ToolContext): ReturnType<typeof tool> {
                     topic,
                     summary,
                     summaryTokens,
-                    totalSessionTokens,
-                    estimatedCompressedTokens,
-                    sessionIds,
-                    weights,
-                    messages.length,
+                    messages,
                     currentParams,
                 )
 

lib/ui/notification.ts

@@ -1,12 +1,13 @@
 import type { Logger } from "../logger"
 import type { SessionState } from "../state"
 import {
+    buildCompressionGraphData,
+    formatCompressionGraph,
     formatPrunedItemsList,
     formatStatsHeader,
     formatTokenCount,
-    formatSessionMap,
 } from "./utils"
-import { ToolParameterEntry } from "../state"
+import { ToolParameterEntry, WithParts } from "../state"
 import { PluginConfig } from "../config"
 import { clog, C } from "../compress-logger"
 
@@ -139,11 +140,7 @@ export async function sendCompressNotification(
     topic: string,
     summary: string,
     summaryTokens: number,
-    totalSessionTokens: number,
-    compressedTokens: number,
-    sessionMessageIds: string[],
-    weights: Map<string, number>,
-    totalMessages: number,
+    messages: WithParts[],
     params: any,
 ): Promise<boolean> {
     if (config.pruneNotification === "off") {
@@ -158,26 +155,24 @@ export async function sendCompressNotification(
     } else {
         message = formatStatsHeader(state.stats.totalPruneTokens, state.stats.pruneTokenCounter)
 
-        const pruneTokenCounterStr = `~${formatTokenCount(compressedTokens)}`
+        const newIds = new Set(messageIds)
+        const newToolIds = new Set(toolIds)
+        const graphData = buildCompressionGraphData(state, messages, newIds, newToolIds)
+        const progressBar = formatCompressionGraph(graphData, 50)
+        const pruneTokenCounterStr = `~${formatTokenCount(graphData.recentCompressedTokens)}`
+        const reduction =
+            graphData.totalSessionTokens > 0
+                ? Math.round(
+                      (graphData.recentCompressedTokens / graphData.totalSessionTokens) * 100,
+                  )
+                : 0
 
-        clog.info(C.COMPRESS, `sendCompressNotification inputs`, {
-            totalSessionTokens,
-            compressedTokens,
+        clog.info(C.COMPRESS, `sendCompressNotification graph`, {
             summaryTokens,
-            ratio:
-                totalSessionTokens > 0 ? (compressedTokens / totalSessionTokens).toFixed(4) : "N/A",
+            reductionPercent: reduction,
+            ...graphData,
         })
 
-        const newIds = new Set(messageIds)
-        const progressBar = formatSessionMap(
-            sessionMessageIds,
-            state.prune.messages,
-            newIds,
-            weights,
-            50,
-        )
-        const reduction =
-            totalSessionTokens > 0 ? Math.round((compressedTokens / totalSessionTokens) * 100) : 0
         message += `\n\n${progressBar}`
         message += `\n▣ Compressing (${pruneTokenCounterStr} removed, ${reduction}% reduction)`
         message += `\n→ Topic: ${topic}`