refactor: centralize compression timing logic

Author: Tarquinen

lib/compress/timing.ts

@@ -0,0 +1,129 @@
+import type { SessionState } from "../state/types"
+import { attachCompressionDuration } from "./state"
+
+export interface CompressionStart {
+    sessionId: string
+    messageId: string
+    startedAt: number
+}
+
+export interface PendingCompressionDuration {
+    callId: string
+    messageId: string
+    durationMs: number
+}
+
+export interface CompressionTimingState {
+    startsByCallId: Map<string, CompressionStart>
+    pendingBySessionId: Map<string, PendingCompressionDuration[]>
+}
+
+export function createCompressionTimingState(): CompressionTimingState {
+    return {
+        startsByCallId: new Map(),
+        pendingBySessionId: new Map(),
+    }
+}
+
+export function recordCompressionStart(
+    state: SessionState,
+    callId: string,
+    sessionId: string,
+    messageId: string,
+    startedAt: number,
+): boolean {
+    if (state.compressionTiming.startsByCallId.has(callId)) {
+        return false
+    }
+
+    state.compressionTiming.startsByCallId.set(callId, {
+        sessionId,
+        messageId,
+        startedAt,
+    })
+    return true
+}
+
+export function consumeCompressionStart(
+    state: SessionState,
+    callId: string,
+): CompressionStart | undefined {
+    const start = state.compressionTiming.startsByCallId.get(callId)
+    state.compressionTiming.startsByCallId.delete(callId)
+    return start
+}
+
+export function clearCompressionStart(state: SessionState, callId: string): void {
+    state.compressionTiming.startsByCallId.delete(callId)
+}
+
+export function resolveCompressionDuration(
+    start: CompressionStart | undefined,
+    eventTime: number | undefined,
+    partTime: { start?: unknown; end?: unknown } | undefined,
+): number | undefined {
+    const runningAt =
+        typeof partTime?.start === "number" && Number.isFinite(partTime.start)
+            ? partTime.start
+            : eventTime
+    const pendingToRunningMs =
+        start && typeof runningAt === "number"
+            ? Math.max(0, runningAt - start.startedAt)
+            : undefined
+
+    const toolStart = partTime?.start
+    const toolEnd = partTime?.end
+    const runtimeMs =
+        typeof toolStart === "number" &&
+        Number.isFinite(toolStart) &&
+        typeof toolEnd === "number" &&
+        Number.isFinite(toolEnd)
+            ? Math.max(0, toolEnd - toolStart)
+            : undefined
+
+    return typeof pendingToRunningMs === "number" ? pendingToRunningMs : runtimeMs
+}
+
+export function queueCompressionDuration(
+    state: SessionState,
+    sessionId: string,
+    callId: string,
+    messageId: string,
+    durationMs: number,
+): void {
+    const queued = state.compressionTiming.pendingBySessionId.get(sessionId) || []
+    const filtered = queued.filter((entry) => entry.callId !== callId)
+    filtered.push({ callId, messageId, durationMs })
+    state.compressionTiming.pendingBySessionId.set(sessionId, filtered)
+}
+
+export function applyPendingCompressionDurations(state: SessionState, sessionId: string): number {
+    const queued = state.compressionTiming.pendingBySessionId.get(sessionId)
+    if (!queued || queued.length === 0) {
+        return 0
+    }
+
+    let updates = 0
+    const remaining = []
+    for (const entry of queued) {
+        const applied = attachCompressionDuration(
+            state.prune.messages,
+            entry.callId,
+            entry.messageId,
+            entry.durationMs,
+        )
+        if (applied > 0) {
+            updates += applied
+            continue
+        }
+        remaining.push(entry)
+    }
+
+    if (remaining.length > 0) {
+        state.compressionTiming.pendingBySessionId.set(sessionId, remaining)
+    } else {
+        state.compressionTiming.pendingBySessionId.delete(sessionId)
+    }
+
+    return updates
+}

lib/hooks.ts

@@ -16,6 +16,14 @@ import {
 } from "./messages"
 import { renderSystemPrompt, type PromptStore } from "./prompts"
 import { buildProtectedToolsExtension } from "./prompts/extensions/system"
