@royalcat/opencode-dcp-rc 0.0.0-stage → 4.0.0

package/lib/compress/search.ts

@@ -0,0 +1,277 @@
+import type { SessionState, WithParts } from "../state"
+import type { PluginConfig } from "../config"
+import { formatBlockRef, parseBoundaryId } from "../message-ids"
+import { isIgnoredUserMessage, isProtectedUserMessage } from "../messages/query"
+import { filterMessages } from "../messages/shape"
+import { countAllMessageTokens } from "../token-utils"
+import type { BoundaryReference, SearchContext, SelectionResolution } from "./types"
+
+export async function fetchSessionMessages(client: any, sessionId: string): Promise<WithParts[]> {
+    const response = await client.session.messages({
+        path: { id: sessionId },
+    })
+
+    return filterMessages(response?.data || response)
+}
+
+export function buildSearchContext(state: SessionState, rawMessages: WithParts[]): SearchContext {
+    const rawMessagesById = new Map<string, WithParts>()
+    const rawIndexById = new Map<string, number>()
+    for (const msg of rawMessages) {
+        rawMessagesById.set(msg.info.id, msg)
+    }
+    for (let index = 0; index < rawMessages.length; index++) {
+        const message = rawMessages[index]
+        if (!message) {
+            continue
+        }
+        rawIndexById.set(message.info.id, index)
+    }
+
+    const summaryByBlockId = new Map()
+    for (const [blockId, block] of state.prune.messages.blocksById) {
+        if (!block.active) {
+            continue
+        }
+        summaryByBlockId.set(blockId, block)
+    }
+
+    return {
+        rawMessages,
+        rawMessagesById,
+        rawIndexById,
+        summaryByBlockId,
+    }
+}
+
+export function resolveBoundaryIds(
+    context: SearchContext,
+    state: SessionState,
+    config: PluginConfig,
+    startId: string,
+    endId: string,
+): { startReference: BoundaryReference; endReference: BoundaryReference } {
+    const lookup = buildBoundaryLookup(context, state, config)
+    const issues: string[] = []
+    const parsedStartId = parseBoundaryId(startId, state.idFormat)
+    const parsedEndId = parseBoundaryId(endId, state.idFormat)
+    const messageExample = state.idFormat === "compact" ? "@4@" : "mNNNN"
+    const blockExample = state.idFormat === "compact" ? "@b1@" : "bN"
+
+    if (parsedStartId === null) {
+        issues.push(
+            `startId is invalid. Use an injected message ID (${messageExample}) or block ID (${blockExample}).`,
+        )
+    }
+
+    if (parsedEndId === null) {
+        issues.push(
+            `endId is invalid. Use an injected message ID (${messageExample}) or block ID (${blockExample}).`,
+        )
+    }
+
+    if (issues.length > 0) {
+        throw new Error(
+            issues.length === 1 ? issues[0] : issues.map((issue) => `- ${issue}`).join("\n"),
+        )
+    }
+
+    if (!parsedStartId || !parsedEndId) {
+        throw new Error("Invalid boundary ID(s)")
+    }
+
+    const startReference = lookup.get(parsedStartId.ref)
+    const endReference = lookup.get(parsedEndId.ref)
+
+    if (!startReference) {
+        issues.push(
+            `startId ${parsedStartId.ref} is not available in the current conversation context. Choose an injected ID visible in context.`,
+        )
+    }
+
+    if (!endReference) {
+        issues.push(
+            `endId ${parsedEndId.ref} is not available in the current conversation context. Choose an injected ID visible in context.`,
+        )
+    }
+
+    if (issues.length > 0) {
+        throw new Error(
+            issues.length === 1 ? issues[0] : issues.map((issue) => `- ${issue}`).join("\n"),
+        )
+    }
+
+    if (!startReference || !endReference) {
+        throw new Error("Failed to resolve boundary IDs")
+    }
+
+    if (startReference.rawIndex > endReference.rawIndex) {
+        throw new Error(
+            `startId ${parsedStartId.ref} appears after endId ${parsedEndId.ref} in the conversation. Start must come before end.`,
+        )
+    }
+
+    return { startReference, endReference }
+}
+
+export function resolveSelection(
+    context: SearchContext,
+    config: PluginConfig,
+    startReference: BoundaryReference,
+    endReference: BoundaryReference,
+): SelectionResolution {
+    const startRawIndex = startReference.rawIndex
+    const endRawIndex = endReference.rawIndex
+    const messageIds: string[] = []
+    const messageSeen = new Set<string>()
+    const toolIds: string[] = []
+    const toolSeen = new Set<string>()
+    const requiredBlockIds: number[] = []
+    const requiredBlockSeen = new Set<number>()
+    const messageTokenById = new Map<string, number>()
+
+    for (let index = startRawIndex; index <= endRawIndex; index++) {
+        const rawMessage = context.rawMessages[index]
+        if (!rawMessage) {
+            continue
+        }
+        if (isIgnoredUserMessage(rawMessage) || isProtectedUserMessage(config, rawMessage)) {
+            continue
+        }
+
+        const messageId = rawMessage.info.id
+        if (!messageSeen.has(messageId)) {
+            messageSeen.add(messageId)
+            messageIds.push(messageId)
+        }
+
+        if (!messageTokenById.has(messageId)) {
+            messageTokenById.set(messageId, countAllMessageTokens(rawMessage))
+        }
+
+        const parts = Array.isArray(rawMessage.parts) ? rawMessage.parts : []
+        for (const part of parts) {
+            if (part.type !== "tool" || !part.callID) {
+                continue
+            }
+            if (toolSeen.has(part.callID)) {
+                continue
+            }
+            toolSeen.add(part.callID)
+            toolIds.push(part.callID)
+        }
+    }
+
+    const selectedMessageIds = new Set(messageIds)
+    const summariesInSelection: Array<{ blockId: number; rawIndex: number }> = []
+    for (const summary of context.summaryByBlockId.values()) {
+        if (!selectedMessageIds.has(summary.anchorMessageId)) {
+            continue
+        }
+
+        const anchorIndex = context.rawIndexById.get(summary.anchorMessageId)
+        if (anchorIndex === undefined) {
+            continue
+        }
+
+        summariesInSelection.push({
+            blockId: summary.blockId,
+            rawIndex: anchorIndex,
+        })
+    }
+
+    summariesInSelection.sort((a, b) => a.rawIndex - b.rawIndex || a.blockId - b.blockId)
+    for (const summary of summariesInSelection) {
+        if (requiredBlockSeen.has(summary.blockId)) {
+            continue
+        }
+        requiredBlockSeen.add(summary.blockId)
+        requiredBlockIds.push(summary.blockId)
+    }
+
+    if (messageIds.length === 0) {
+        throw new Error(
+            "Failed to map boundary matches back to raw messages. Choose boundaries that include original conversation messages.",
+        )
+    }
+
+    return {
+        startReference,
+        endReference,
+        messageIds,
+        messageTokenById,
+        toolIds,
+        requiredBlockIds,
+    }
+}
+
+export function resolveAnchorMessageId(startReference: BoundaryReference): string {
+    if (startReference.kind === "compressed-block") {
+        if (!startReference.anchorMessageId) {
+            throw new Error("Failed to map boundary matches back to raw messages")
+        }
+        return startReference.anchorMessageId
+    }
+
+    if (!startReference.messageId) {
+        throw new Error("Failed to map boundary matches back to raw messages")
+    }
+    return startReference.messageId
+}
+
+function buildBoundaryLookup(
+    context: SearchContext,
+    state: SessionState,
+    config: PluginConfig,
+): Map<string, BoundaryReference> {
+    const lookup = new Map<string, BoundaryReference>()
+
+    for (const [messageRef, messageId] of state.messageIds.byRef) {
+        const rawMessage = context.rawMessagesById.get(messageId)
+        if (!rawMessage) {
+            continue
+        }
+        if (isIgnoredUserMessage(rawMessage) || isProtectedUserMessage(config, rawMessage)) {
+            continue
+        }
+
+        const rawIndex = context.rawIndexById.get(messageId)
+        if (rawIndex === undefined) {
+            continue
+        }
+        lookup.set(messageRef, {
+            kind: "message",
+            rawIndex,
+            messageId,
+        })
+    }
+
+    const summaries = Array.from(context.summaryByBlockId.values()).sort(
+        (a, b) => a.blockId - b.blockId,
+    )
+    for (const summary of summaries) {
+        const anchorMessage = context.rawMessagesById.get(summary.anchorMessageId)
+        if (!anchorMessage) {
+            continue
+        }
+        if (isIgnoredUserMessage(anchorMessage)) {
+            continue
+        }
+
+        const rawIndex = context.rawIndexById.get(summary.anchorMessageId)
+        if (rawIndex === undefined) {
+            continue
+        }
+        const blockRef = formatBlockRef(summary.blockId, state.idFormat)
+        if (!lookup.has(blockRef)) {
+            lookup.set(blockRef, {
+                kind: "compressed-block",
+                rawIndex,
+                blockId: summary.blockId,
+                anchorMessageId: summary.anchorMessageId,
+            })
+        }
+    }
+
+    return lookup
+}

