Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
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
5 changes: 3 additions & 2 deletions index.ts
Original file line number Diff line number Diff line change
Expand Up @@ -2,8 +2,8 @@ import type { Plugin } from "@opencode-ai/plugin"
import { getConfig } from "./lib/config"
import { Logger } from "./lib/logger"
import { createSessionState } from "./lib/state"
import { createPruneTool } from "./lib/strategies/prune-tool"
import { createChatMessageTransformHandler } from "./lib/hooks"
import { createPruneTool } from "./lib/strategies"
import { createChatMessageTransformHandler, createEventHandler } from "./lib/hooks"

const plugin: Plugin = (async (ctx) => {
const config = getConfig(ctx)
Expand Down Expand Up @@ -54,6 +54,7 @@ const plugin: Plugin = (async (ctx) => {
logger.info("Added 'prune' to experimental.primary_tools via config mutation")
}
},
event: createEventHandler(ctx.client, config, state, logger, ctx.directory),
}
}) satisfies Plugin

Expand Down
42 changes: 40 additions & 2 deletions lib/hooks.ts
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@ import { syncToolCache } from "./state/tool-cache"
import { deduplicate } from "./strategies"
import { prune, insertPruneToolContext } from "./messages"
import { checkSession } from "./state"
import { runOnIdle } from "./strategies/on-idle"


export function createChatMessageTransformHandler(
Expand All @@ -24,11 +25,48 @@ export function createChatMessageTransformHandler(

syncToolCache(state, config, logger, output.messages);


deduplicate(client, state, logger, config, output.messages)
deduplicate(state, logger, config, output.messages)

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

insertPruneToolContext(state, config, logger, output.messages)
}
}

export function createEventHandler(
client: any,
config: PluginConfig,
state: SessionState,
logger: Logger,
workingDirectory?: string
) {
return async (
{ event }: { event: any }
) => {
if (state.sessionId === null || state.isSubAgent) {
return
}

if (event.type === "session.status" && event.properties.status.type === "idle") {
if (!config.strategies.onIdle.enabled) {
return
}
if (state.lastToolPrune) {
logger.info("Skipping OnIdle pruning - last tool was prune")
return
}

try {
await runOnIdle(
client,
state,
logger,
config,
workingDirectory
)
} catch (err: any) {
logger.error("OnIdle pruning failed", { error: err.message })
}
}
}
}
6 changes: 5 additions & 1 deletion lib/messages/prune.ts
Original file line number Diff line number Diff line change
Expand Up @@ -40,6 +40,10 @@ export const insertPruneToolContext = (
logger: Logger,
messages: WithParts[]
): void => {
if (!config.strategies.pruneTool.enabled) {
return
}

const lastUserMessage = getLastUserMessage(messages)
if (!lastUserMessage || lastUserMessage.info.role !== 'user') {
return
Expand All @@ -48,7 +52,7 @@ export const insertPruneToolContext = (
const prunableToolsList = buildPrunableToolsList(state, config, logger, messages)

let nudgeString = ""
if (config.strategies.pruneTool.nudge.enabled && state.nudgeCounter >= config.strategies.pruneTool.nudge.frequency) {
if (state.nudgeCounter >= config.strategies.pruneTool.nudge.frequency) {
logger.info("Inserting prune nudge message")
nudgeString = "\n" + NUDGE_STRING
}
Expand Down
4 changes: 3 additions & 1 deletion lib/state/state.ts
Original file line number Diff line number Diff line change
Expand Up @@ -43,7 +43,8 @@ export function createSessionState(): SessionState {
totalPruneTokens: 0,
},
toolParameters: new Map<string, ToolParameterEntry>(),
nudgeCounter: 0
nudgeCounter: 0,
lastToolPrune: false
}
}

Expand All @@ -59,6 +60,7 @@ export function resetSessionState(state: SessionState): void {
}
state.toolParameters.clear()
state.nudgeCounter = 0
state.lastToolPrune = false
}

export async function ensureSessionInitialized(
Expand Down
4 changes: 4 additions & 0 deletions lib/state/tool-cache.ts
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@ export async function syncToolCache(
): Promise<void> {
try {
logger.info("Syncing tool parameters from OpenCode messages")

for (const msg of messages) {
for (const part of msg.parts) {
if (part.type !== "tool" || !part.callID || state.toolParameters.has(part.callID)) {
Expand All @@ -36,6 +37,9 @@ export async function syncToolCache(
if (!config.strategies.pruneTool.protectedTools.includes(part.tool)) {
state.nudgeCounter++
}

state.lastToolPrune = part.tool === "prune"
logger.info("lastToolPrune=" + String(state.lastToolPrune))
}
}

Expand Down
1 change: 1 addition & 0 deletions lib/state/types.ts
Original file line number Diff line number Diff line change
Expand Up @@ -30,4 +30,5 @@ export interface SessionState {
stats: SessionStats
toolParameters: Map<string, ToolParameterEntry>
nudgeCounter: number
lastToolPrune: boolean
}
1 change: 0 additions & 1 deletion lib/strategies/deduplication.ts
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,6 @@ import { calculateTokensSaved } from "../utils"
* Modifies the session state in place to add pruned tool call IDs.
*/
export const deduplicate = (
client: any,
state: SessionState,
logger: Logger,
config: PluginConfig,
Expand Down
3 changes: 2 additions & 1 deletion lib/strategies/index.ts
Original file line number Diff line number Diff line change
@@ -1,2 +1,3 @@
export { deduplicate } from "./deduplication"

export { runOnIdle } from "./on-idle"
export { createPruneTool } from "./prune-tool"
Loading