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

package/lib/messages/inject/inject.ts

@@ -0,0 +1,215 @@
+import type { SessionState, WithParts } from "../../state"
+import type { Logger } from "../../logger"
+import type { PluginConfig } from "../../config"
+import type { RuntimePrompts } from "../../prompts/store"
+import { formatMessageIdTag } from "../../message-ids"
+import { compressPermission } from "../../compress-permission"
+import {
+    getLastUserMessage,
+    isIgnoredUserMessage,
+    isProtectedUserMessage,
+    messageHasCompress,
+} from "../query"
+import { saveSessionState } from "../../state/persistence"
+import {
+    appendToTextPart,
+    appendToLastTextPart,
+    appendToAllToolParts,
+    createSyntheticTextPart,
+    hasContent,
+} from "../utils"
+import {
+    addAnchor,
+    applyAnchoredNudges,
+    countMessagesAfterIndex,
+    findLastNonIgnoredMessage,
+    getIterationNudgeThreshold,
+    getNudgeFrequency,
+    getModelInfo,
+    isContextOverLimits,
+} from "./utils"
+
+export const injectCompressNudges = (
+    state: SessionState,
+    config: PluginConfig,
+    logger: Logger,
+    messages: WithParts[],
+    prompts: RuntimePrompts,
+    createAnchors = true,
+): void => {
+    if (compressPermission(state, config) === "deny") {
+        return
+    }
+
+    if (state.manualMode) {
+        return
+    }
+
+    if (!createAnchors) {
+        applyAnchoredNudges(state, config, messages, prompts)
+        return
+    }
+
+    const lastMessage = findLastNonIgnoredMessage(messages)
+    const lastAssistantMessage = messages.findLast((message) => message.info.role === "assistant")
+
+    if (lastAssistantMessage && messageHasCompress(lastAssistantMessage)) {
+        state.nudges.contextLimitAnchors.clear()
+        state.nudges.turnNudgeAnchors.clear()
+        state.nudges.iterationNudgeAnchors.clear()
+        void saveSessionState(state, logger)
+        return
+    }
+
+    const { providerId, modelId } = getModelInfo(messages)
+    let anchorsChanged = false
+
+    const { overMaxLimit, overMinLimit } = isContextOverLimits(
+        config,
+        state,
+        providerId,
+        modelId,
+        messages,
+    )
+
+    if (!overMinLimit) {
+        const hadTurnAnchors = state.nudges.turnNudgeAnchors.size > 0
+        const hadIterationAnchors = state.nudges.iterationNudgeAnchors.size > 0
+
+        if (hadTurnAnchors || hadIterationAnchors) {
+            state.nudges.turnNudgeAnchors.clear()
+            state.nudges.iterationNudgeAnchors.clear()
+            anchorsChanged = true
+        }
+    }
+
+    if (overMaxLimit) {
+        if (lastMessage) {
+            const interval = getNudgeFrequency(config)
+            const added = addAnchor(
+                state.nudges.contextLimitAnchors,
+                lastMessage.message.info.id,
+                lastMessage.index,
+                messages,
+                interval,
+            )
+            if (added) {
+                anchorsChanged = true
+            }
+        }
+    } else if (overMinLimit) {
+        const isLastMessageUser = lastMessage?.message.info.role === "user"
+
+        if (isLastMessageUser && lastAssistantMessage) {
+            const previousSize = state.nudges.turnNudgeAnchors.size
+            state.nudges.turnNudgeAnchors.add(lastMessage.message.info.id)
+            state.nudges.turnNudgeAnchors.add(lastAssistantMessage.info.id)
+            if (state.nudges.turnNudgeAnchors.size !== previousSize) {
+                anchorsChanged = true
+            }
+        }
+
+        const lastUserMessage = getLastUserMessage(messages)
+        if (lastUserMessage && lastMessage) {
+            const lastUserMessageIndex = messages.findIndex(
+                (message) => message.info.id === lastUserMessage.info.id,
+            )
+            if (lastUserMessageIndex >= 0) {
+                const messagesSinceUser = countMessagesAfterIndex(messages, lastUserMessageIndex)
+                const iterationThreshold = getIterationNudgeThreshold(config)
+
+                if (
+                    lastMessage.index > lastUserMessageIndex &&
+                    messagesSinceUser >= iterationThreshold
+                ) {
+                    const interval = getNudgeFrequency(config)
+                    const added = addAnchor(
+                        state.nudges.iterationNudgeAnchors,
+                        lastMessage.message.info.id,
+                        lastMessage.index,
+                        messages,
+                        interval,
+                    )
+
+                    if (added) {
+                        anchorsChanged = true
+                    }
+                }
+            }
+        }
+    }
+
+    applyAnchoredNudges(state, config, messages, prompts)
+
+    if (anchorsChanged) {
+        void saveSessionState(state, logger)
+    }
+}
+
+export const injectMessageIds = (
+    state: SessionState,
+    config: PluginConfig,
+    messages: WithParts[],
+): void => {
+    if (compressPermission(state, config) === "deny") {
+        return
+    }
+
+    for (const message of messages) {
+        if (isIgnoredUserMessage(message)) {
+            continue
+        }
+
+        const messageRef = state.messageIds.byRawId.get(message.info.id)
+        if (!messageRef) {
+            continue
+        }
+
+        const isBlockedMessage = isProtectedUserMessage(config, message)
+        const tag = formatMessageIdTag(
+            isBlockedMessage ? "BLOCKED" : messageRef,
+            undefined,
+            state.idFormat,
+        )
+
+        if (message.info.role === "user") {
+            let injected = false
+            for (const part of message.parts) {
+                if (part.type === "text") {
+                    injected = appendToTextPart(part, tag) || injected
+                }
+            }
+
+            if (injected) {
+                continue
+            }
+
+            message.parts.push(createSyntheticTextPart(message, tag))
+            continue
+        }
+
+        if (message.info.role !== "assistant") {
+            continue
+        }
+
+        if (!hasContent(message)) {
+            continue
+        }
+
+        if (appendToAllToolParts(message, tag)) {
+            continue
+        }
+
+        if (appendToLastTextPart(message, tag)) {
+            continue
+        }
+
+        const syntheticPart = createSyntheticTextPart(message, tag)
+        const firstToolIndex = message.parts.findIndex((p) => p.type === "tool")
+        if (firstToolIndex === -1) {
+            message.parts.push(syntheticPart)
+        } else {
+            message.parts.splice(firstToolIndex, 0, syntheticPart)
+        }
+    }
+}

