refactor: move compression duration updates into compress state

Author: Tarquinen

lib/compress/state.ts

@@ -1,4 +1,4 @@
-import type { CompressionBlock, SessionState } from "../state"
+import type { CompressionBlock, PruneMessagesState, SessionState } from "../state"
 import { formatBlockRef, formatMessageIdTag } from "../message-ids"
 import type { AppliedCompressionResult, CompressionStateInput, SelectionResolution } from "./types"
 
@@ -26,6 +26,31 @@ export function allocateRunId(state: SessionState): number {
     return next
 }
 
+export function attachCompressionDuration(
+    messagesState: PruneMessagesState,
+    callId: string,
+    messageId: string,
+    durationMs: number,
+): number {
+    if (typeof durationMs !== "number" || !Number.isFinite(durationMs)) {
+        return 0
+    }
+
+    let updates = 0
+    for (const block of messagesState.blocksById.values()) {
+        const matchesCall = block.compressCallId === callId
+        const matchesMessage = !block.compressCallId && block.compressMessageId === messageId
+        if (!matchesCall && !matchesMessage) {
+            continue
+        }
+
+        block.durationMs = durationMs
+        updates++
+    }
+
+    return updates
+}
+
 export function wrapCompressedSummary(blockId: number, summary: string): string {
     const header = COMPRESSED_BLOCK_HEADER
     const footer = formatMessageIdTag(formatBlockRef(blockId))

lib/state/state.ts

@@ -1,8 +1,8 @@
 import type { CompressionTimingState, SessionState, ToolParameterEntry, WithParts } from "./types"
 import type { Logger } from "../logger"
+import { attachCompressionDuration } from "../compress/state"
 import { loadSessionState, saveSessionState } from "./persistence"
 import {
-    attachCompressionDuration,
     isSubAgentSession,
     findLastCompactionTimestamp,
     countTurns,

lib/state/utils.ts

@@ -8,31 +8,6 @@ import type {
 import { isIgnoredUserMessage, messageHasCompress } from "../messages/query"
 import { countTokens } from "../token-utils"
 
-export function attachCompressionDuration(
-    messagesState: PruneMessagesState,
-    callId: string,
-    messageId: string,
-    durationMs: number,
-): number {
-    if (typeof durationMs !== "number" || !Number.isFinite(durationMs)) {
-        return 0
-    }
-
-    let updates = 0
-    for (const block of messagesState.blocksById.values()) {
-        const matchesCall = block.compressCallId === callId
-        const matchesMessage = !block.compressCallId && block.compressMessageId === messageId
-        if (!matchesCall && !matchesMessage) {
-            continue
-        }
-
-        block.durationMs = durationMs
-        updates++
-    }
-
-    return updates
-}
-
 export const isMessageCompacted = (state: SessionState, msg: WithParts): boolean => {
     if (msg.info.time.created < state.lastCompaction) {
         return true