|
| 1 | +/** |
| 2 | + * Tests for protecting compress tool calls (which carry summaries) from being |
| 3 | + * included in subsequent compression ranges. |
| 4 | + * |
| 5 | + * Background: compress tool calls live inside assistant messages a few positions |
| 6 | + * after the range they compressed. When the model issues a new sequential |
| 7 | + * compress whose range starts right after the previous one's end, the previous |
| 8 | + * compress call falls inside the new range and gets pruned — destroying the |
| 9 | + * accumulated summary chain. Adding "compress" to the default protectedTools |
| 10 | + * (COMPRESS_DEFAULT_PROTECTED_TOOLS) makes filterProtectedToolMessages |
| 11 | + * hard-exclude those messages (Bug 39 mechanism), so summaries survive. |
| 12 | + */ |
| 13 | +import assert from "node:assert/strict" |
| 14 | +import test from "node:test" |
| 15 | +import { messageContainsProtectedTool, filterProtectedToolMessages } from "../lib/compress/protected-content" |
| 16 | +import type { SelectionResolution, SearchContext } from "../lib/compress/types" |
| 17 | +import type { WithParts } from "../lib/state" |
| 18 | + |
| 19 | +const DEFAULT_PROTECTED = ["skill", "compress"] |
| 20 | + |
| 21 | +function makeCompressCallPart(callID: string, summary: string) { |
| 22 | + return { |
| 23 | + type: "tool" as const, |
| 24 | + callID, |
| 25 | + tool: "compress", |
| 26 | + state: { |
| 27 | + status: "completed" as const, |
| 28 | + input: { content: [{ startId: "m00001", endId: "m00010", summary }] }, |
| 29 | + output: "compressed", |
| 30 | + }, |
| 31 | + } |
| 32 | +} |
| 33 | + |
| 34 | +function makeTextPart(id: string, text: string) { |
| 35 | + return { type: "text" as const, id, text } |
| 36 | +} |
| 37 | + |
| 38 | +function makeMessage(id: string, role: "user" | "assistant", parts: any[]): WithParts { |
| 39 | + return { |
| 40 | + info: { id, role, sessionID: "ses-test", time: { created: 1 } } as any, |
| 41 | + parts, |
| 42 | + } |
| 43 | +} |
| 44 | + |
| 45 | +function makeSearchContext(messages: WithParts[]): SearchContext { |
| 46 | + const rawMessagesById = new Map<string, WithParts>() |
| 47 | + const rawIndexById = new Map<string, number>() |
| 48 | + messages.forEach((m, i) => { |
| 49 | + rawMessagesById.set(m.info.id, m) |
| 50 | + rawIndexById.set(m.info.id, i) |
| 51 | + }) |
| 52 | + return { |
| 53 | + rawMessages: messages, |
| 54 | + rawMessagesById, |
| 55 | + rawIndexById, |
| 56 | + summaryByBlockId: new Map(), |
| 57 | + } |
| 58 | +} |
| 59 | + |
| 60 | +function makeSelection(messageIds: string[]): SelectionResolution { |
| 61 | + return { |
| 62 | + startReference: { kind: "message", rawIndex: 0, messageId: messageIds[0] }, |
| 63 | + endReference: { kind: "message", rawIndex: messageIds.length - 1, messageId: messageIds[messageIds.length - 1] }, |
| 64 | + messageIds, |
| 65 | + messageTokenById: new Map(messageIds.map((id) => [id, 100])), |
| 66 | + toolIds: [], |
| 67 | + requiredBlockIds: [], |
| 68 | + } |
| 69 | +} |
| 70 | + |
| 71 | +test("messageContainsProtectedTool: compress tool call is protected when 'compress' is in the list", () => { |
| 72 | + const msg = makeMessage("msg-compress-call", "assistant", [ |
| 73 | + makeTextPart("p1", "Let me compress the earlier findings."), |
| 74 | + makeCompressCallPart("call-1", "Summary of earlier work..."), |
| 75 | + ]) |
| 76 | + assert.equal(messageContainsProtectedTool(msg, DEFAULT_PROTECTED, []), true) |
| 77 | +}) |
| 78 | + |
| 79 | +test("messageContainsProtectedTool: compress tool call is NOT protected when 'compress' is absent (opt-out)", () => { |
| 80 | + const msg = makeMessage("msg-compress-call", "assistant", [ |
| 81 | + makeCompressCallPart("call-1", "Summary of earlier work..."), |
| 82 | + ]) |
| 83 | + assert.equal(messageContainsProtectedTool(msg, ["skill"], []), false) |
| 84 | + assert.equal(messageContainsProtectedTool(msg, [], []), false) |
| 85 | +}) |
| 86 | + |
| 87 | +test("messageContainsProtectedTool: plain text message is never protected", () => { |
| 88 | + const msg = makeMessage("msg-text", "user", [makeTextPart("p1", "Hello world")]) |
| 89 | + assert.equal(messageContainsProtectedTool(msg, DEFAULT_PROTECTED, []), false) |
| 90 | +}) |
| 91 | + |
| 92 | +test("filterProtectedToolMessages: removes compress-call message from selection, keeps surrounding messages", () => { |
| 93 | + const compressMsg = makeMessage("msg-compress", "assistant", [ |
| 94 | + makeTextPart("p1", "Compressing now."), |
| 95 | + makeCompressCallPart("call-compress", "Previous summary content..."), |
| 96 | + ]) |
| 97 | + const plainMsg1 = makeMessage("msg-plain-1", "user", [makeTextPart("p2", "User question")]) |
| 98 | + const plainMsg2 = makeMessage("msg-plain-2", "assistant", [makeTextPart("p3", "Assistant answer")]) |
| 99 | + |
| 100 | + const ctx = makeSearchContext([plainMsg1, compressMsg, plainMsg2]) |
| 101 | + const selection = makeSelection(["msg-plain-1", "msg-compress", "msg-plain-2"]) |
| 102 | + |
| 103 | + const result = filterProtectedToolMessages(selection, ctx, DEFAULT_PROTECTED, []) |
| 104 | + |
| 105 | + assert.deepEqual(result.messageIds, ["msg-plain-1", "msg-plain-2"]) |
| 106 | + assert.equal(result.messageTokenById.size, 2) |
| 107 | + assert.ok(result.messageTokenById.has("msg-plain-1")) |
| 108 | + assert.ok(result.messageTokenById.has("msg-plain-2")) |
| 109 | + assert.ok(!result.messageTokenById.has("msg-compress")) |
| 110 | +}) |
| 111 | + |
| 112 | +test("filterProtectedToolMessages: no-op when 'compress' is not in protectedTools (old behavior)", () => { |
| 113 | + const compressMsg = makeMessage("msg-compress", "assistant", [ |
| 114 | + makeCompressCallPart("call-1", "Summary..."), |
| 115 | + ]) |
| 116 | + const ctx = makeSearchContext([compressMsg]) |
| 117 | + const selection = makeSelection(["msg-compress"]) |
| 118 | + |
| 119 | + const result = filterProtectedToolMessages(selection, ctx, ["skill"], []) |
| 120 | + assert.deepEqual(result.messageIds, ["msg-compress"]) |
| 121 | +}) |
| 122 | + |
| 123 | +test("filterProtectedToolMessages: all-compress-call selection becomes empty (all excluded)", () => { |
| 124 | + const msg1 = makeMessage("msg-c1", "assistant", [makeCompressCallPart("c1", "Summary A")]) |
| 125 | + const msg2 = makeMessage("msg-c2", "assistant", [makeCompressCallPart("c2", "Summary B")]) |
| 126 | + |
| 127 | + const ctx = makeSearchContext([msg1, msg2]) |
| 128 | + const selection = makeSelection(["msg-c1", "msg-c2"]) |
| 129 | + |
| 130 | + const result = filterProtectedToolMessages(selection, ctx, DEFAULT_PROTECTED, []) |
| 131 | + assert.equal(result.messageIds.length, 0) |
| 132 | + assert.equal(result.messageTokenById.size, 0) |
| 133 | +}) |
0 commit comments