package/lib/messages/inject/subagent-results.ts

@@ -0,0 +1,84 @@
+import type { Logger } from "../../logger"
+import type { SessionState, WithParts } from "../../state"
+import { filterMessages } from "../shape"
+import {
+    buildSubagentResultText,
+    getSubAgentId,
+    mergeSubagentResult,
+} from "../../subagents/subagent-results"
+import { stripHallucinationsFromString } from "../utils"
+
+async function fetchSubAgentMessages(client: any, sessionId: string): Promise<WithParts[]> {
+    const response = await client.session.messages({
+        path: { id: sessionId },
+    })
+
+    return filterMessages(response?.data || response)
+}
+
+export const injectExtendedSubAgentResults = async (
+    client: any,
+    state: SessionState,
+    logger: Logger,
+    messages: WithParts[],
+    allowSubAgents: boolean,
+): Promise<void> => {
+    if (!allowSubAgents) {
+        return
+    }
+
+    for (const message of messages) {
+        const parts = Array.isArray(message.parts) ? message.parts : []
+
+        for (const part of parts) {
+            if (part.type !== "tool" || !["task", "subagent"].includes(part.tool) || !part.callID) {
+                continue
+            }
+            if (state.prune.tools.has(part.callID)) {
+                continue
+            }
+            if (part.state?.status !== "completed" || typeof part.state.output !== "string") {
+                continue
+            }
+
+            const cachedResult = state.subAgentResultCache.get(part.callID)
+            if (cachedResult !== undefined) {
+                if (cachedResult) {
+                    part.state.output = stripHallucinationsFromString(
+                        mergeSubagentResult(part.state.output, cachedResult),
+                        state.idFormat,
+                    )
+                }
+                continue
+            }
+
+            const subAgentSessionId = getSubAgentId(part)
+            if (!subAgentSessionId) {
+                continue
+            }
+
+            let subAgentMessages: WithParts[] = []
+            try {
+                subAgentMessages = await fetchSubAgentMessages(client, subAgentSessionId)
+            } catch (error) {
+                logger.warn("Failed to fetch subagent session for output expansion", {
+                    subAgentSessionId,
+                    callID: part.callID,
+                    error: error instanceof Error ? error.message : String(error),
+                })
+                continue
+            }
+
+            const subAgentResultText = buildSubagentResultText(subAgentMessages)
+            if (!subAgentResultText) {
+                continue
+            }
+
+            state.subAgentResultCache.set(part.callID, subAgentResultText)
+            part.state.output = stripHallucinationsFromString(
+                mergeSubagentResult(part.state.output, subAgentResultText),
+                state.idFormat,
+            )
+        }
+    }
+}

package/lib/messages/inject/utils.ts

