Merge pull request #375 from Opencode-DCP/dev

fix: search transformed message view for compress boundaries

Author: Tarquinen

lib/messages/prune.ts

@@ -38,6 +38,12 @@ const pruneFullTool = (state: SessionState, logger: Logger, messages: WithParts[
             if (part.type !== "tool") {
                 continue
             }
+
+            if (part.tool === "compress" && part.state.status === "completed") {
+                partsToRemove.push(part.callID)
+                continue
+            }
+
             if (!state.prune.tools.has(part.callID)) {
                 continue
             }

lib/tools/compress.ts

@@ -7,6 +7,7 @@ import { loadPrompt } from "../prompts"
 import { getCurrentParams, getTotalToolTokens, countMessageTextTokens } from "../strategies/utils"
 import { findStringInMessages, collectToolIdsInRange, collectMessageIdsInRange } from "./utils"
 import { sendCompressNotification } from "../ui/notification"
+import { prune as applyPruneTransforms } from "../messages/prune"
 
 const COMPRESS_TOOL_DESCRIPTION = loadPrompt("compress-tool-spec")
 
@@ -82,37 +83,56 @@ export function createCompressTool(ctx: PruneToolContext): ReturnType<typeof too
                 ctx.config.manualMode.enabled,
             )
 
+            const transformedMessages = structuredClone(messages) as WithParts[]
+            applyPruneTransforms(state, logger, ctx.config, transformedMessages)
+
             const startResult = findStringInMessages(
-                messages,
+                transformedMessages,
                 startString,
                 logger,
-                state.compressSummaries,
                 "startString",
             )
             const endResult = findStringInMessages(
-                messages,
+                transformedMessages,
                 endString,
                 logger,
-                state.compressSummaries,
                 "endString",
             )
 
-            if (startResult.messageIndex > endResult.messageIndex) {
+            let rawStartIndex = messages.findIndex((m) => m.info.id === startResult.messageId)
+            let rawEndIndex = messages.findIndex((m) => m.info.id === endResult.messageId)
+
+            // If a boundary matched inside a synthetic compress summary message,
+            // resolve it back to the summary's anchor message in the raw messages
+            if (rawStartIndex === -1) {
+                const summary = state.compressSummaries.find((s) => s.summary.includes(startString))
+                if (summary) {
+                    rawStartIndex = messages.findIndex((m) => m.info.id === summary.anchorMessageId)
+                }
+            }
+            if (rawEndIndex === -1) {
+                const summary = state.compressSummaries.find((s) => s.summary.includes(endString))
+                if (summary) {
+                    rawEndIndex = messages.findIndex((m) => m.info.id === summary.anchorMessageId)
+                }
+            }
+
+            if (rawStartIndex === -1 || rawEndIndex === -1) {
+                throw new Error(`Failed to map boundary matches back to raw messages`)
+            }
+
+            if (rawStartIndex > rawEndIndex) {
                 throw new Error(
                     `startString appears after endString in the conversation. Start must come before end.`,
                 )
             }
 
-            const containedToolIds = collectToolIdsInRange(
-                messages,
-                startResult.messageIndex,
-                endResult.messageIndex,
-            )
+            const containedToolIds = collectToolIdsInRange(messages, rawStartIndex, rawEndIndex)
 
             const containedMessageIds = collectMessageIdsInRange(
                 messages,
-                startResult.messageIndex,
-                endResult.messageIndex,
+                rawStartIndex,
+                rawEndIndex,
             )
 
             // Remove any existing summaries whose anchors are now inside this range
@@ -132,38 +152,45 @@ export function createCompressTool(ctx: PruneToolContext): ReturnType<typeof too
             }
             state.compressSummaries.push(compressSummary)
 
+            const compressedMessageIds = containedMessageIds.filter(
+                (id) => !state.prune.messages.has(id),
+            )
+            const compressedToolIds = containedToolIds.filter((id) => !state.prune.tools.has(id))
+
             let textTokens = 0
-            for (let i = startResult.messageIndex; i <= endResult.messageIndex; i++) {
-                const msgId = messages[i].info.id
-                if (!state.prune.messages.has(msgId)) {
-                    const tokens = countMessageTextTokens(messages[i])
+            for (const msgId of compressedMessageIds) {
+                const msg = messages.find((m) => m.info.id === msgId)
+                if (msg) {
+                    const tokens = countMessageTextTokens(msg)
                     textTokens += tokens
                     state.prune.messages.set(msgId, tokens)
                 }
             }
-            const newToolIds = containedToolIds.filter((id) => !state.prune.tools.has(id))
-            const toolTokens = getTotalToolTokens(state, newToolIds)
-            for (const id of newToolIds) {
+            const toolTokens = getTotalToolTokens(state, compressedToolIds)
+            for (const id of compressedToolIds) {
                 const entry = state.toolParameters.get(id)
                 state.prune.tools.set(id, entry?.tokenCount ?? 0)
             }
             const estimatedCompressedTokens = textTokens + toolTokens
 
             state.stats.pruneTokenCounter += estimatedCompressedTokens
 
+            const rawStartResult = { messageId: startResult.messageId, messageIndex: rawStartIndex }
+            const rawEndResult = { messageId: endResult.messageId, messageIndex: rawEndIndex }
+
             const currentParams = getCurrentParams(state, messages, logger)
             await sendCompressNotification(
                 client,
                 logger,
                 ctx.config,
                 state,
                 sessionId,
-                containedToolIds,
-                containedMessageIds,
+                compressedToolIds,
+                compressedMessageIds,
                 topic,
                 summary,
-                startResult,
-                endResult,
+                rawStartResult,
+                rawEndResult,
                 messages.length,
                 currentParams,
             )
@@ -184,8 +211,7 @@ export function createCompressTool(ctx: PruneToolContext): ReturnType<typeof too
                 logger.error("Failed to persist state", { error: err.message }),
             )
 
-            const messagesCompressed = endResult.messageIndex - startResult.messageIndex + 1
-            return `Compressed ${messagesCompressed} messages (${containedToolIds.length} tool calls) into summary. The content will be replaced with your summary.`
+            return `Compressed ${compressedMessageIds.length} messages (${compressedToolIds.length} tool calls) into summary. The content will be replaced with your summary.`
         },
     })
 }

