fix: attach compression timing to blocks

Author: spoons-and-mirrors

lib/compress/message.ts

@@ -7,8 +7,6 @@ import { finalizeSession, prepareSession, type NotificationEntry } from "./pipel
 import { appendProtectedTools } from "./protected-content"
 import {
     allocateBlockId,
-    consumeCompressionDuration,
-    getCompressionMessageKey,
     allocateRunId,
     applyCompressionState,
     wrapCompressedSummary,
@@ -96,11 +94,6 @@ export function createCompressMessageTool(ctx: ToolContext): ReturnType<typeof t
                 })
             }
 
-            const durationMs = consumeCompressionDuration(
-                ctx.state,
-                getCompressionMessageKey(toolCtx.sessionID, toolCtx.messageID),
-                callId,
-            )
             const runId = allocateRunId(ctx.state)
 
             for (const { plan, summaryWithTools } of preparedPlans) {
@@ -118,8 +111,9 @@ export function createCompressMessageTool(ctx: ToolContext): ReturnType<typeof t
                         mode: "message",
                         runId,
                         compressMessageId: toolCtx.messageID,
+                        compressCallId: callId,
                         summaryTokens,
-                        durationMs,
+                        durationMs: 0,
                     },
                     plan.selection,
                     plan.anchorMessageId,

lib/compress/range.ts

@@ -16,8 +16,6 @@ import {
 import {
     COMPRESSED_BLOCK_HEADER,
     allocateBlockId,
-    consumeCompressionDuration,
-    getCompressionMessageKey,
     allocateRunId,
     applyCompressionState,
     wrapCompressedSummary,
@@ -137,11 +135,6 @@ export function createCompressRangeTool(ctx: ToolContext): ReturnType<typeof too
                 })
             }
 
-            const durationMs = consumeCompressionDuration(
-                ctx.state,
-                getCompressionMessageKey(toolCtx.sessionID, toolCtx.messageID),
-                callId,
-            )
             const runId = allocateRunId(ctx.state)
 
             for (const preparedPlan of preparedPlans) {
@@ -159,8 +152,9 @@ export function createCompressRangeTool(ctx: ToolContext): ReturnType<typeof too
                         mode: "range",
                         runId,
                         compressMessageId: toolCtx.messageID,
+                        compressCallId: callId,
                         summaryTokens,
-                        durationMs,
+                        durationMs: 0,
                     },
                     preparedPlan.selection,
                     preparedPlan.anchorMessageId,

lib/compress/state.ts

