feat: integrate context annotation into hook and compress tool

Author: spoons-and-mirrors

lib/hooks.ts

@@ -4,7 +4,7 @@ import type { PluginConfig } from "./config"
 import { syncToolCache } from "./state/tool-cache"
 import { deduplicate, supersedeWrites, purgeErrors } from "./strategies"
 import { prune, insertCompressToolContext } from "./messages"
-import { buildToolIdList, isIgnoredUserMessage } from "./messages/utils"
+import { buildToolIdList, isIgnoredUserMessage, annotateContext } from "./messages/utils"
 import { checkSession } from "./state"
 import { renderSystemPrompt } from "./prompts"
 import { handleStatsCommand } from "./commands/stats"
@@ -13,49 +13,14 @@ import { handleHelpCommand } from "./commands/help"
 import { handleSweepCommand } from "./commands/sweep"
 import { handleManualToggleCommand, handleManualTriggerCommand } from "./commands/manual"
 import { ensureSessionInitialized } from "./state/state"
-import { getCurrentParams, countToolTokens } from "./strategies/utils"
+import { getCurrentParams } from "./strategies/utils"
 
 const INTERNAL_AGENT_SIGNATURES = [
     "You are a title generator",
     "You are a helpful AI assistant tasked with summarizing conversations",
     "Summarize what was done in this conversation",
 ]
 
-function formatTokenCount(n: number): string {
-    if (n >= 1000) return `${(n / 1000).toFixed(1)}k tokens`
-    return `${n} tokens`
-}
-
-function annotateContext(messages: WithParts[]): void {
-    let uid = 0
-    let muid = 0
-    for (const msg of messages) {
-        const parts = Array.isArray(msg.parts) ? msg.parts : []
-
-        if (msg.info.role === "user" && !isIgnoredUserMessage(msg)) {
-            muid++
-            const text = parts.find(
-                (p) => p.type === "text" && !(p as any).ignored && !(p as any).synthetic,
-            )
-            if (text && text.type === "text") {
-                text.text = `[muid_${muid}]\n${text.text}`
-            }
-        }
-
-        for (const part of parts) {
-            if (part.type !== "tool") continue
-            uid++
-            const tokens = countToolTokens(part)
-            const tag = tokens > 0 ? `[uid_${uid}, ${formatTokenCount(tokens)}]` : `[uid_${uid}]`
-            if (part.state?.status === "completed" && typeof part.state.output === "string") {
-                part.state.output = `${tag}\n${part.state.output}`
-            } else if (part.state?.status === "error" && typeof part.state.error === "string") {
-                part.state.error = `${tag}\n${part.state.error}`
-            }
-        }
-    }
-}
-
 function applyPendingManualTriggerPrompt(
     state: SessionState,
     messages: WithParts[],

lib/messages/utils.ts

@@ -3,6 +3,7 @@ import { isMessageCompacted } from "../shared-utils"
 import { Logger } from "../logger"
 import type { SessionState, WithParts } from "../state"
 import type { UserMessage } from "@opencode-ai/sdk/v2"
+import { countToolTokens } from "../strategies/utils"
 
 export const COMPRESS_SUMMARY_PREFIX = "[Compressed conversation block]\n\n"
 
@@ -267,3 +268,38 @@ export const isIgnoredUserMessage = (message: WithParts): boolean => {
 export const findMessageIndex = (messages: WithParts[], messageId: string): number => {
     return messages.findIndex((msg) => msg.info.id === messageId)
 }
+
+function formatTokenCount(n: number): string {
+    if (n >= 1000) return `${(n / 1000).toFixed(1)}k tokens`
+    return `${n} tokens`
+}
+
+export function annotateContext(messages: WithParts[]): void {
+    let uid = 0
+    let muid = 0
+    for (const msg of messages) {
+        const parts = Array.isArray(msg.parts) ? msg.parts : []
+
+        if (msg.info.role === "user" && !isIgnoredUserMessage(msg)) {
+            muid++
+            const text = parts.find(
+                (p) => p.type === "text" && !(p as any).ignored && !(p as any).synthetic,
+            )
+            if (text && text.type === "text") {
+                text.text = `[muid_${muid}]\n${text.text}`
+            }
+        }
+
+        for (const part of parts) {
+            if (part.type !== "tool") continue
+            uid++
+            const tokens = countToolTokens(part)
+            const tag = tokens > 0 ? `[uid_${uid}, ${formatTokenCount(tokens)}]` : `[uid_${uid}]`
+            if (part.state?.status === "completed" && typeof part.state.output === "string") {
+                part.state.output = `${tag}\n${part.state.output}`
+            } else if (part.state?.status === "error" && typeof part.state.error === "string") {
+                part.state.error = `${tag}\n${part.state.error}`
+            }
+        }
+    }
+}

lib/tools/compress.ts

@@ -6,6 +6,7 @@ import { saveSessionState } from "../state/persistence"
 import { loadPrompt } from "../prompts"
 import { getCurrentParams, getTotalToolTokens, countMessageTextTokens } from "../strategies/utils"
 import { findStringInMessages, collectToolIdsInRange, collectMessageIdsInRange } from "./utils"
+import { annotateContext } from "../messages/utils"
 import { sendCompressNotification } from "../ui/notification"
 import { prune as applyPruneTransforms } from "../messages/prune"
 
@@ -82,6 +83,7 @@ export function createCompressTool(ctx: ToolContext): ReturnType<typeof tool> {
 
             const transformedMessages = structuredClone(messages) as WithParts[]
             applyPruneTransforms(state, logger, ctx.config, transformedMessages)
+            annotateContext(transformedMessages)
 
             const startResult = findStringInMessages(
                 transformedMessages,