lib/tools/utils.ts

@@ -1,5 +1,5 @@
 import { partial_ratio } from "fuzzball"
-import type { WithParts, CompressSummary } from "../state"
+import type { WithParts } from "../state"
 import type { Logger } from "../logger"
 
 export interface FuzzyConfig {
@@ -76,38 +76,13 @@ function extractMessageContent(msg: WithParts): string {
     return content
 }
 
-function findExactMatches(
-    messages: WithParts[],
-    searchString: string,
-    compressSummaries: CompressSummary[],
-): MatchResult[] {
+function findExactMatches(messages: WithParts[], searchString: string): MatchResult[] {
     const matches: MatchResult[] = []
-    const seenMessageIds = new Set<string>()
-
-    // Search compress summaries first
-    for (const summary of compressSummaries) {
-        if (summary.summary.includes(searchString)) {
-            const anchorIndex = messages.findIndex((m) => m.info.id === summary.anchorMessageId)
-            if (anchorIndex !== -1 && !seenMessageIds.has(summary.anchorMessageId)) {
-                seenMessageIds.add(summary.anchorMessageId)
-                matches.push({
-                    messageId: summary.anchorMessageId,
-                    messageIndex: anchorIndex,
-                    score: 100,
-                    matchType: "exact",
-                })
-            }
-        }
-    }
 
-    // Search raw messages
     for (let i = 0; i < messages.length; i++) {
         const msg = messages[i]
-        if (seenMessageIds.has(msg.info.id)) continue
-
         const content = extractMessageContent(msg)
         if (content.includes(searchString)) {
-            seenMessageIds.add(msg.info.id)
             matches.push({
                 messageId: msg.info.id,
                 messageIndex: i,
@@ -123,38 +98,15 @@ function findExactMatches(
 function findFuzzyMatches(
     messages: WithParts[],
     searchString: string,
-    compressSummaries: CompressSummary[],
     minScore: number,
 ): MatchResult[] {
     const matches: MatchResult[] = []
-    const seenMessageIds = new Set<string>()
 
-    // Search compress summaries first
-    for (const summary of compressSummaries) {
-        const score = partial_ratio(searchString, summary.summary)
-        if (score >= minScore) {
-            const anchorIndex = messages.findIndex((m) => m.info.id === summary.anchorMessageId)
-            if (anchorIndex !== -1 && !seenMessageIds.has(summary.anchorMessageId)) {
-                seenMessageIds.add(summary.anchorMessageId)
-                matches.push({
-                    messageId: summary.anchorMessageId,
-                    messageIndex: anchorIndex,
-                    score,
-                    matchType: "fuzzy",
-                })
-            }
-        }
-    }
-
-    // Search raw messages
     for (let i = 0; i < messages.length; i++) {
         const msg = messages[i]
-        if (seenMessageIds.has(msg.info.id)) continue
-
         const content = extractMessageContent(msg)
         const score = partial_ratio(searchString, content)
         if (score >= minScore) {
-            seenMessageIds.add(msg.info.id)
             matches.push({
                 messageId: msg.info.id,
                 messageIndex: i,
@@ -171,14 +123,13 @@ export function findStringInMessages(
     messages: WithParts[],
     searchString: string,
     logger: Logger,
-    compressSummaries: CompressSummary[] = [],
     stringType: "startString" | "endString",
     fuzzyConfig: FuzzyConfig = DEFAULT_FUZZY_CONFIG,
 ): { messageId: string; messageIndex: number } {
     const searchableMessages = messages.length > 1 ? messages.slice(0, -1) : messages
     const lastMessage = messages.length > 0 ? messages[messages.length - 1] : undefined
 
-    const exactMatches = findExactMatches(searchableMessages, searchString, compressSummaries)
+    const exactMatches = findExactMatches(searchableMessages, searchString)
 
     if (exactMatches.length === 1) {
         return { messageId: exactMatches[0].messageId, messageIndex: exactMatches[0].messageIndex }
@@ -191,12 +142,7 @@ export function findStringInMessages(
         )
     }
 
-    const fuzzyMatches = findFuzzyMatches(
-        searchableMessages,
-        searchString,
-        compressSummaries,
-        fuzzyConfig.minScore,
-    )
+    const fuzzyMatches = findFuzzyMatches(searchableMessages, searchString, fuzzyConfig.minScore)
 
     if (fuzzyMatches.length === 0) {
         if (lastMessage) {