+import {
+    applyPendingCompressionDurations,
+    clearCompressionStart,
+    consumeCompressionStart,
+    queueCompressionDuration,
+    recordCompressionStart,
+    resolveCompressionDuration,
+} from "./compress/timing"
 import {
     applyPendingManualTrigger,
     handleContextCommand,
@@ -29,14 +37,7 @@ import {
 } from "./commands"
 import { type HostPermissionSnapshot } from "./host-permissions"
 import { compressPermission, syncCompressPermissionState } from "./compress-permission"
-import {
-    checkSession,
-    ensureSessionInitialized,
-    applyPendingCompressionDurations,
-    queueCompressionDuration,
-    saveSessionState,
-    syncToolCache,
-} from "./state"
+import { checkSession, ensureSessionInitialized, saveSessionState, syncToolCache } from "./state"
 import { cacheSystemPromptTokens } from "./ui/utils"
 
 const INTERNAL_AGENT_SIGNATURES = [
@@ -307,16 +308,18 @@ export function createEventHandler(state: SessionState, logger: Logger) {
                 return
             }
 
-            if (state.compressionTiming.startsByCallId.has(part.callID)) {
+            const startedAt = eventTime ?? Date.now()
+            if (
+                !recordCompressionStart(
+                    state,
+                    part.callID,
+                    eventSessionId,
+                    part.messageID,
+                    startedAt,
+                )
+            ) {
                 return
             }
-
-            const startedAt = eventTime ?? Date.now()
-            state.compressionTiming.startsByCallId.set(part.callID, {
-                sessionId: eventSessionId,
-                messageId: part.messageID,
-                startedAt,
-            })
             logger.debug("Recorded compression start", {
                 sessionID: eventSessionId,
                 callID: part.callID,
@@ -335,30 +338,8 @@ export function createEventHandler(state: SessionState, logger: Logger) {
                 return
             }
 
-            const start = state.compressionTiming.startsByCallId.get(part.callID)
-            state.compressionTiming.startsByCallId.delete(part.callID)
-
-            const runningAt =
-                typeof part.state.time?.start === "number" && Number.isFinite(part.state.time.start)
-                    ? part.state.time.start
-                    : eventTime
-            const pendingToRunningMs =
-                start && typeof runningAt === "number"
-                    ? Math.max(0, runningAt - start.startedAt)
-                    : undefined
-
-            const toolStart = part.state.time?.start
-            const toolEnd = part.state.time?.end
-            const runtimeMs =
-                typeof toolStart === "number" &&
-                Number.isFinite(toolStart) &&
-                typeof toolEnd === "number" &&
-                Number.isFinite(toolEnd)
-                    ? Math.max(0, toolEnd - toolStart)
-                    : undefined
-
-            const durationMs =
-                typeof pendingToRunningMs === "number" ? pendingToRunningMs : runtimeMs
+            const start = consumeCompressionStart(state, part.callID)
+            const durationMs = resolveCompressionDuration(start, eventTime, part.state.time)
             if (typeof durationMs !== "number") {
                 return
             }
@@ -390,7 +371,7 @@ export function createEventHandler(state: SessionState, logger: Logger) {
         }
 
         if (typeof part.callID === "string") {
-            state.compressionTiming.startsByCallId.delete(part.callID)
+            clearCompressionStart(state, part.callID)
         }
     }
 }

lib/state/state.ts

@@ -1,6 +1,6 @@
 import type { SessionState, ToolParameterEntry, WithParts } from "./types"
 import type { Logger } from "../logger"
-import { attachCompressionDuration } from "../compress/state"
+import { applyPendingCompressionDurations, createCompressionTimingState } from "../compress/timing"
 import { loadSessionState, saveSessionState } from "./persistence"
 import {
     isSubAgentSession,
@@ -82,10 +82,7 @@ export function createSessionState(): SessionState {
             pruneTokenCounter: 0,
             totalPruneTokens: 0,
         },
-        compressionTiming: {
-            startsByCallId: new Map(),
-            pendingBySessionId: new Map(),
-        },
+        compressionTiming: createCompressionTimingState(),
         toolParameters: new Map<string, ToolParameterEntry>(),
         subAgentResultCache: new Map<string, string>(),
         toolIdList: [],
@@ -188,47 +185,3 @@ export async function ensureSessionInitialized(
         await saveSessionState(state, logger)
     }
 }
-
-export function queueCompressionDuration(
-    state: SessionState,
-    sessionId: string,
-    callId: string,
-    messageId: string,
-    durationMs: number,
-): void {
-    const queued = state.compressionTiming.pendingBySessionId.get(sessionId) || []
-    const filtered = queued.filter((entry) => entry.callId !== callId)
-    filtered.push({ callId, messageId, durationMs })
-    state.compressionTiming.pendingBySessionId.set(sessionId, filtered)
-}
-
-export function applyPendingCompressionDurations(state: SessionState, sessionId: string): number {
-    const queued = state.compressionTiming.pendingBySessionId.get(sessionId)
-    if (!queued || queued.length === 0) {
-        return 0
-    }
-
-    let updates = 0
-    const remaining = []
-    for (const entry of queued) {
-        const applied = attachCompressionDuration(
-            state.prune.messages,
-            entry.callId,
-            entry.messageId,
-            entry.durationMs,
-        )
-        if (applied > 0) {
-            updates += applied
-            continue
-        }
-        remaining.push(entry)
-    }
-
-    if (remaining.length > 0) {
-        state.compressionTiming.pendingBySessionId.set(sessionId, remaining)
-    } else {
-        state.compressionTiming.pendingBySessionId.delete(sessionId)
-    }
-
-    return updates
-}

lib/state/types.ts

@@ -1,3 +1,4 @@
+import type { CompressionTimingState } from "../compress/timing"
 import { Message, Part } from "@opencode-ai/sdk/v2"
 
 export interface WithParts {
@@ -21,23 +22,6 @@ export interface SessionStats {
     totalPruneTokens: number
 }
 
-export interface CompressionStart {
-    sessionId: string
-    messageId: string
-    startedAt: number
-}
-
-export interface PendingCompressionDuration {
-    callId: string
-    messageId: string
-    durationMs: number
-}
-
-export interface CompressionTimingState {
-    startsByCallId: Map<string, CompressionStart>
-    pendingBySessionId: Map<string, PendingCompressionDuration[]>
-}
-
 export interface PrunedMessageEntry {
     tokenCount: number
     allBlockIds: number[]