feat: annotate tool outputs with token counts

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"
@@ -115,8 +115,10 @@ export function createChatMessageTransformHandler(
         purgeErrors(state, logger, config, output.messages)
 
         prune(state, logger, config, output.messages)
+
         insertCompressToolContext(state, config, logger, output.messages)
 
+        annotateContext(output.messages)
         applyPendingManualTriggerPrompt(state, output.messages, logger)
 
         if (state.sessionId) {

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,26 @@ 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 {
+    for (const msg of messages) {
+        const parts = Array.isArray(msg.parts) ? msg.parts : []
+        for (const part of parts) {
+            if (part.type !== "tool") continue
+            const tokens = countToolTokens(part)
+            if (tokens <= 0) continue
+            const tag = `[${formatTokenCount(tokens)}]`
+            if (part.state?.status === "completed" && typeof part.state.output === "string") {
+                part.state.output = `${tag}\n${part.state.output}`
+            }
+            if (part.state?.status === "error" && typeof part.state.error === "string") {
+                part.state.error = `${tag}\n${part.state.error}`
+            }
+        }
+    }
+}