@@ -4,10 +4,6 @@ import type { AppliedCompressionResult, CompressionStateInput, SelectionResoluti
 
 export const COMPRESSED_BLOCK_HEADER = "[Compressed conversation section]"
 
-export function getCompressionMessageKey(sessionId: string, messageId: string): string {
-    return `${sessionId}:${messageId}`
-}
-
 export function allocateBlockId(state: SessionState): number {
     const next = state.prune.messages.nextBlockId
     if (!Number.isInteger(next) || next < 1) {
@@ -33,54 +29,36 @@ export function allocateRunId(state: SessionState): number {
 export function recordCompressionDuration(
     state: SessionState,
     callId: string,
-    messageKey: string,
     durationMs: number,
 ): void {
     state.compressionDurations.set(callId, durationMs)
-
-    const queue = state.compressionDurationQueue.get(messageKey)
-    if (!queue) {
-        state.compressionDurationQueue.set(messageKey, [callId])
-        return
-    }
-
-    queue.push(callId)
 }
 
-export function consumeCompressionDuration(
+export function attachCompressionDuration(
     state: SessionState,
-    messageKey: string,
-    callId?: string,
+    callId: string,
+    messageId: string,
 ): number {
-    if (callId && state.compressionDurations.has(callId)) {
-        const queue = state.compressionDurationQueue.get(messageKey)
-        if (queue) {
-            const next = queue.filter((id) => id !== callId)
-            if (next.length === 0) {
-                state.compressionDurationQueue.delete(messageKey)
-            } else {
-                state.compressionDurationQueue.set(messageKey, next)
-            }
-        }
-
-        const durationMs = state.compressionDurations.get(callId) || 0
-        state.compressionDurations.delete(callId)
-        return durationMs
+    const durationMs = state.compressionDurations.get(callId)
+    if (typeof durationMs !== "number" || !Number.isFinite(durationMs)) {
+        return 0
     }
 
-    const queue = state.compressionDurationQueue.get(messageKey)
-    const queuedCallId = queue?.shift()
-    if (queue && queue.length === 0) {
-        state.compressionDurationQueue.delete(messageKey)
-    }
+    state.compressionDurations.delete(callId)
 
-    if (!queuedCallId) {
-        return 0
+    let updates = 0
+    for (const block of state.prune.messages.blocksById.values()) {
+        const matchesCall = block.compressCallId === callId
+        const matchesMessage = !block.compressCallId && block.compressMessageId === messageId
+        if (!matchesCall && !matchesMessage) {
+            continue
+        }
+
+        block.durationMs = durationMs
+        updates++
     }
 
-    const durationMs = state.compressionDurations.get(queuedCallId) || 0
-    state.compressionDurations.delete(queuedCallId)
-    return durationMs
+    return updates
 }
 
 export function wrapCompressedSummary(blockId: number, summary: string): string {
@@ -158,6 +136,7 @@ export function applyCompressionState(
         endId: input.endId,
         anchorMessageId,
         compressMessageId: input.compressMessageId,
+        compressCallId: input.compressCallId,
         includedBlockIds: included,
         consumedBlockIds: consumed,
         parentBlockIds: [],

lib/compress/types.ts

@@ -102,6 +102,7 @@ export interface CompressionStateInput {
     mode: CompressionMode
     runId: number
     compressMessageId: string
+    compressCallId?: string
     summaryTokens: number
     durationMs: number
 }

lib/hooks.ts

@@ -16,7 +16,7 @@ import {
 } from "./messages"
 import { renderSystemPrompt, type PromptStore } from "./prompts"
 import { buildProtectedToolsExtension } from "./prompts/extensions/system"
-import { getCompressionMessageKey, recordCompressionDuration } from "./compress/state"
+import { attachCompressionDuration, recordCompressionDuration } from "./compress/state"
 import {
     applyPendingManualTrigger,
     handleContextCommand,
@@ -30,7 +30,7 @@ import {
 } from "./commands"
 import { type HostPermissionSnapshot } from "./host-permissions"
 import { compressPermission, syncCompressPermissionState } from "./compress-permission"
-import { checkSession, ensureSessionInitialized, syncToolCache } from "./state"
+import { checkSession, ensureSessionInitialized, saveSessionState, syncToolCache } from "./state"
 import { cacheSystemPromptTokens } from "./ui/utils"
 
 const INTERNAL_AGENT_SIGNATURES = [
@@ -317,12 +317,7 @@ export function createEventHandler(state: SessionState, logger: Logger) {
 
             state.compressionStarts.delete(part.callID)
             const durationMs = Math.max(0, runningAt - start.startedAt)
-            recordCompressionDuration(
-                state,
-                part.callID,
-                getCompressionMessageKey(part.sessionID, start.messageId),
-                durationMs,
-            )
+            recordCompressionDuration(state, part.callID, durationMs)
 
             logger.info("Recorded compression time", {
                 callID: part.callID,
@@ -332,8 +327,33 @@ export function createEventHandler(state: SessionState, logger: Logger) {
             return
         }
 
+        if (part.state.status === "completed") {
+            if (typeof part.callID !== "string" || typeof part.messageID !== "string") {
+                return
+            }
+
+            const updates = attachCompressionDuration(state, part.callID, part.messageID)
+            if (updates === 0) {
+                return
+            }
+
+            logger.info("Attached compression time to blocks", {
+                callID: part.callID,
+                messageID: part.messageID,
+                blocks: updates,
+            })
+
+            saveSessionState(state, logger).catch((error) => {
+                logger.warn("Failed to persist compression time update", {
+                    error: error instanceof Error ? error.message : String(error),
+                })
+            })
+            return
+        }
+
         if (typeof part.callID === "string") {
             state.compressionStarts.delete(part.callID)
+            state.compressionDurations.delete(part.callID)
         }
     }
 }

lib/state/state.ts

@@ -83,7 +83,6 @@ export function createSessionState(): SessionState {
         },
         compressionStarts: new Map<string, CompressionStart>(),
         compressionDurations: new Map<string, number>(),
-        compressionDurationQueue: new Map<string, string[]>(),
         toolParameters: new Map<string, ToolParameterEntry>(),
         subAgentResultCache: new Map<string, string>(),
         toolIdList: [],

lib/state/types.ts

@@ -49,6 +49,7 @@ export interface CompressionBlock {
     endId: string
     anchorMessageId: string
     compressMessageId: string
+    compressCallId?: string
     includedBlockIds: number[]
     consumedBlockIds: number[]
     parentBlockIds: number[]
@@ -104,7 +105,6 @@ export interface SessionState {
     stats: SessionStats
     compressionStarts: Map<string, CompressionStart>
     compressionDurations: Map<string, number>
-    compressionDurationQueue: Map<string, string[]>
     toolParameters: Map<string, ToolParameterEntry>
     subAgentResultCache: Map<string, string>
     toolIdList: string[]

lib/state/utils.ts

@@ -196,6 +196,8 @@ export function loadPruneMessagesState(
                     typeof block.anchorMessageId === "string" ? block.anchorMessageId : "",
                 compressMessageId:
                     typeof block.compressMessageId === "string" ? block.compressMessageId : "",
+                compressCallId:
+                    typeof block.compressCallId === "string" ? block.compressCallId : undefined,
                 includedBlockIds: toNumberArray(block.includedBlockIds),
                 consumedBlockIds: toNumberArray(block.consumedBlockIds),
                 parentBlockIds: toNumberArray(block.parentBlockIds),