fix: stabilize compression timing correlation

Author: spoons-and-mirrors

lib/compress/message.ts

@@ -7,6 +7,7 @@ import { finalizeSession, prepareSession, type NotificationEntry } from "./pipel
 import { appendProtectedTools } from "./protected-content"
 import {
     allocateBlockId,
+    consumeCompressionDuration,
     allocateRunId,
     applyCompressionState,
     wrapCompressedSummary,
@@ -46,92 +47,93 @@ export function createCompressMessageTool(ctx: ToolContext): ReturnType<typeof t
         description: runtimePrompts.compressMessage + MESSAGE_FORMAT_EXTENSION,
         args: buildSchema(),
         async execute(args, toolCtx) {
-            try {
-                const input = args as CompressMessageToolArgs
-                validateArgs(input)
-                const { rawMessages, searchContext } = await prepareSession(
-                    ctx,
-                    toolCtx,
-                    `Compress Message: ${input.topic}`,
-                )
-                const { plans, skippedIssues } = resolveMessages(
-                    input,
-                    searchContext,
-                    ctx.state,
-                    ctx.config,
-                )
+            const input = args as CompressMessageToolArgs
+            validateArgs(input)
+            const callId =
+                typeof (toolCtx as unknown as { callID?: unknown }).callID === "string"
+                    ? (toolCtx as unknown as { callID: string }).callID
+                    : undefined
+            const durationMs = consumeCompressionDuration(ctx.state, toolCtx.messageID, callId)
 
-                if (plans.length === 0 && skippedIssues.length > 0) {
-                    throw new Error(formatIssues(skippedIssues))
-                }
+            const { rawMessages, searchContext } = await prepareSession(
+                ctx,
+                toolCtx,
+                `Compress Message: ${input.topic}`,
+            )
+            const { plans, skippedIssues } = resolveMessages(
+                input,
+                searchContext,
+                ctx.state,
+                ctx.config,
+            )
 
-                const notifications: NotificationEntry[] = []
+            if (plans.length === 0 && skippedIssues.length > 0) {
+                throw new Error(formatIssues(skippedIssues))
+            }
 
-                const preparedPlans: Array<{
-                    plan: (typeof plans)[number]
-                    summaryWithTools: string
-                }> = []
+            const notifications: NotificationEntry[] = []
 
-                for (const plan of plans) {
-                    const summaryWithTools = await appendProtectedTools(
-                        ctx.client,
-                        ctx.state,
-                        ctx.config.experimental.allowSubAgents,
-                        plan.entry.summary,
-                        plan.selection,
-                        searchContext,
-                        ctx.config.compress.protectedTools,
-                        ctx.config.protectedFilePatterns,
-                    )
+            const preparedPlans: Array<{
+                plan: (typeof plans)[number]
+                summaryWithTools: string
+            }> = []
 
-                    preparedPlans.push({
-                        plan,
-                        summaryWithTools,
-                    })
-                }
+            for (const plan of plans) {
+                const summaryWithTools = await appendProtectedTools(
+                    ctx.client,
+                    ctx.state,
+                    ctx.config.experimental.allowSubAgents,
+                    plan.entry.summary,
+                    plan.selection,
+                    searchContext,
+                    ctx.config.compress.protectedTools,
+                    ctx.config.protectedFilePatterns,
+                )
 
-                const runId = allocateRunId(ctx.state)
-                const durationMs = ctx.state.compressionDurations.get(toolCtx.messageID) || 0
+                preparedPlans.push({
+                    plan,
+                    summaryWithTools,
+                })
+            }
 
-                for (const { plan, summaryWithTools } of preparedPlans) {
-                    const blockId = allocateBlockId(ctx.state)
-                    const storedSummary = wrapCompressedSummary(blockId, summaryWithTools)
-                    const summaryTokens = countTokens(storedSummary)
+            const runId = allocateRunId(ctx.state)
 
-                    applyCompressionState(
-                        ctx.state,
-                        {
-                            topic: plan.entry.topic,
-                            batchTopic: input.topic,
-                            startId: plan.entry.messageId,
-                            endId: plan.entry.messageId,
-                            mode: "message",
-                            runId,
-                            compressMessageId: toolCtx.messageID,
-                            summaryTokens,
-                            durationMs,
-                        },
-                        plan.selection,
-                        plan.anchorMessageId,
-                        blockId,
-                        storedSummary,
-                        [],
-                    )
+            for (const { plan, summaryWithTools } of preparedPlans) {
+                const blockId = allocateBlockId(ctx.state)
+                const storedSummary = wrapCompressedSummary(blockId, summaryWithTools)
+                const summaryTokens = countTokens(storedSummary)
 
-                    notifications.push({
-                        blockId,
+                applyCompressionState(
+                    ctx.state,
+                    {
+                        topic: plan.entry.topic,
+                        batchTopic: input.topic,
+                        startId: plan.entry.messageId,
+                        endId: plan.entry.messageId,
+                        mode: "message",
                         runId,
-                        summary: summaryWithTools,
+                        compressMessageId: toolCtx.messageID,
                         summaryTokens,
-                    })
-                }
-
-                await finalizeSession(ctx, toolCtx, rawMessages, notifications, input.topic)
+                        durationMs,
+                    },
+                    plan.selection,
+                    plan.anchorMessageId,
+                    blockId,
+                    storedSummary,
+                    [],
+                )
 
-                return formatResult(plans.length, skippedIssues)
-            } finally {
-                ctx.state.compressionDurations.delete(toolCtx.messageID)
+                notifications.push({
+                    blockId,
+                    runId,
+                    summary: summaryWithTools,
+                    summaryTokens,
+                })
             }
+
+            await finalizeSession(ctx, toolCtx, rawMessages, notifications, input.topic)
+
+            return formatResult(plans.length, skippedIssues)
         },
     })
 }