@@ -0,0 +1,304 @@
+import type { SessionState, WithParts } from "../../state"
+import type { PluginConfig } from "../../config"
+import {
+    appendGuidanceToDcpTag,
+    buildCompressedBlockGuidance,
+} from "../../prompts/extensions/nudge"
+import type { RuntimePrompts } from "../../prompts/store"
+import type { UserMessage } from "@opencode-ai/sdk/v2"
+import {
+    appendToTextPart,
+    appendToLastTextPart,
+    createSyntheticTextPart,
+    hasContent,
+} from "../utils"
+import { getLastUserMessage, isIgnoredUserMessage } from "../query"
+import { getCurrentTokenUsage } from "../../token-utils"
+import { getActiveSummaryTokenUsage } from "../../state/utils"
+
+export interface LastUserModelContext {
+    providerId: string | undefined
+    modelId: string | undefined
+}
+
+export interface LastNonIgnoredMessage {
+    message: WithParts
+    index: number
+}
+
+export function getNudgeFrequency(config: PluginConfig): number {
+    return Math.max(1, Math.floor(config.compress.nudgeFrequency || 1))
+}
+
+export function getIterationNudgeThreshold(config: PluginConfig): number {
+    return Math.max(1, Math.floor(config.compress.iterationNudgeThreshold || 1))
+}
+
+export function findLastNonIgnoredMessage(messages: WithParts[]): LastNonIgnoredMessage | null {
+    for (let i = messages.length - 1; i >= 0; i--) {
+        const message = messages[i]
+        if (isIgnoredUserMessage(message)) {
+            continue
+        }
+        return { message, index: i }
+    }
+
+    return null
+}
+
+export function countMessagesAfterIndex(messages: WithParts[], index: number): number {
+    let count = 0
+
+    for (let i = index + 1; i < messages.length; i++) {
+        const message = messages[i]
+        if (isIgnoredUserMessage(message)) {
+            continue
+        }
+        count++
+    }
+
+    return count
+}
+
+export function getModelInfo(messages: WithParts[]): LastUserModelContext {
+    const lastUserMessage = getLastUserMessage(messages)
+    if (!lastUserMessage) {
+        return {
+            providerId: undefined,
+            modelId: undefined,
+        }
+    }
+
+    const userInfo = lastUserMessage.info as UserMessage
+    return {
+        providerId: userInfo.model.providerID,
+        modelId: userInfo.model.modelID,
+    }
+}
+
+function resolveContextTokenLimit(
+    config: PluginConfig,
+    state: SessionState,
+    providerId: string | undefined,
+    modelId: string | undefined,
+    threshold: "max" | "min",
+): number | undefined {
+    const parseLimitValue = (limit: number | `${number}%` | undefined): number | undefined => {
+        if (limit === undefined) {
+            return undefined
+        }
+
+        if (typeof limit === "number") {
+            return limit
+        }
+
+        if (!limit.endsWith("%") || state.modelContextLimit === undefined) {
+            return undefined
+        }
+
+        const parsedPercent = parseFloat(limit.slice(0, -1))
+        if (isNaN(parsedPercent)) {
+            return undefined
+        }
+
+        const roundedPercent = Math.round(parsedPercent)
+        const clampedPercent = Math.max(0, Math.min(100, roundedPercent))
+        return Math.round((clampedPercent / 100) * state.modelContextLimit)
+    }
+
+    const modelLimits =
+        threshold === "max" ? config.compress.modelMaxLimits : config.compress.modelMinLimits
+    if (modelLimits && providerId !== undefined && modelId !== undefined) {
+        const providerModelId = `${providerId}/${modelId}`
+        const modelLimit = modelLimits[providerModelId]
+        if (modelLimit !== undefined) {
+            return parseLimitValue(modelLimit)
+        }
+    }
+
+    const globalLimit =
+        threshold === "max" ? config.compress.maxContextLimit : config.compress.minContextLimit
+    return parseLimitValue(globalLimit)
+}
+
+export function isContextOverLimits(
+    config: PluginConfig,
+    state: SessionState,
+    providerId: string | undefined,
+    modelId: string | undefined,
+    messages: WithParts[],
+) {
+    const summaryTokenExtension = config.compress.summaryBuffer
+        ? getActiveSummaryTokenUsage(state)
+        : 0
+    const resolvedMaxContextLimit = resolveContextTokenLimit(
+        config,
+        state,
+        providerId,
+        modelId,
+        "max",
+    )
+    const maxContextLimit =
+        resolvedMaxContextLimit === undefined
+            ? undefined
+            : resolvedMaxContextLimit + summaryTokenExtension
+    const minContextLimit = resolveContextTokenLimit(config, state, providerId, modelId, "min")
+    const currentTokens = getCurrentTokenUsage(state, messages)
+
+    const overMaxLimit = maxContextLimit === undefined ? false : currentTokens > maxContextLimit
+    const overMinLimit = minContextLimit === undefined ? true : currentTokens >= minContextLimit
+
+    return {
+        overMaxLimit,
+        overMinLimit,
+    }
+}
+
+export function addAnchor(
+    anchorMessageIds: Set<string>,
+    anchorMessageId: string,
+    anchorMessageIndex: number,
+    messages: WithParts[],
+    interval: number,
+): boolean {
+    if (anchorMessageIndex < 0) {
+        return false
+    }
+
+    let latestAnchorMessageIndex = -1
+    for (let i = messages.length - 1; i >= 0; i--) {
+        if (anchorMessageIds.has(messages[i].info.id)) {
+            latestAnchorMessageIndex = i
+            break
+        }
+    }
+
+    const shouldAdd =
+        latestAnchorMessageIndex < 0 || anchorMessageIndex - latestAnchorMessageIndex >= interval
+    if (!shouldAdd) {
+        return false
+    }
+
+    const previousSize = anchorMessageIds.size
+    anchorMessageIds.add(anchorMessageId)
+    return anchorMessageIds.size !== previousSize
+}
+
+function injectAnchoredNudge(message: WithParts, nudgeText: string): void {
+    if (!nudgeText.trim()) {
+        return
+    }
+
+    if (message.info.role === "user") {
+        if (appendToLastTextPart(message, nudgeText)) {
+            return
+        }
+
+        message.parts.push(createSyntheticTextPart(message, nudgeText))
+        return
+    }
+
+    if (message.info.role !== "assistant") {
+        return
+    }
+
+    if (!hasContent(message)) {
+        return
+    }
+
+    for (const part of message.parts) {
+        if (part.type === "text") {
+            if (appendToTextPart(part, nudgeText)) {
+                return
+            }
+        }
+    }
+
+    const syntheticPart = createSyntheticTextPart(message, nudgeText)
+    const firstToolIndex = message.parts.findIndex((p) => p.type === "tool")
+    if (firstToolIndex === -1) {
+        message.parts.push(syntheticPart)
+    } else {
+        message.parts.splice(firstToolIndex, 0, syntheticPart)
+    }
+}
+
+function collectAnchoredMessages(
+    anchorMessageIds: Set<string>,
+    messages: WithParts[],
+): Array<{ message: WithParts; index: number }> {
+    const anchoredMessages: Array<{ message: WithParts; index: number }> = []
+
+    for (const anchorMessageId of anchorMessageIds) {
+        const index = messages.findIndex((message) => message.info.id === anchorMessageId)
+        if (index === -1) {
+            continue
+        }
+
+        anchoredMessages.push({
+            message: messages[index],
+            index,
+        })
+    }
+
+    return anchoredMessages
+}
+
+function collectTurnNudgeAnchors(
+    state: SessionState,
+    config: PluginConfig,
+    messages: WithParts[],
+): Set<string> {
+    const turnNudgeAnchors = new Set<string>()
+    const targetRole = config.compress.nudgeForce === "strong" ? "user" : "assistant"
+
+    for (const message of messages) {
+        if (!state.nudges.turnNudgeAnchors.has(message.info.id)) continue
+
+        if (message.info.role === targetRole) {
+            turnNudgeAnchors.add(message.info.id)
+        }
+    }
+
+    return turnNudgeAnchors
+}
+
+function applyGuidedAnchoredNudge(
+    anchorMessageIds: Set<string>,
+    messages: WithParts[],
+    baseNudgeText: string,
+    compressedBlockGuidance: string,
+): void {
+    const nudgeText = appendGuidanceToDcpTag(baseNudgeText, compressedBlockGuidance)
+    if (!nudgeText.trim()) {
+        return
+    }
+
+    for (const { message } of collectAnchoredMessages(anchorMessageIds, messages)) {
+        injectAnchoredNudge(message, nudgeText)
+    }
+}
+
+export function applyAnchoredNudges(
+    state: SessionState,
+    config: PluginConfig,
+    messages: WithParts[],
+    prompts: RuntimePrompts,
+): void {
+    const turnNudgeAnchors = collectTurnNudgeAnchors(state, config, messages)
+    const compressedBlockGuidance = buildCompressedBlockGuidance(state)
+
+    applyGuidedAnchoredNudge(
+        state.nudges.contextLimitAnchors,
+        messages,
+        prompts.contextLimitNudge,
+        compressedBlockGuidance,
+    )
+    applyGuidedAnchoredNudge(turnNudgeAnchors, messages, prompts.turnNudge, compressedBlockGuidance)
+    applyGuidedAnchoredNudge(
+        state.nudges.iterationNudgeAnchors,
+        messages,
+        prompts.iterationNudge,
+        compressedBlockGuidance,
+    )
+}