refactor: simplify compression timing state

Author: Tarquinen

lib/compress/pipeline.ts

@@ -9,6 +9,7 @@ import { sendCompressNotification } from "../ui/notification"
 import type { ToolContext } from "./types"
 import { buildSearchContext, fetchSessionMessages } from "./search"
 import type { SearchContext } from "./types"
+import { applyPendingCompressionDurations } from "./timing"
 
 interface RunContext {
     ask(input: {
@@ -83,6 +84,7 @@ export async function finalizeSession(
     batchTopic: string | undefined,
 ): Promise<void> {
     ctx.state.manualMode = ctx.state.manualMode ? "active" : false
+    applyPendingCompressionDurations(ctx.state)
     await saveSessionState(ctx.state, ctx.logger)
 
     const params = getCurrentParams(ctx.state, rawMessages, ctx.logger)

lib/compress/state.ts

@@ -29,7 +29,6 @@ export function allocateRunId(state: SessionState): number {
 export function attachCompressionDuration(
     messagesState: PruneMessagesState,
     callId: string,
-    messageId: string,
     durationMs: number,
 ): number {
     if (typeof durationMs !== "number" || !Number.isFinite(durationMs)) {
@@ -38,9 +37,7 @@ export function attachCompressionDuration(
 
     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) {
+        if (block.compressCallId !== callId) {
             continue
         }
 

lib/compress/timing.ts

@@ -1,53 +1,37 @@
 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[]>
+    startsByCallId: Map<string, number>
+    pendingByCallId: Map<string, PendingCompressionDuration>
 }
 
 export function createCompressionTimingState(): CompressionTimingState {
     return {
         startsByCallId: new Map(),
-        pendingBySessionId: new Map(),
+        pendingByCallId: 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,
-    })
+    state.compressionTiming.startsByCallId.set(callId, startedAt)
     return true
 }
 
-export function consumeCompressionStart(
-    state: SessionState,
-    callId: string,
-): CompressionStart | undefined {
+export function consumeCompressionStart(state: SessionState, callId: string): number | undefined {
     const start = state.compressionTiming.startsByCallId.get(callId)
     state.compressionTiming.startsByCallId.delete(callId)
     return start
@@ -58,7 +42,7 @@ export function clearCompressionStart(state: SessionState, callId: string): void
 }
 
 export function resolveCompressionDuration(
-    start: CompressionStart | undefined,
+    startedAt: number | undefined,
     eventTime: number | undefined,
     partTime: { start?: unknown; end?: unknown } | undefined,
 ): number | undefined {
@@ -67,8 +51,8 @@ export function resolveCompressionDuration(
             ? partTime.start
             : eventTime
     const pendingToRunningMs =
-        start && typeof runningAt === "number"
-            ? Math.max(0, runningAt - start.startedAt)
+        typeof startedAt === "number" && typeof runningAt === "number"
+            ? Math.max(0, runningAt - startedAt)
             : undefined
 
     const toolStart = partTime?.start
@@ -86,43 +70,28 @@ export function resolveCompressionDuration(
 
 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)
+    state.compressionTiming.pendingByCallId.set(callId, { callId, durationMs })
 }
 
-export function applyPendingCompressionDurations(state: SessionState, sessionId: string): number {
-    const queued = state.compressionTiming.pendingBySessionId.get(sessionId)
-    if (!queued || queued.length === 0) {
+export function applyPendingCompressionDurations(state: SessionState): number {
+    if (state.compressionTiming.pendingByCallId.size === 0) {
         return 0
     }
 
     let updates = 0
-    const remaining = []
-    for (const entry of queued) {
+    for (const [callId, entry] of state.compressionTiming.pendingByCallId) {
         const applied = attachCompressionDuration(
             state.prune.messages,
             entry.callId,
-            entry.messageId,
             entry.durationMs,
         )
         if (applied > 0) {
             updates += applied
-            continue
+            state.compressionTiming.pendingByCallId.delete(callId)
         }
-        remaining.push(entry)
-    }
-
-    if (remaining.length > 0) {
-        state.compressionTiming.pendingBySessionId.set(sessionId, remaining)
-    } else {
-        state.compressionTiming.pendingBySessionId.delete(sessionId)
     }
 
     return updates

lib/hooks.ts

@@ -276,12 +276,6 @@ export function createTextCompleteHandler() {
 
 export function createEventHandler(state: SessionState, logger: Logger) {
     return async (input: { event: any }) => {
-        const eventSessionId =
-            typeof input.event?.properties?.sessionID === "string"
-                ? input.event.properties.sessionID
-                : typeof input.event?.properties?.part?.sessionID === "string"
-                  ? input.event.properties.part.sessionID
-                  : undefined
         const eventTime =
             typeof input.event?.time === "number" && Number.isFinite(input.event.time)
                 ? input.event.time
@@ -300,41 +294,23 @@ export function createEventHandler(state: SessionState, logger: Logger) {
         }
 
         if (part.state.status === "pending") {
-            if (
-                typeof part.callID !== "string" ||
-                typeof part.messageID !== "string" ||
-                typeof eventSessionId !== "string"
-            ) {
+            if (typeof part.callID !== "string") {
                 return
             }
 
             const startedAt = eventTime ?? Date.now()
-            if (
-                !recordCompressionStart(
-                    state,
-                    part.callID,
-                    eventSessionId,
-                    part.messageID,
-                    startedAt,
-                )
-            ) {
+            if (!recordCompressionStart(state, part.callID, startedAt)) {
                 return
             }
             logger.debug("Recorded compression start", {
-                sessionID: eventSessionId,
                 callID: part.callID,
-                messageID: part.messageID,
                 startedAt,
             })
             return
         }
 
         if (part.state.status === "completed") {
-            if (
-                typeof part.callID !== "string" ||
-                typeof part.messageID !== "string" ||
-                typeof eventSessionId !== "string"
-            ) {
+            if (typeof part.callID !== "string") {
                 return
             }
 
@@ -344,22 +320,17 @@ export function createEventHandler(state: SessionState, logger: Logger) {
                 return
             }
 
-            queueCompressionDuration(state, eventSessionId, part.callID, part.messageID, durationMs)
+            queueCompressionDuration(state, part.callID, durationMs)
 
-            const updates =
-                state.sessionId === eventSessionId
-                    ? applyPendingCompressionDurations(state, eventSessionId)
-                    : 0
+            const updates = applyPendingCompressionDurations(state)
             if (updates === 0) {
                 return
             }
 
             await saveSessionState(state, logger)
 
             logger.info("Attached compression time to blocks", {
-                sessionID: eventSessionId,
                 callID: part.callID,
-                messageID: part.messageID,
                 blocks: updates,
                 durationMs,
             })

lib/state/state.ts

@@ -180,7 +180,7 @@ export async function ensureSessionInitialized(
         totalPruneTokens: persisted.stats?.totalPruneTokens || 0,
     }
 
-    const applied = applyPendingCompressionDurations(state, sessionId)
+    const applied = applyPendingCompressionDurations(state)
     if (applied > 0) {
         await saveSessionState(state, logger)
     }

tests/hooks-permission.test.ts

@@ -372,7 +372,7 @@ test("event hook falls back to completed runtime when running duration missing",
         endId: "m0001",
         anchorMessageId: "msg-a",
         compressMessageId: "message-1",
-        compressCallId: undefined,
+        compressCallId: "call-3",
         includedBlockIds: [],
         consumedBlockIds: [],
         parentBlockIds: [],
@@ -492,7 +492,7 @@ test("event hook queues duration updates until the matching session is loaded",
         },
     })
 
-    assert.equal(liveState.compressionTiming.pendingBySessionId.get(targetSessionId)?.length, 1)
+    assert.equal(liveState.compressionTiming.pendingByCallId.has("call-remote"), true)
     assert.equal(liveState.compressionTiming.startsByCallId.has("call-remote"), false)
 
     await ensureSessionInitialized(
@@ -520,5 +520,5 @@ test("event hook queues duration updates until the matching session is loaded",
     )
 
     assert.equal(liveState.prune.messages.blocksById.get(1)?.durationMs, 250)
-    assert.equal(liveState.compressionTiming.pendingBySessionId.has(targetSessionId), false)
+    assert.equal(liveState.compressionTiming.pendingByCallId.has("call-remote"), false)
 })