fix: restore compression timing durability

Author: spoons-and-mirrors

lib/compress/message.ts

@@ -8,6 +8,7 @@ import { appendProtectedTools } from "./protected-content"
 import {
     allocateBlockId,
     consumeCompressionDuration,
+    getCompressionMessageKey,
     allocateRunId,
     applyCompressionState,
     wrapCompressedSummary,
@@ -53,7 +54,6 @@ export function createCompressMessageTool(ctx: ToolContext): ReturnType<typeof t
                 typeof (toolCtx as unknown as { callID?: unknown }).callID === "string"
                     ? (toolCtx as unknown as { callID: string }).callID
                     : undefined
-            const durationMs = consumeCompressionDuration(ctx.state, callId)
 
             const { rawMessages, searchContext } = await prepareSession(
                 ctx,
@@ -96,6 +96,11 @@ 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) {

lib/compress/range.ts

@@ -17,6 +17,7 @@ import {
     COMPRESSED_BLOCK_HEADER,
     allocateBlockId,
     consumeCompressionDuration,
+    getCompressionMessageKey,
     allocateRunId,
     applyCompressionState,
     wrapCompressedSummary,
@@ -64,7 +65,6 @@ export function createCompressRangeTool(ctx: ToolContext): ReturnType<typeof too
                 typeof (toolCtx as unknown as { callID?: unknown }).callID === "string"
                     ? (toolCtx as unknown as { callID: string }).callID
                     : undefined
-            const durationMs = consumeCompressionDuration(ctx.state, callId)
 
             const { rawMessages, searchContext } = await prepareSession(
                 ctx,
@@ -137,6 +137,11 @@ 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) {

lib/compress/state.ts

@@ -4,6 +4,10 @@ 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) {
@@ -29,18 +33,53 @@ 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(state: SessionState, callId?: string): number {
-    if (!callId || !state.compressionDurations.has(callId)) {
+export function consumeCompressionDuration(
+    state: SessionState,
+    messageKey: string,
+    callId?: 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 queue = state.compressionDurationQueue.get(messageKey)
+    const queuedCallId = queue?.shift()
+    if (queue && queue.length === 0) {
+        state.compressionDurationQueue.delete(messageKey)
+    }
+
+    if (!queuedCallId) {
         return 0
     }
 
-    const durationMs = state.compressionDurations.get(callId) || 0
-    state.compressionDurations.delete(callId)
+    const durationMs = state.compressionDurations.get(queuedCallId) || 0
+    state.compressionDurations.delete(queuedCallId)
     return durationMs
 }
 

lib/hooks.ts

@@ -16,7 +16,7 @@ import {
 } from "./messages"
 import { renderSystemPrompt, type PromptStore } from "./prompts"
 import { buildProtectedToolsExtension } from "./prompts/extensions/system"
-import { recordCompressionDuration } from "./compress/state"
+import { getCompressionMessageKey, recordCompressionDuration } from "./compress/state"
 import {
     applyPendingManualTrigger,
     handleContextCommand,
@@ -269,20 +269,12 @@ export function createTextCompleteHandler() {
 
 export function createEventHandler(state: SessionState, logger: Logger) {
     return async (input: { event: any }) => {
-        if (!state.sessionId) {
-            return
-        }
-
         if (input.event.type !== "message.part.updated") {
             return
         }
 
         const part = input.event.properties?.part
-        if (
-            part?.sessionID !== state.sessionId ||
-            part.type !== "tool" ||
-            part.tool !== "compress"
-        ) {
+        if (part?.type !== "tool" || part.tool !== "compress") {
             return
         }
 
@@ -325,7 +317,12 @@ export function createEventHandler(state: SessionState, logger: Logger) {
 
             state.compressionStarts.delete(part.callID)
             const durationMs = Math.max(0, runningAt - start.startedAt)
-            recordCompressionDuration(state, part.callID, durationMs)
+            recordCompressionDuration(
+                state,
+                part.callID,
+                getCompressionMessageKey(part.sessionID, start.messageId),
+                durationMs,
+            )
 
             logger.info("Recorded compression time", {
                 callID: part.callID,

lib/state/state.ts

@@ -83,6 +83,7 @@ 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: [],
@@ -118,8 +119,6 @@ export function resetSessionState(state: SessionState): void {
         pruneTokenCounter: 0,
         totalPruneTokens: 0,
     }
-    state.compressionStarts.clear()
-    state.compressionDurations.clear()
     state.toolParameters.clear()
     state.subAgentResultCache.clear()
     state.toolIdList = []

lib/state/types.ts

@@ -104,6 +104,7 @@ 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[]

tests/compress-message.test.ts

@@ -226,6 +226,55 @@ test("compress message mode batches individual message summaries", async () => {
     assert.match(blocks[1]?.summary || "", /task output body/)
 })
 
+test("compress message mode consumes queued duration without call id", async () => {
+    const sessionID = `ses_message_compress_duration_${Date.now()}`
+    const rawMessages = buildMessages(sessionID)
+    const state = createSessionState()
+    const logger = new Logger(false)
+    state.compressionDurations.set("call-1", 225)
+    state.compressionDurationQueue.set(`${sessionID}:msg-compress-message`, ["call-1"])
+
+    const tool = createCompressMessageTool({
+        client: {
+            session: {
+                messages: async () => ({ data: rawMessages }),
+                get: async () => ({ data: { parentID: null } }),
+            },
+        },
+        state,
+        logger,
+        config: buildConfig(),
+        prompts: {
+            reload() {},
+            getRuntimePrompts() {
+                return { compressMessage: "", compressRange: "" }
+            },
+        },
+    } as any)
+
+    await tool.execute(
+        {
+            topic: "Batch stale notes",
+            content: [
+                {
+                    messageId: "m0002",
+                    topic: "Code path note",
+                    summary: "Captured the assistant's code-path findings.",
+                },
+            ],
+        },
+        {
+            ask: async () => {},
+            metadata: () => {},
+            sessionID,
+            messageID: "msg-compress-message",
+        },
+    )
+
+    const block = Array.from(state.prune.messages.blocksById.values())[0]
+    assert.equal(block?.durationMs, 225)
+})
+
 test("compress message mode does not partially apply when preparation fails", async () => {
     const sessionID = `ses_message_compress_prepare_fail_${Date.now()}`
     const rawMessages = buildMessages(sessionID)

tests/hooks-permission.test.ts

@@ -8,7 +8,7 @@ import {
     createEventHandler,
     createTextCompleteHandler,
 } from "../lib/hooks"
-import { consumeCompressionDuration } from "../lib/compress/state"
+import { consumeCompressionDuration, getCompressionMessageKey } from "../lib/compress/state"
 import { Logger } from "../lib/logger"
 import { createSessionState, type WithParts } from "../lib/state"
 
@@ -207,6 +207,10 @@ test("event hook records compress input generation duration", async () => {
     }
 
     assert.equal(state.compressionDurations.get("call-1"), 225)
+    assert.deepEqual(
+        state.compressionDurationQueue.get(getCompressionMessageKey("session-1", "message-1")),
+        ["call-1"],
+    )
     assert.equal(state.compressionStarts.has("call-1"), false)
 })
 
@@ -301,12 +305,13 @@ test("compression duration consumption handles multiple calls in one message", a
         Date.now = originalNow
     }
 
-    assert.equal(consumeCompressionDuration(state, "call-1"), 225)
-    assert.equal(consumeCompressionDuration(state, "call-2"), 310)
-    assert.equal(consumeCompressionDuration(state, "call-1"), 0)
+    const key = getCompressionMessageKey("session-1", "message-1")
+    assert.equal(consumeCompressionDuration(state, key), 225)
+    assert.equal(consumeCompressionDuration(state, key), 310)
+    assert.equal(consumeCompressionDuration(state, key), 0)
 })
 
-test("compression duration consumption returns zero for unknown call id", async () => {
+test("compression duration consumption supports explicit call id lookup", async () => {
     const state = createSessionState()
     state.sessionId = "session-1"
     const handler = createEventHandler(state, new Logger(false))
@@ -397,10 +402,10 @@ test("compression duration consumption returns zero for unknown call id", async
         Date.now = originalNow
     }
 
-    assert.equal(consumeCompressionDuration(state, "call-2"), 310)
-    assert.equal(consumeCompressionDuration(state, "missing-call"), 0)
-    assert.equal(consumeCompressionDuration(state, undefined), 0)
-    assert.equal(consumeCompressionDuration(state, "call-1"), 225)
+    const key = getCompressionMessageKey("session-1", "message-1")
+    assert.equal(consumeCompressionDuration(state, key, "call-2"), 310)
+    assert.equal(consumeCompressionDuration(state, key), 225)
+    assert.equal(consumeCompressionDuration(state, key), 0)
 })
 
 test("event hook ignores non-compress tool parts", async () => {