Skip to content

Commit c5d9287

Browse files
authored
Merge pull request #153 from Opencode-DCP/fix-prune-after-compact-bug
fix
2 parents 53190b6 + 7f50ce1 commit c5d9287

File tree

12 files changed

+79
-28
lines changed

12 files changed

+79
-28
lines changed

lib/hooks.ts

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -48,6 +48,11 @@ export function createEventHandler(
4848
return
4949
}
5050

51+
if (event.type === "session.compacted") {
52+
logger.info("Session compaction detected - updating state")
53+
state.lastCompaction = Date.now()
54+
}
55+
5156
if (event.type === "session.status" && event.properties.status.type === "idle") {
5257
if (!config.strategies.onIdle.enabled) {
5358
return

lib/messages/prune.ts

Lines changed: 6 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@ import type { Logger } from "../logger"
33
import type { PluginConfig } from "../config"
44
import { loadPrompt } from "../prompt"
55
import { extractParameterKey, buildToolIdList } from "./utils"
6-
import { getLastUserMessage } from "../shared-utils"
6+
import { getLastUserMessage, isMessageCompacted } from "../shared-utils"
77
import { UserMessage } from "@opencode-ai/sdk"
88

99
const PRUNED_TOOL_INPUT_REPLACEMENT = '[Input removed to save context]'
@@ -17,7 +17,7 @@ const buildPrunableToolsList = (
1717
messages: WithParts[],
1818
): string => {
1919
const lines: string[] = []
20-
const toolIdList: string[] = buildToolIdList(messages)
20+
const toolIdList: string[] = buildToolIdList(state, messages, logger)
2121

2222
state.toolParameters.forEach((toolParameterEntry, toolCallId) => {
2323
if (state.prune.toolIds.includes(toolCallId)) {
@@ -26,9 +26,6 @@ const buildPrunableToolsList = (
2626
if (config.strategies.pruneTool.protectedTools.includes(toolParameterEntry.tool)) {
2727
return
2828
}
29-
if (toolParameterEntry.compacted) {
30-
return
31-
}
3229
const numericId = toolIdList.indexOf(toolCallId)
3330
const paramKey = extractParameterKey(toolParameterEntry.tool, toolParameterEntry.parameters)
3431
const description = paramKey ? `${toolParameterEntry.tool}, ${paramKey}` : toolParameterEntry.tool
@@ -111,6 +108,10 @@ const pruneToolOutputs = (
111108
messages: WithParts[]
112109
): void => {
113110
for (const msg of messages) {
111+
if (isMessageCompacted(state, msg)) {
112+
continue
113+
}
114+
114115
for (const part of msg.parts) {
115116
if (part.type !== 'tool') {
116117
continue

lib/messages/utils.ts

Lines changed: 11 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,6 @@
1-
import type { WithParts } from "../state"
1+
import { Logger } from "../logger"
2+
import { isMessageCompacted } from "../shared-utils"
3+
import type { SessionState, WithParts } from "../state"
24

35
/**
46
* Extracts a human-readable key from tool metadata for display purposes.
@@ -71,9 +73,16 @@ export const extractParameterKey = (tool: string, parameters: any): string => {
7173
return paramStr.substring(0, 50)
7274
}
7375

74-
export function buildToolIdList(messages: WithParts[]): string[] {
76+
export function buildToolIdList(
77+
state: SessionState,
78+
messages: WithParts[],
79+
logger: Logger
80+
): string[] {
7581
const toolIds: string[] = []
7682
for (const msg of messages) {
83+
if (isMessageCompacted(state, msg)) {
84+
continue
85+
}
7786
if (msg.parts) {
7887
for (const part of msg.parts) {
7988
if (part.type === 'tool' && part.callID && part.tool) {

lib/shared-utils.ts

Lines changed: 19 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,12 @@
1-
import { WithParts } from "./state"
1+
import { Logger } from "./logger"
2+
import { SessionState, WithParts } from "./state"
3+
4+
export const isMessageCompacted = (
5+
state: SessionState,
6+
msg: WithParts
7+
): boolean => {
8+
return msg.info.time.created < state.lastCompaction
9+
}
210

311
export const getLastUserMessage = (
412
messages: WithParts[]
@@ -11,3 +19,13 @@ export const getLastUserMessage = (
1119
}
1220
return null
1321
}
22+
23+
export const checkForCompaction = (
24+
state: SessionState,
25+
messages: WithParts[],
26+
logger: Logger
27+
): void => {
28+
for (const msg of messages) {
29+
30+
}
31+
}

lib/state/persistence.ts

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@ export interface PersistedSessionState {
1616
prune: Prune
1717
stats: SessionStats;
1818
lastUpdated: string;
19+
lastCompacted: number
1920
}
2021

2122
const STORAGE_DIR = join(
@@ -55,6 +56,7 @@ export async function saveSessionState(
5556
prune: sessionState.prune,
5657
stats: sessionState.stats,
5758
lastUpdated: new Date().toISOString(),
59+
lastCompacted: sessionState.lastCompaction
5860
};
5961

6062
const filePath = getSessionFilePath(sessionState.sessionId);
@@ -99,8 +101,7 @@ export async function loadSessionState(
99101
}
100102

101103
logger.info("Loaded session state from disk", {
102-
sessionId: sessionId,
103-
totalTokensSaved: state.stats.totalPruneTokens
104+
sessionId: sessionId
104105
});
105106

106107
return state;

lib/state/state.ts

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -41,7 +41,8 @@ export function createSessionState(): SessionState {
4141
},
4242
toolParameters: new Map<string, ToolParameterEntry>(),
4343
nudgeCounter: 0,
44-
lastToolPrune: false
44+
lastToolPrune: false,
45+
lastCompaction: 0
4546
}
4647
}
4748

@@ -58,6 +59,7 @@ export function resetSessionState(state: SessionState): void {
5859
state.toolParameters.clear()
5960
state.nudgeCounter = 0
6061
state.lastToolPrune = false
62+
state.lastCompaction = 0
6163
}
6264

6365
export async function ensureSessionInitialized(
@@ -95,4 +97,5 @@ export async function ensureSessionInitialized(
9597
pruneTokenCounter: persisted.stats?.pruneTokenCounter || 0,
9698
totalPruneTokens: persisted.stats?.totalPruneTokens || 0,
9799
}
100+
state.lastCompaction = persisted.lastCompacted || 0
98101
}

lib/state/tool-cache.ts

Lines changed: 10 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
import type { SessionState, ToolStatus, WithParts } from "./index"
22
import type { Logger } from "../logger"
33
import { PluginConfig } from "../config"
4+
import { isMessageCompacted } from "../shared-utils"
45

56
const MAX_TOOL_CACHE_SIZE = 1000
67

@@ -19,10 +20,17 @@ export async function syncToolCache(
1920
state.nudgeCounter = 0
2021

2122
for (const msg of messages) {
23+
if (isMessageCompacted(state, msg)) {
24+
continue
25+
}
26+
2227
for (const part of msg.parts) {
2328
if (part.type !== "tool" || !part.callID) {
2429
continue
2530
}
31+
if (state.toolParameters.has(part.callID)) {
32+
continue
33+
}
2634

2735
if (part.tool === "prune") {
2836
state.nudgeCounter = 0
@@ -31,25 +39,19 @@ export async function syncToolCache(
3139
}
3240
state.lastToolPrune = part.tool === "prune"
3341

34-
if (state.toolParameters.has(part.callID)) {
35-
continue
36-
}
37-
3842
state.toolParameters.set(
3943
part.callID,
4044
{
4145
tool: part.tool,
4246
parameters: part.state?.input ?? {},
4347
status: part.state.status as ToolStatus | undefined,
4448
error: part.state.status === "error" ? part.state.error : undefined,
45-
compacted: part.state.status === "completed" && !!part.state.time.compacted,
4649
}
4750
)
51+
logger.info("Cached tool id: " + part.callID)
4852
}
4953
}
50-
51-
// logger.info(`nudgeCounter=${state.nudgeCounter}, lastToolPrune=${state.lastToolPrune}`)
52-
54+
logger.info("Synced cache - size: " + state.toolParameters.size)
5355
trimToolParametersCache(state)
5456
} catch (error) {
5557
logger.warn("Failed to sync tool parameters from OpenCode", {

lib/state/types.ts

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,6 @@ export interface ToolParameterEntry {
1212
parameters: any
1313
status?: ToolStatus
1414
error?: string
15-
compacted?: boolean
1615
}
1716

1817
export interface SessionStats {
@@ -32,4 +31,5 @@ export interface SessionState {
3231
toolParameters: Map<string, ToolParameterEntry>
3332
nudgeCounter: number
3433
lastToolPrune: boolean
34+
lastCompaction: number
3535
}

lib/strategies/deduplication.ts

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,7 @@ export const deduplicate = (
2020
}
2121

2222
// Build list of all tool call IDs from messages (chronological order)
23-
const allToolIds = buildToolIdList(messages)
23+
const allToolIds = buildToolIdList(state, messages, logger)
2424
if (allToolIds.length === 0) {
2525
return
2626
}
@@ -68,7 +68,7 @@ export const deduplicate = (
6868
}
6969
}
7070

71-
state.stats.totalPruneTokens += calculateTokensSaved(messages, newPruneIds)
71+
state.stats.totalPruneTokens += calculateTokensSaved(state, messages, newPruneIds)
7272

7373
if (newPruneIds.length > 0) {
7474
state.prune.toolIds.push(...newPruneIds)

lib/strategies/on-idle.ts

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@ import { selectModel, ModelInfo } from "../model-selector"
77
import { saveSessionState } from "../state/persistence"
88
import { sendUnifiedNotification } from "../ui/notification"
99
import { calculateTokensSaved, getCurrentParams } from "./utils"
10+
import { isMessageCompacted } from "../shared-utils"
1011

1112
export interface OnIdleResult {
1213
prunedCount: number
@@ -18,6 +19,7 @@ export interface OnIdleResult {
1819
* Parse messages to extract tool information.
1920
*/
2021
function parseMessages(
22+
state: SessionState,
2123
messages: WithParts[],
2224
toolParametersCache: Map<string, ToolParameterEntry>
2325
): {
@@ -28,6 +30,9 @@ function parseMessages(
2830
const toolMetadata = new Map<string, ToolParameterEntry>()
2931

3032
for (const msg of messages) {
33+
if (isMessageCompacted(state, msg)) {
34+
continue
35+
}
3136
if (msg.parts) {
3237
for (const part of msg.parts) {
3338
if (part.type === "tool" && part.callID) {
@@ -224,7 +229,7 @@ export async function runOnIdle(
224229
}
225230

226231
const currentParams = getCurrentParams(messages, logger)
227-
const { toolCallIds, toolMetadata } = parseMessages(messages, state.toolParameters)
232+
const { toolCallIds, toolMetadata } = parseMessages(state, messages, state.toolParameters)
228233

229234
const alreadyPrunedIds = state.prune.toolIds
230235
const unprunedToolCallIds = toolCallIds.filter(id => !alreadyPrunedIds.includes(id))
@@ -273,7 +278,7 @@ export async function runOnIdle(
273278
const allPrunedIds = [...new Set([...alreadyPrunedIds, ...newlyPrunedIds])]
274279
state.prune.toolIds = allPrunedIds
275280

276-
state.stats.pruneTokenCounter += calculateTokensSaved(messages, newlyPrunedIds)
281+
state.stats.pruneTokenCounter += calculateTokensSaved(state, messages, newlyPrunedIds)
277282

278283
// Build tool metadata map for notification
279284
const prunedToolMetadata = new Map<string, ToolParameterEntry>()

0 commit comments

Comments
 (0)