fix: move assignMessageRefs after prune to prevent orphan alias leak

lib/hooks.ts

@@ -124,11 +124,11 @@ export function createChatMessageTransformHandler(
 
         stripHallucinations(output.messages)
         cacheSystemPromptTokens(state, output.messages)
-        assignMessageRefs(state, output.messages)
         syncCompressionBlocks(state, logger, output.messages)
         syncToolCache(state, config, logger, output.messages)
         buildToolIdList(state, output.messages)
         prune(state, logger, config, output.messages)
+        assignMessageRefs(state, output.messages)
         await injectExtendedSubAgentResults(
             client,
             state,

tests/message-ids.test.ts

@@ -1,9 +1,42 @@
 import assert from "node:assert/strict"
 import test from "node:test"
+import type { PluginConfig } from "../lib/config"
+import { createChatMessageTransformHandler } from "../lib/hooks"
 import { Logger } from "../lib/logger"
 import { assignMessageRefs } from "../lib/message-ids"
 import { checkSession, createSessionState, type WithParts } from "../lib/state"
 
+function buildConfig(permission: "allow" | "ask" | "deny" = "allow"): PluginConfig {
+    return {
+        enabled: true,
+        debug: false,
+        pruneNotification: "off",
+        pruneNotificationType: "chat",
+        commands: { enabled: true, protectedTools: [] },
+        manualMode: { enabled: false, automaticStrategies: true },
+        turnProtection: { enabled: false, turns: 4 },
+        experimental: { allowSubAgents: false, customPrompts: false },
+        protectedFilePatterns: [],
+        compress: {
+            mode: "message",
+            permission,
+            showCompression: false,
+            maxContextLimit: 150000,
+            minContextLimit: 50000,
+            nudgeFrequency: 5,
+            iterationNudgeThreshold: 15,
+            nudgeForce: "soft",
+            protectedTools: ["task"],
+            protectTags: false,
+            protectUserMessages: false,
+        },
+        strategies: {
+            deduplication: { enabled: true, protectedTools: [] },
+            purgeErrors: { enabled: true, turns: 4, protectedTools: [] },
+        },
+    }
+}
+
 function textPart(messageID: string, sessionID: string, id: string, text: string) {
     return {
         id,
@@ -87,3 +120,106 @@ test("checkSession resets message id aliases after native compaction", async ()
     assert.equal(state.messageIds.byRef.get("m0002"), "msg-user-follow-up")
     assert.equal(state.messageIds.nextRef, 3)
 })
+
+test("assignMessageRefs after prune prevents orphan aliases (RED->GREEN: fails on old pipeline order)", async () => {
+    const sessionID = "ses_alias_leak_test"
+    const state = createSessionState()
+    const logger = new Logger(false)
+
+    // Prevent checkSession from resetting state via ensureSessionInitialized
+    state.sessionId = sessionID
+
+    // Simulate a completed compress — this message should be removed by prune
+    state.prune.messages.byMessageId.set("msg-pruned", {
+        tokenCount: 100,
+        allBlockIds: [1],
+        activeBlockIds: [1],
+    })
+
+    const messages: WithParts[] = [
+        {
+            info: {
+                id: "msg-user",
+                role: "user",
+                sessionID,
+                agent: "assistant",
+                model: { providerID: "anthropic", modelID: "claude-test" },
+                time: { created: 1 },
+            } as WithParts["info"],
+            parts: [
+                {
+                    id: "msg-user-part",
+                    messageID: "msg-user",
+                    sessionID,
+                    type: "text",
+                    text: "Test message",
+                },
+            ],
+        },
+        {
+            info: {
+                id: "msg-assistant-survivor",
+                role: "assistant",
+                sessionID,
+                agent: "assistant",
+                time: { created: 2 },
+            } as WithParts["info"],
+            parts: [
+                {
+                    id: "msg-assistant-survivor-part",
+                    messageID: "msg-assistant-survivor",
+                    sessionID,
+                    type: "text",
+                    text: "Survivor response",
+                },
+            ],
+        },
+        {
+            info: {
+                id: "msg-pruned",
+                role: "assistant",
+                sessionID,
+                agent: "assistant",
+                time: { created: 3 },
+            } as WithParts["info"],
+            parts: [
+                {
+                    id: "msg-pruned-part",
+                    messageID: "msg-pruned",
+                    sessionID,
+                    type: "text",
+                    text: "Will be removed",
+                },
+            ],
+        },
+    ]
+
+    const handler = createChatMessageTransformHandler(
+        {} as any,
+        state,
+        logger,
+        buildConfig("allow"),
+        {
+            reload() {},
+            getRuntimePrompts() {
+                return {} as any
+            },
+        } as any,
+        { global: undefined, agents: {} },
+    )
+
+    await handler({}, { messages })
+
+    // Survivors get aliases
+    assert.ok(state.messageIds.byRawId.has("msg-user"), "user message should have alias")
+    assert.ok(state.messageIds.byRawId.has("msg-assistant-survivor"), "survivor should have alias")
+
+    // RED on old pipeline (assignMessageRefs before prune): msg-pruned got an alias before
+    //   being removed, and no cleanup follows → assertion FAILS (alias leak detected)
+    // GREEN on new pipeline (assignMessageRefs after prune): prune removes msg-pruned first,
+    //   then assignMessageRefs never sees it → assertion PASSES (no leak)
+    assert.ok(
+        !state.messageIds.byRawId.has("msg-pruned"),
+        "pruned message should NOT have an alias — leak if it does",
+    )
+})