feat: add uid + token counts to tool outputs and muid to user message

Author: spoons-and-mirrors

lib/hooks.ts

@@ -13,14 +13,49 @@ import { handleHelpCommand } from "./commands/help"
 import { handleSweepCommand } from "./commands/sweep"
 import { handleManualToggleCommand, handleManualTriggerCommand } from "./commands/manual"
 import { ensureSessionInitialized } from "./state/state"
-import { getCurrentParams } from "./strategies/utils"
+import { getCurrentParams, countToolTokens } 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[],
@@ -118,6 +153,7 @@ export function createChatMessageTransformHandler(
         insertCompressToolContext(state, config, logger, output.messages)
 
         applyPendingManualTriggerPrompt(state, output.messages, logger)
+        annotateContext(output.messages)
 
         if (state.sessionId) {
             await logger.saveContext(state.sessionId, output.messages)