Skip to content

feat: add pinning mode for context management #231

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
AlbertoCentonze wants to merge 1 commit into
base: master
Choose a base branch
from
Draft

feat: add pinning mode for context management #231

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
40 changes: 33 additions & 7 deletions index.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@ import { getConfig } from "./lib/config"
import { Logger } from "./lib/logger"
import { loadPrompt } from "./lib/prompts"
import { createSessionState } from "./lib/state"
import { createDiscardTool, createExtractTool } from "./lib/strategies"
import { createDiscardTool, createExtractTool, createPinTool } from "./lib/strategies"
import { createChatMessageTransformHandler } from "./lib/hooks"

const plugin: Plugin = (async (ctx) => {
Expand Down Expand Up @@ -43,9 +43,16 @@ const plugin: Plugin = (async (ctx) => {

const discardEnabled = config.tools.discard.enabled
const extractEnabled = config.tools.extract.enabled
const pinEnabled = config.tools.pin.enabled
const pinningModeEnabled = config.tools.pinningMode.enabled

let promptName: string
if (discardEnabled && extractEnabled) {
if (pinningModeEnabled || pinEnabled) {
// Pinning mode: use pin prompt (with optional extract)
promptName = extractEnabled
? "user/system/system-prompt-pin-extract"
: "user/system/system-prompt-pin"
} else if (discardEnabled && extractEnabled) {
promptName = "user/system/system-prompt-both"
} else if (discardEnabled) {
promptName = "user/system/system-prompt-discard"
Expand All @@ -63,6 +70,7 @@ const plugin: Plugin = (async (ctx) => {
state,
logger,
config,
ctx.directory,
),
"chat.message": async (
input: {
Expand All @@ -80,17 +88,30 @@ const plugin: Plugin = (async (ctx) => {
logger.debug("Cached variant from chat.message hook", { variant: input.variant })
},
tool: {
...(config.tools.discard.enabled && {
discard: createDiscardTool({
// Discard tool only available in non-pinning mode
...(!config.tools.pinningMode.enabled &&
config.tools.discard.enabled && {
discard: createDiscardTool({
client: ctx.client,
state,
logger,
config,
workingDirectory: ctx.directory,
}),
}),
// Extract tool available in both modes
...(config.tools.extract.enabled && {
extract: createExtractTool({
client: ctx.client,
state,
logger,
config,
workingDirectory: ctx.directory,
}),
}),
...(config.tools.extract.enabled && {
extract: createExtractTool({
// Pin tool only available in pinning mode
...(config.tools.pin.enabled && {
pin: createPinTool({
client: ctx.client,
state,
logger,
Expand All @@ -103,7 +124,12 @@ const plugin: Plugin = (async (ctx) => {
// Add enabled tools to primary_tools by mutating the opencode config
// This works because config is cached and passed by reference
const toolsToAdd: string[] = []
if (config.tools.discard.enabled) toolsToAdd.push("discard")
// In pinning mode, add pin instead of discard
if (config.tools.pinningMode.enabled) {
if (config.tools.pin.enabled) toolsToAdd.push("pin")
} else {
if (config.tools.discard.enabled) toolsToAdd.push("discard")
}
if (config.tools.extract.enabled) toolsToAdd.push("extract")

if (toolsToAdd.length > 0) {
Expand Down
105 changes: 105 additions & 0 deletions lib/config.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,17 @@ export interface ExtractTool {
showDistillation: boolean
}

export interface PinTool {
enabled: boolean
}

export interface PinningMode {
enabled: boolean
pruneFrequency: number // Auto-prune every N turns
pinDuration: number // Pins expire after M turns
warningTurns: number // Warn N turns before auto-prune
}

export interface ToolSettings {
nudgeEnabled: boolean
nudgeFrequency: number
Expand All @@ -28,6 +39,8 @@ export interface Tools {
settings: ToolSettings
discard: DiscardTool
extract: ExtractTool
pin: PinTool
pinningMode: PinningMode
}

export interface SupersedeWrites {
Expand Down Expand Up @@ -89,6 +102,13 @@ export const VALID_CONFIG_KEYS = new Set([
"tools.extract",
"tools.extract.enabled",
"tools.extract.showDistillation",
"tools.pin",
"tools.pin.enabled",
"tools.pinningMode",
"tools.pinningMode.enabled",
"tools.pinningMode.pruneFrequency",
"tools.pinningMode.pinDuration",
"tools.pinningMode.warningTurns",
"strategies",
// strategies.deduplication
"strategies.deduplication",
Expand Down Expand Up @@ -238,6 +258,57 @@ function validateConfigTypes(config: Record<string, any>): ValidationError[] {
})
}
}
if (tools.pin) {
if (tools.pin.enabled !== undefined && typeof tools.pin.enabled !== "boolean") {
errors.push({
key: "tools.pin.enabled",
expected: "boolean",
actual: typeof tools.pin.enabled,
})
}
}
if (tools.pinningMode) {
if (
tools.pinningMode.enabled !== undefined &&
typeof tools.pinningMode.enabled !== "boolean"
) {
errors.push({
key: "tools.pinningMode.enabled",
expected: "boolean",
actual: typeof tools.pinningMode.enabled,
})
}
if (
tools.pinningMode.pruneFrequency !== undefined &&
typeof tools.pinningMode.pruneFrequency !== "number"
) {
errors.push({
key: "tools.pinningMode.pruneFrequency",
expected: "number",
actual: typeof tools.pinningMode.pruneFrequency,
})
}
if (
tools.pinningMode.pinDuration !== undefined &&
typeof tools.pinningMode.pinDuration !== "number"
) {
errors.push({
key: "tools.pinningMode.pinDuration",
expected: "number",
actual: typeof tools.pinningMode.pinDuration,
})
}
if (
tools.pinningMode.warningTurns !== undefined &&
typeof tools.pinningMode.warningTurns !== "number"
) {
errors.push({
key: "tools.pinningMode.warningTurns",
expected: "number",
actual: typeof tools.pinningMode.warningTurns,
})
}
}
}

// Strategies validators
Expand Down Expand Up @@ -384,6 +455,15 @@ const defaultConfig: PluginConfig = {
enabled: true,
showDistillation: false,
},
pin: {
enabled: false,
},
pinningMode: {
enabled: false,
pruneFrequency: 10,
pinDuration: 8,
warningTurns: 2,
},
},
strategies: {
deduplication: {
Expand Down Expand Up @@ -499,6 +579,20 @@ function createDefaultConfig(): void {
"enabled": true,
// Show distillation content as an ignored message notification
"showDistillation": false
},
// Pin tool to protect context from auto-pruning (used with pinningMode)
"pin": {
"enabled": false
},
// Pinning mode: auto-prune everything not pinned every N turns
"pinningMode": {
"enabled": false,
// Auto-prune every N turns
Copy link
Collaborator

@Tarquinen Tarquinen Jan 10, 2026 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

this should match the pattern set in tools.settings.nudgeFrequency, something like
auto-prune everything not pinned every <pruneFrequency> turns
same for other occurances below

"pruneFrequency": 10,
// Pins expire after M turns
"pinDuration": 8,
// Warn N turns before auto-prune
"warningTurns": 2
}
},
// Automatic pruning strategies
Expand Down Expand Up @@ -608,6 +702,15 @@ function mergeTools(
enabled: override.extract?.enabled ?? base.extract.enabled,
showDistillation: override.extract?.showDistillation ?? base.extract.showDistillation,
},
pin: {
enabled: override.pin?.enabled ?? base.pin.enabled,
},
pinningMode: {
enabled: override.pinningMode?.enabled ?? base.pinningMode.enabled,
pruneFrequency: override.pinningMode?.pruneFrequency ?? base.pinningMode.pruneFrequency,
pinDuration: override.pinningMode?.pinDuration ?? base.pinningMode.pinDuration,
warningTurns: override.pinningMode?.warningTurns ?? base.pinningMode.warningTurns,
},
}
}

Expand All @@ -622,6 +725,8 @@ function deepCloneConfig(config: PluginConfig): PluginConfig {
},
discard: { ...config.tools.discard },
extract: { ...config.tools.extract },
pin: { ...config.tools.pin },
pinningMode: { ...config.tools.pinningMode },
},
strategies: {
deduplication: {
Expand Down
8 changes: 7 additions & 1 deletion lib/hooks.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@ import type { SessionState, WithParts } from "./state"
import type { Logger } from "./logger"
import type { PluginConfig } from "./config"
import { syncToolCache } from "./state/tool-cache"
import { deduplicate, supersedeWrites, purgeErrors } from "./strategies"
import { deduplicate, supersedeWrites, purgeErrors, autoPrune } from "./strategies"
import { prune, insertPruneToolContext } from "./messages"
import { checkSession } from "./state"

Expand All @@ -11,6 +11,7 @@ export function createChatMessageTransformHandler(
state: SessionState,
logger: Logger,
config: PluginConfig,
workingDirectory: string,
) {
return async (input: {}, output: { messages: WithParts[] }) => {
await checkSession(client, state, logger, output.messages)
Expand All @@ -25,6 +26,11 @@ export function createChatMessageTransformHandler(
supersedeWrites(state, logger, config, output.messages)
purgeErrors(state, logger, config, output.messages)

// Run auto-prune if pinning mode is enabled
if (config.tools.pinningMode.enabled) {
Copy link
Collaborator

@Tarquinen Tarquinen Jan 10, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

this "if" should be handled within autoPrune

await autoPrune(client, state, logger, config, output.messages, workingDirectory)
}

prune(state, logger, config, output.messages)

insertPruneToolContext(state, config, logger, output.messages)
Expand Down
48 changes: 44 additions & 4 deletions lib/messages/inject.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@ import type { UserMessage } from "@opencode-ai/sdk/v2"
import { loadPrompt } from "../prompts"
import { extractParameterKey, buildToolIdList, createSyntheticUserMessage } from "./utils"
import { getLastUserMessage } from "../shared-utils"
import { turnsUntilAutoPrune, shouldShowAutoPruneWarning } from "../strategies/auto-prune"

const getNudgeString = (config: PluginConfig): string => {
const discardEnabled = config.tools.discard.enabled
Expand All @@ -28,9 +29,12 @@ ${content}
const getCooldownMessage = (config: PluginConfig): string => {
const discardEnabled = config.tools.discard.enabled
const extractEnabled = config.tools.extract.enabled
const pinEnabled = config.tools.pin.enabled

let toolName: string
if (discardEnabled && extractEnabled) {
if (pinEnabled) {
toolName = extractEnabled ? "pin or extract tools" : "pin tool"
Copy link
Collaborator

@Tarquinen Tarquinen Jan 10, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

this should be something like "pin, discard or extract tools" and should also handle if only one of discard / extract is enabled, so if pin and extract are enabled but discard disabled, "pin or extract tools".

} else if (discardEnabled && extractEnabled) {
toolName = "discard or extract tools"
} else if (discardEnabled) {
toolName = "discard tool"
Expand All @@ -43,6 +47,16 @@ Context management was just performed. Do not use the ${toolName} again. A fresh
</prunable-tools>`
}

const getAutoPruneWarningMessage = (state: SessionState, config: PluginConfig): string => {
const turns = turnsUntilAutoPrune(state, config)
const pinnedCount = state.pins.size

return `<auto-prune-warning>
Auto-prune in ${turns} turn(s). All unpinned tool outputs will be discarded.
Currently ${pinnedCount} tool(s) pinned. Use the \`pin\` tool NOW to preserve any context you need.
Copy link
Collaborator

@Tarquinen Tarquinen Jan 10, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Might need to tell the model what is pinned or what isn't, I don't think telling the model how many tools are pinned is useful information

</auto-prune-warning>`
}

const buildPrunableToolsList = (
state: SessionState,
config: PluginConfig,
Expand All @@ -51,6 +65,7 @@ const buildPrunableToolsList = (
): string => {
const lines: string[] = []
const toolIdList: string[] = buildToolIdList(state, messages, logger)
const pinningModeEnabled = config.tools.pinningMode.enabled

state.toolParameters.forEach((toolParameterEntry, toolCallId) => {
if (state.prune.toolIds.includes(toolCallId)) {
Expand All @@ -74,7 +89,18 @@ const buildPrunableToolsList = (
const description = paramKey
? `${toolParameterEntry.tool}, ${paramKey}`
: toolParameterEntry.tool
lines.push(`${numericId}: ${description}`)

// Show pin status in pinning mode
let line = `${numericId}: ${description}`
if (pinningModeEnabled) {
const pin = state.pins.get(toolCallId)
if (pin) {
const turnsRemaining = pin.expiresAtTurn - state.currentTurn
line += ` [PINNED, expires in ${turnsRemaining} turn(s)]`
}
}

lines.push(line)
logger.debug(
`Prunable tool found - ID: ${numericId}, Tool: ${toolParameterEntry.tool}, Call ID: ${toolCallId}`,
)
Expand All @@ -93,7 +119,12 @@ export const insertPruneToolContext = (
logger: Logger,
messages: WithParts[],
): void => {
if (!config.tools.discard.enabled && !config.tools.extract.enabled) {
const pinEnabled = config.tools.pin.enabled
const discardEnabled = config.tools.discard.enabled
const extractEnabled = config.tools.extract.enabled

// Need at least one pruning tool enabled
if (!pinEnabled && !discardEnabled && !extractEnabled) {
return
}

Expand All @@ -111,15 +142,24 @@ export const insertPruneToolContext = (
logger.debug("prunable-tools: \n" + prunableToolsList)

let nudgeString = ""
// Only show nudge in non-pinning mode
if (
!config.tools.pinningMode.enabled &&
config.tools.settings.nudgeEnabled &&
state.nudgeCounter >= config.tools.settings.nudgeFrequency
) {
logger.info("Inserting prune nudge message")
nudgeString = "\n" + getNudgeString(config)
}

prunableToolsContent = prunableToolsList + nudgeString
// Add auto-prune warning if approaching prune cycle
let warningString = ""
if (shouldShowAutoPruneWarning(state, config)) {
logger.info("Inserting auto-prune warning message")
warningString = "\n" + getAutoPruneWarningMessage(state, config)
}

prunableToolsContent = prunableToolsList + nudgeString + warningString
}

const lastUserMessage = getLastUserMessage(messages)
Expand Down
Loading