package/lib/compress/state.ts

@@ -0,0 +1,272 @@
+import type { CompressionBlock, PruneMessagesState, SessionState } from "../state"
+import { formatBlockRef, formatMessageIdTag, type IdFormat } from "../message-ids"
+import type { AppliedCompressionResult, CompressionStateInput, SelectionResolution } from "./types"
+
+export const COMPRESSED_BLOCK_HEADER = "[Compressed conversation section]"
+
+export function allocateBlockId(state: SessionState): number {
+    const next = state.prune.messages.nextBlockId
+    if (!Number.isInteger(next) || next < 1) {
+        state.prune.messages.nextBlockId = 2
+        return 1
+    }
+
+    state.prune.messages.nextBlockId = next + 1
+    return next
+}
+
+export function allocateRunId(state: SessionState): number {
+    const next = state.prune.messages.nextRunId
+    if (!Number.isInteger(next) || next < 1) {
+        state.prune.messages.nextRunId = 2
+        return 1
+    }
+
+    state.prune.messages.nextRunId = next + 1
+    return next
+}
+
+export function attachCompressionDuration(
+    messagesState: PruneMessagesState,
+    messageId: string,
+    callId: string,
+    durationMs: number,
+): number {
+    if (typeof durationMs !== "number" || !Number.isFinite(durationMs)) {
+        return 0
+    }
+
+    let updates = 0
+    for (const block of messagesState.blocksById.values()) {
+        if (block.compressMessageId !== messageId || block.compressCallId !== callId) {
+            continue
+        }
+
+        block.durationMs = durationMs
+        updates++
+    }
+
+    return updates
+}
+
+export function wrapCompressedSummary(
+    blockId: number,
+    summary: string,
+    format: IdFormat = "xml",
+): string {
+    const header = COMPRESSED_BLOCK_HEADER
+    const footer = formatMessageIdTag(formatBlockRef(blockId, format), undefined, format)
+    const body = summary.trim()
+    if (body.length === 0) {
+        return `${header}\n${footer}`
+    }
+    return `${header}\n${body}\n\n${footer}`
+}
+
+export function applyCompressionState(
+    state: SessionState,
+    input: CompressionStateInput,
+    selection: SelectionResolution,
+    anchorMessageId: string,
+    blockId: number,
+    summary: string,
+    consumedBlockIds: number[],
+): AppliedCompressionResult {
+    const messagesState = state.prune.messages
+    const consumed = [...new Set(consumedBlockIds.filter((id) => Number.isInteger(id) && id > 0))]
+    const included = [...consumed]
+
+    const effectiveMessageIds = new Set<string>(selection.messageIds)
+    const effectiveToolIds = new Set<string>(selection.toolIds)
+
+    for (const consumedBlockId of consumed) {
+        const consumedBlock = messagesState.blocksById.get(consumedBlockId)
+        if (!consumedBlock) {
+            continue
+        }
+        for (const messageId of consumedBlock.effectiveMessageIds) {
+            effectiveMessageIds.add(messageId)
+        }
+        for (const toolId of consumedBlock.effectiveToolIds) {
+            effectiveToolIds.add(toolId)
+        }
+    }
+
+    const initiallyActiveMessages = new Set<string>()
+    for (const messageId of effectiveMessageIds) {
+        const entry = messagesState.byMessageId.get(messageId)
+        if (entry && entry.activeBlockIds.length > 0) {
+            initiallyActiveMessages.add(messageId)
+        }
+    }
+
+    const initiallyActiveToolIds = new Set<string>()
+    for (const activeBlockId of messagesState.activeBlockIds) {
+        const activeBlock = messagesState.blocksById.get(activeBlockId)
+        if (!activeBlock || !activeBlock.active) {
+            continue
+        }
+
+        for (const toolId of activeBlock.effectiveToolIds) {
+            initiallyActiveToolIds.add(toolId)
+        }
+    }
+
+    const createdAt = Date.now()
+    const block: CompressionBlock = {
+        blockId,
+        runId: input.runId,
+        active: true,
+        deactivatedByUser: false,
+        compressedTokens: 0,
+        summaryTokens: input.summaryTokens,
+        durationMs: 0,
+        mode: input.mode,
+        topic: input.topic,
+        batchTopic: input.batchTopic,
+        startId: input.startId,
+        endId: input.endId,
+        anchorMessageId,
+        compressMessageId: input.compressMessageId,
+        compressCallId: input.compressCallId,
+        includedBlockIds: included,
+        consumedBlockIds: consumed,
+        parentBlockIds: [],
+        directMessageIds: [],
+        directToolIds: [],
+        effectiveMessageIds: [...effectiveMessageIds],
+        effectiveToolIds: [...effectiveToolIds],
+        createdAt,
+        summary,
+    }
+
+    messagesState.blocksById.set(blockId, block)
+    messagesState.activeBlockIds.add(blockId)
+    messagesState.activeByAnchorMessageId.set(anchorMessageId, blockId)
+
+    const deactivatedAt = Date.now()
+    for (const consumedBlockId of consumed) {
+        const consumedBlock = messagesState.blocksById.get(consumedBlockId)
+        if (!consumedBlock || !consumedBlock.active) {
+            continue
+        }
+
+        consumedBlock.active = false
+        consumedBlock.deactivatedAt = deactivatedAt
+        consumedBlock.deactivatedByBlockId = blockId
+        if (!consumedBlock.parentBlockIds.includes(blockId)) {
+            consumedBlock.parentBlockIds.push(blockId)
+        }
+
+        messagesState.activeBlockIds.delete(consumedBlockId)
+        const mappedBlockId = messagesState.activeByAnchorMessageId.get(
+            consumedBlock.anchorMessageId,
+        )
+        if (mappedBlockId === consumedBlockId) {
+            messagesState.activeByAnchorMessageId.delete(consumedBlock.anchorMessageId)
+        }
+    }
+
+    const removeActiveBlockId = (
+        entry: { activeBlockIds: number[] },
+        blockIdToRemove: number,
+    ): void => {
+        if (entry.activeBlockIds.length === 0) {
+            return
+        }
+        entry.activeBlockIds = entry.activeBlockIds.filter((id) => id !== blockIdToRemove)
+    }
+
+    for (const consumedBlockId of consumed) {
+        const consumedBlock = messagesState.blocksById.get(consumedBlockId)
+        if (!consumedBlock) {
+            continue
+        }
+        for (const messageId of consumedBlock.effectiveMessageIds) {
+            const entry = messagesState.byMessageId.get(messageId)
+            if (!entry) {
+                continue
+            }
+            removeActiveBlockId(entry, consumedBlockId)
+        }
+    }
+
+    for (const messageId of selection.messageIds) {
+        const tokenCount = selection.messageTokenById.get(messageId) || 0
+        const existing = messagesState.byMessageId.get(messageId)
+
+        if (!existing) {
+            messagesState.byMessageId.set(messageId, {
+                tokenCount,
+                allBlockIds: [blockId],
+                activeBlockIds: [blockId],
+            })
+            continue
+        }
+
+        existing.tokenCount = Math.max(existing.tokenCount, tokenCount)
+        if (!existing.allBlockIds.includes(blockId)) {
+            existing.allBlockIds.push(blockId)
+        }
+        if (!existing.activeBlockIds.includes(blockId)) {
+            existing.activeBlockIds.push(blockId)
+        }
+    }
+
+    for (const messageId of block.effectiveMessageIds) {
+        if (selection.messageTokenById.has(messageId)) {
+            continue
+        }
+
+        const existing = messagesState.byMessageId.get(messageId)
+        if (!existing) {
+            continue
+        }
+        if (!existing.allBlockIds.includes(blockId)) {
+            existing.allBlockIds.push(blockId)
+        }
+        if (!existing.activeBlockIds.includes(blockId)) {
+            existing.activeBlockIds.push(blockId)
+        }
+    }
+
+    let compressedTokens = 0
+    const newlyCompressedMessageIds: string[] = []
+    for (const messageId of effectiveMessageIds) {
+        const entry = messagesState.byMessageId.get(messageId)
+        if (!entry) {
+            continue
+        }
+
+        const isNowActive = entry.activeBlockIds.length > 0
+        const wasActive = initiallyActiveMessages.has(messageId)
+
+        if (isNowActive && !wasActive) {
+            compressedTokens += entry.tokenCount
+            newlyCompressedMessageIds.push(messageId)
+        }
+    }
+
+    const newlyCompressedToolIds: string[] = []
+    for (const toolId of effectiveToolIds) {
+        if (!initiallyActiveToolIds.has(toolId)) {
+            newlyCompressedToolIds.push(toolId)
+        }
+    }
+
+    block.directMessageIds = [...newlyCompressedMessageIds]
+    block.directToolIds = [...newlyCompressedToolIds]
+
+    block.compressedTokens = compressedTokens
+
+    state.stats.pruneTokenCounter += compressedTokens
+    state.stats.totalPruneTokens += state.stats.pruneTokenCounter
+    state.stats.pruneTokenCounter = 0
+
+    return {
+        compressedTokens,
+        messageIds: selection.messageIds,
+        newlyCompressedMessageIds,
+        newlyCompressedToolIds,
+    }
+}