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

package/lib/compress/rc.ts

@@ -0,0 +1,494 @@
+import { tool } from "@opencode-ai/plugin"
+import type { ToolContext, SelectionResolution } from "./types"
+import type { WithParts } from "../state"
+import { refreshManualMode } from "../state"
+import { formatBlockRef, formatMessageRef, parseBoundaryId, type IdFormat } from "../message-ids"
+import {
+    buildSearchContext,
+    fetchSessionMessages,
+    resolveAnchorMessageId,
+    resolveBoundaryIds,
+    resolveSelection,
+} from "./search"
+import {
+    allocateBlockId,
+    allocateRunId,
+    applyCompressionState,
+    wrapCompressedSummary,
+} from "./state"
+import { finalizeSession, prepareSession, type NotificationEntry } from "./pipeline"
+import { countTokens } from "../token-utils"
+import { saveSessionState } from "../state/persistence"
+import { rcFormat } from "../prompts/extensions/tool"
+import { appendProtectedPromptInfo, appendProtectedTools } from "./protected-content"
+import {
+    addCompressionUsage,
+    addCompressionUsageTotals,
+    emptyCompressionUsage,
+    estimatedCompressionUsage,
+    type CompressionUsageTotals,
+} from "./usage"
+import {
+    buildSummaryPrompt,
+    parseSummarySections,
+    restoreSummary,
+    type PriorSummary,
+    type SummarySection,
+    type SummarySegment,
+} from "./summary"
+
+export interface RcCompressToolArgs {
+    ids: string[]
+}
+
+let summaryCallCounter = 0
+
+/** Unique per hidden summary request so provider usage can be attributed. */
+function nextSummaryCallId(): string {
+    summaryCallCounter += 1
+    return `c${summaryCallCounter.toString(36)}-${Date.now().toString(36)}-${Math.random()
+        .toString(36)
+        .slice(2, 8)}`
+}
+
+/**
+ * Compress selections with a hidden, synchronous summary request.
+ *
+ * The model only selects message IDs (`ids`); summaries are generated by the
+ * plugin in a transient `session.generate` call during this tool call and are
+ * applied before the model's next step. Nothing except the `compress` tool
+ * call itself is user-visible.
+ */
+
+function validateArgs(args: RcCompressToolArgs): void {
+    if (!args || typeof args !== "object") {
+        throw new Error("compress requires an object argument with an `ids` array")
+    }
+    if (!Array.isArray(args.ids)) {
+        throw new Error("`ids` must be an array of message IDs or inclusive ranges")
+    }
+    if (args.ids.length === 0) {
+        throw new Error("`ids` must not be empty")
+    }
+}
+
+interface ParsedSelector {
+    selector: string
+    startRef: string
+    endRef: string
+}
+
+export function parseSelector(selector: string, format: IdFormat): ParsedSelector {
+    const raw = selector.trim().toLowerCase()
+    const parts = raw.split("-").map((part) => part.trim())
+    if (parts.length === 1 || parts.length === 2) {
+        const parseAny = (part: string) =>
+            parseBoundaryId(part, format) ??
+            parseBoundaryId(part, format === "compact" ? "xml" : "compact")
+        const toRef = (parsed: NonNullable<ReturnType<typeof parseAny>>): string =>
+            parsed.kind === "message"
+                ? formatMessageRef(parsed.index, format)
+                : formatBlockRef(parsed.blockId, format)
+
+        const start = parseAny(parts[0]!)
+        const end = parts.length === 2 ? parseAny(parts[1]!) : start
+        if (start && end) {
+            if (start.kind === "message" && end.kind === "message" && start.index > end.index) {
+                throw new Error(`Invalid selector ${selector}: start comes after end`)
+            }
+            const startRef = toRef(start)
+            const endRef = toRef(end)
+            return {
+                selector: startRef === endRef ? startRef : `${startRef}-${endRef}`,
+                startRef,
+                endRef,
+            }
+        }
+    }
+    throw new Error(
+        `Invalid selector ${JSON.stringify(selector)}. Use an injected message ID (e.g. ${
+            format === "compact" ? "" : "m0004"
+        }), a block ID (e.g. ${
+            format === "compact" ? "" : "b2"
+        }), or an inclusive range of either.`,
+    )
+}
+
+interface RcPlan {
+    selector: ParsedSelector
+    messageIds: string[]
+    messages: WithParts[]
+    selection: SelectionResolution
+    anchorMessageId: string
+    priorSummaries: PriorSummary[]
+}
+
+function planSelections(
+    ctx: ToolContext,
+    rawMessages: WithParts[],
+    selectors: ParsedSelector[],
+): RcPlan[] {
+    const context = buildSearchContext(ctx.state, rawMessages)
+
+    return selectors.map((selector) => {
+        const { startReference, endReference } = resolveBoundaryIds(
+            context,
+            ctx.state,
+            ctx.config,
+            selector.startRef,
+            selector.endRef,
+        )
+        const selection = resolveSelection(context, ctx.config, startReference, endReference)
+        const messages = selection.messageIds
+            .map((messageId) => context.rawMessagesById.get(messageId))
+            .filter((message): message is WithParts => Boolean(message))
+        const anchorMessageId = resolveAnchorMessageId(startReference)
+
+        const priorSummaries: PriorSummary[] = []
+        for (const blockId of selection.requiredBlockIds) {
+            const block = ctx.state.prune.messages.blocksById.get(blockId)
+            if (!block) {
+                continue
+            }
+            const body = restoreSummary(block.summary)
+            if (body.length > 0) {
+                priorSummaries.push({
+                    ref: formatBlockRef(blockId, ctx.state.idFormat),
+                    body,
+                })
+            }
+        }
+
+        return {
+            selector,
+            messageIds: selection.messageIds,
+            messages,
+            selection,
+            anchorMessageId,
+            priorSummaries,
+        }
+    })
+}
+
+function validateNonOverlapping(plans: RcPlan[]): void {
+    const sorted = [...plans].sort(
+        (left, right) =>
+            left.selection.startReference.rawIndex - right.selection.startReference.rawIndex ||
+            left.selection.endReference.rawIndex - right.selection.endReference.rawIndex,
+    )
+
+    for (let index = 1; index < sorted.length; index++) {
+        const previous = sorted[index - 1]
+        const current = sorted[index]
+        if (!previous || !current) {
+            continue
+        }
+        if (current.selection.startReference.rawIndex <= previous.selection.endReference.rawIndex) {
+            throw new Error(
+                `Selectors ${previous.selector.selector} and ${current.selector.selector} overlap. Overlapping selections cannot be compressed in one call.`,
+            )
+        }
+    }
+}
+
+function assertBlockCoverage(
+    ctx: ToolContext,
+    selector: string,
+    selection: SelectionResolution,
+): void {
+    const required = new Set(selection.requiredBlockIds)
+    for (const messageId of selection.messageIds) {
+        const entry = ctx.state.prune.messages.byMessageId.get(messageId)
+        if (!entry || entry.activeBlockIds.length === 0) {
+            continue
+        }
+        for (const blockId of entry.activeBlockIds) {
+            if (required.has(blockId)) {
+                continue
+            }
+            const ref = formatBlockRef(blockId, ctx.state.idFormat)
+            throw new Error(
+                `Selector ${selector} only partially covers compressed block ${ref}. Include ${ref} as a boundary so its summary can be merged, or exclude its messages.`,
+            )
+        }
+    }
+}
+
+export function createRcCompressTool(ctx: ToolContext): ReturnType<typeof tool> {
+    ctx.prompts.reload()
+    const runtimePrompts = ctx.prompts.getRuntimePrompts()
+    return tool({
+        description: runtimePrompts.compressRc + rcFormat(ctx.state.idFormat),
+        args: {
+            ids: tool.schema
+                .array(tool.schema.string())
+                .describe(
+                    "The injected message or block IDs, or inclusive ranges, to compress. Each item becomes one summary.",
+                ),
+        },
+        async execute(args, toolCtx) {
+            const input = (args ?? {}) as RcCompressToolArgs
+            validateArgs(input)
+
+            const sessionID =
+                typeof (toolCtx as any)?.sessionID === "string" ? (toolCtx as any).sessionID : ""
+            if (!sessionID) {
+                throw new Error("compress could not determine the session")
+            }
+            const callId =
+                typeof (toolCtx as any)?.callID === "string" ? (toolCtx as any).callID : undefined
+            const toolMessageID = String((toolCtx as any)?.messageID ?? "")
+
+            try {
+                const { rawMessages } = await prepareSession(
+                    ctx,
+                    toolCtx as any,
+                    `Compress: ${input.ids.join(", ")}`,
+                )
+
+                const selectors = input.ids.map((selector) =>
+                    parseSelector(selector, ctx.state.idFormat),
+                )
+                const plans = planSelections(ctx, rawMessages, selectors)
+                validateNonOverlapping(plans)
+                for (const plan of plans) {
+                    assertBlockCoverage(ctx, plan.selector.selector, plan.selection)
+                }
+
+                const segments: SummarySegment[] = plans.map((plan) => ({
+                    selector: plan.selector.selector,
+                    messages: plan.messages,
+                    priorSummaries: plan.priorSummaries,
+                }))
+
+                const { text, usage } = await generateSummaries(ctx, sessionID, segments)
+                ctx.state.stats.compressionUsage = addCompressionUsageTotals(
+                    ctx.state.stats.compressionUsage,
+                    usage,
+                )
+                await saveSessionState(ctx.state, ctx.logger)
+                const accepted = await resolveAcceptedSections(
+                    ctx,
+                    plans.map((plan) => plan.selector.selector),
+                    text,
+                )
+                if (accepted.length === 0) {
+                    throw new Error("summary reply contained no usable sections")
+                }
+
+                const { entries, applied, missing } = await applySummaries(
+                    ctx,
+                    sessionID,
+                    toolMessageID,
+                    callId,
+                    plans,
+                    accepted,
+                )
+                if (entries.length === 0) {
+                    throw new Error("summary reply could not be applied to the session")
+                }
+
+                await finalizeSession(
+                    ctx,
+                    toolCtx as any,
+                    rawMessages,
+                    entries,
+                    plans.map((plan) => plan.selector.selector).join(", "),
+                )
+
+                if (missing.length === 0) {
+                    return `Compressed: ${applied.join(", ")}.`
+                }
+                return `Compressed: ${applied.join(", ")}. Not compressed: ${missing.join(", ")}.`
+            } catch (error) {
+                await resetManualMode(ctx, sessionID)
+                throw error
+            }
+        },
+    })
+}
+
+async function generateSummaries(
+    ctx: ToolContext,
+    sessionID: string,
+    segments: SummarySegment[],
+): Promise<{ text: string; usage: CompressionUsageTotals }> {
+    let usage = emptyCompressionUsage()
+
+    const summarize = async (promptSegments: SummarySegment[]): Promise<string> => {
+        const callId = nextSummaryCallId()
+        const prompt = buildSummaryPrompt(promptSegments, ctx.state.idFormat, callId)
+        const result = await ctx.generate({ sessionID, prompt, callId })
+        const text = result.text ?? ""
+        // Provider usage is authoritative; fall back to a local estimate when
+        // the plugin could not observe the underlying model call.
+        usage = addCompressionUsage(
+            usage,
+            result.usage ?? estimatedCompressionUsage(countTokens(prompt), countTokens(text)),
+        )
+        return text
+    }
+
+    let text = await summarize(segments)
+
+    // One retry for sections the model failed to emit.
+    for (let retry = 0; retry < 2; retry++) {
+        const parsed = parseSummarySections(
+            text,
+            segments.map((segment) => segment.selector),
+            ctx.state.idFormat,
+        )
+        const retrySelectors = [...parsed.missing, ...parsed.empty]
+        if (retrySelectors.length === 0) {
+            break
+        }
+        const retrySegments = segments.filter((segment) =>
+            retrySelectors.includes(segment.selector.toLowerCase()),
+        )
+        if (retrySegments.length === 0) {
+            break
+        }
+        try {
+            text += "\n\n" + (await summarize(retrySegments))
+        } catch {
+            break
+        }
+    }
+    return { text, usage }
+}
+
+async function resolveAcceptedSections(
+    ctx: ToolContext,
+    selectors: string[],
+    text: string,
+): Promise<SummarySection[]> {
+    const parsed = parseSummarySections(text, selectors, ctx.state.idFormat)
+    ctx.logger.info("rc summary reply parsed", {
+        replyChars: text.length,
+        accepted: parsed.sections.map((section) => section.selector).join(","),
+        missing: parsed.missing.join(","),
+        empty: parsed.empty.join(","),
+    })
+    if (parsed.missing.length > 0 || parsed.empty.length > 0) {
+        ctx.logger.warn("rc compression reply was incomplete", {
+            missing: parsed.missing.join(","),
+            empty: parsed.empty.join(","),
+        })
+    }
+    return parsed.sections
+}
+
+async function applySummaries(
+    ctx: ToolContext,
+    sessionID: string,
+    toolMessageID: string,
+    callId: string | undefined,
+    plans: RcPlan[],
+    sections: SummarySection[],
+): Promise<{ entries: NotificationEntry[]; applied: string[]; missing: string[] }> {
+    // Re-read the session so we only apply summaries whose source messages are
+    // still present (a compaction or revert may have happened meanwhile).
+    const rawMessages = await fetchSessionMessages(ctx.client, sessionID)
+    const context = buildSearchContext(ctx.state, rawMessages)
+    const planBySelector = new Map(plans.map((plan) => [plan.selector.selector, plan]))
+
+    const runId = allocateRunId(ctx.state)
+    const batchTopic = plans.map((plan) => plan.selector.selector).join(", ")
+    const entries: NotificationEntry[] = []
+    const applied: string[] = []
+    const appliedSet = new Set<string>()
+
+    for (const section of sections) {
+        const plan = planBySelector.get(section.selector)
+        if (!plan) {
+            continue
+        }
+
+        try {
+            const { startReference, endReference } = resolveBoundaryIds(
+                context,
+                ctx.state,
+                ctx.config,
+                plan.selector.startRef,
+                plan.selector.endRef,
+            )
+            const selection = resolveSelection(context, ctx.config, startReference, endReference)
+            if (selection.messageIds.length === 0) {
+                continue
+            }
+            assertBlockCoverage(ctx, section.selector, selection)
+
+            const anchorMessageId = resolveAnchorMessageId(startReference)
+            const blockId = allocateBlockId(ctx.state)
+
+            const withPromptInfo = appendProtectedPromptInfo(
+                section.summary,
+                selection,
+                context,
+                ctx.state,
+                ctx.config.compress.protectTags,
+            )
+            const finalSummary = await appendProtectedTools(
+                ctx.client,
+                ctx.state,
+                ctx.config.experimental.allowSubAgents,
+                withPromptInfo,
+                selection,
+                context,
+                ctx.config.compress.protectedTools,
+                ctx.config.protectedFilePatterns,
+            )
+
+            const storedSummary = wrapCompressedSummary(blockId, finalSummary, ctx.state.idFormat)
+            const summaryTokens = countTokens(storedSummary)
+
+            applyCompressionState(
+                ctx.state,
+                {
+                    topic: section.selector,
+                    batchTopic,
+                    startId: plan.selector.startRef,
+                    endId: plan.selector.endRef,
+                    mode: "rc",
+                    runId,
+                    compressMessageId: toolMessageID,
+                    compressCallId: callId,
+                    summaryTokens,
+                },
+                selection,
+                anchorMessageId,
+                blockId,
+                storedSummary,
+                selection.requiredBlockIds,
+            )
+
+            entries.push({
+                blockId,
+                runId,
+                summary: finalSummary,
+                summaryTokens,
+            })
+            applied.push(section.selector)
+            appliedSet.add(section.selector)
+        } catch (error) {
+            ctx.logger.warn("rc compression could not apply a summary", {
+                selector: section.selector,
+                error: error instanceof Error ? error.message : String(error),
+            })
+        }
+    }
+
+    const missing = plans
+        .map((plan) => plan.selector.selector)
+        .filter((selector) => !appliedSet.has(selector))
+    return { entries, applied, missing }
+}
+
+async function resetManualMode(ctx: ToolContext, sessionID: string): Promise<void> {
+    if (ctx.state.manualMode !== "compress-pending") {
+        return
+    }
+    ctx.state.manualMode = false
+    await refreshManualMode(ctx.state, sessionID, ctx.logger, ctx.config.manualMode.enabled)
+    await saveSessionState(ctx.state, ctx.logger)
+}