Skip to content

Commit 4e03a4b

Browse files
TarquinenTarquinen
committed
refactor: use Set for prune ID storage
- Change toolIds and messageIds from string[] to Set<string> - Update all .includes() to .has(), .push() to .add(), .length to .size - Add serialization layer in persistence.ts for JSON compatibility - Remove redundant Set wrappers now that state is already a Set - Fix session initialization order in dcp commands - Fix token calculation for pruned tools
1 parent 7453ed4 commit 4e03a4b

16 files changed

Lines changed: 69 additions & 47 deletions

‎lib/commands/context.ts‎

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -74,8 +74,8 @@ function analyzeTokens(state: SessionState, messages: WithParts[]): TokenBreakdo
7474
tools: 0,
7575
toolCount: 0,
7676
prunedTokens: state.stats.totalPruneTokens,
77-
prunedCount: state.prune.toolIds.length,
78-
prunedMessageCount: state.prune.messageIds.length,
77+
prunedCount: state.prune.toolIds.size,
78+
prunedMessageCount: state.prune.messageIds.size,
7979
total: 0,
8080
}
8181

@@ -129,7 +129,8 @@ function analyzeTokens(state: SessionState, messages: WithParts[]): TokenBreakdo
129129
foundToolIds.add(toolPart.callID)
130130
}
131131

132-
if (!isCompacted) {
132+
const isPruned = toolPart.callID && state.prune.toolIds.has(toolPart.callID)
133+
if (!isCompacted && !isPruned) {
133134
if (toolPart.state?.input) {
134135
const inputStr =
135136
typeof toolPart.state.input === "string"
@@ -177,7 +178,7 @@ function analyzeTokens(state: SessionState, messages: WithParts[]): TokenBreakdo
177178
breakdown.system = Math.max(0, firstInput - firstUserTokens)
178179
}
179180

180-
breakdown.tools = Math.max(0, toolInputTokens + toolOutputTokens - breakdown.prunedTokens)
181+
breakdown.tools = toolInputTokens + toolOutputTokens
181182
breakdown.assistant = Math.max(
182183
0,
183184
breakdown.total - breakdown.system - breakdown.user - breakdown.tools,

‎lib/commands/stats.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -48,7 +48,7 @@ export async function handleStatsCommand(ctx: StatsCommandContext): Promise<void
4848

4949
// Session stats from in-memory state
5050
const sessionTokens = state.stats.totalPruneTokens
51-
const sessionTools = state.prune.toolIds.length
51+
const sessionTools = state.prune.toolIds.size
5252

5353
// All-time stats from storage files
5454
const allTime = await loadAllSessionStats(logger)

‎lib/commands/sweep.ts‎

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -160,9 +160,8 @@ export async function handleSweepCommand(ctx: SweepCommandContext): Promise<void
160160
}
161161

162162
// Filter out already-pruned tools, protected tools, and protected file paths
163-
const existingPrunedSet = new Set(state.prune.toolIds)
164163
const newToolIds = toolIdsToSweep.filter((id) => {
165-
if (existingPrunedSet.has(id)) {
164+
if (state.prune.toolIds.has(id)) {
166165
return false
167166
}
168167
const entry = state.toolParameters.get(id)
@@ -213,7 +212,9 @@ export async function handleSweepCommand(ctx: SweepCommandContext): Promise<void
213212
}
214213

215214
// Add to prune list
216-
state.prune.toolIds.push(...newToolIds)
215+
for (const id of newToolIds) {
216+
state.prune.toolIds.add(id)
217+
}
217218

218219
// Calculate tokens saved
219220
const tokensSaved = calculateTokensSaved(state, messages, newToolIds)

‎lib/hooks.ts‎

Lines changed: 7 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@ import { handleStatsCommand } from "./commands/stats"
1010
import { handleContextCommand } from "./commands/context"
1111
import { handleHelpCommand } from "./commands/help"
1212
import { handleSweepCommand } from "./commands/sweep"
13+
import { ensureSessionInitialized } from "./state/state"
1314

1415
const INTERNAL_AGENT_SIGNATURES = [
1516
"You are a title generator",
@@ -92,15 +93,17 @@ export function createCommandExecuteHandler(
9293
}
9394

9495
if (input.command === "dcp") {
95-
const args = (input.arguments || "").trim().split(/\s+/).filter(Boolean)
96-
const subcommand = args[0]?.toLowerCase() || ""
97-
const _subArgs = args.slice(1)
98-
9996
const messagesResponse = await client.session.messages({
10097
path: { id: input.sessionID },
10198
})
10299
const messages = (messagesResponse.data || messagesResponse) as WithParts[]
103100

101+
await ensureSessionInitialized(client, state, input.sessionID, logger, messages)
102+
103+
const args = (input.arguments || "").trim().split(/\s+/).filter(Boolean)
104+
const subcommand = args[0]?.toLowerCase() || ""
105+
const _subArgs = args.slice(1)
106+
104107
if (subcommand === "context") {
105108
await handleContextCommand({
106109
client,

‎lib/messages/inject.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -88,7 +88,7 @@ const buildPrunableToolsList = (
8888
const toolIdList: string[] = buildToolIdList(state, messages, logger)
8989

9090
state.toolParameters.forEach((toolParameterEntry, toolCallId) => {
91-
if (state.prune.toolIds.includes(toolCallId)) {
91+
if (state.prune.toolIds.has(toolCallId)) {
9292
return
9393
}
9494

‎lib/messages/prune.ts‎

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -38,7 +38,7 @@ const pruneFullTool = (state: SessionState, logger: Logger, messages: WithParts[
3838
if (part.type !== "tool") {
3939
continue
4040
}
41-
if (!state.prune.toolIds.includes(part.callID)) {
41+
if (!state.prune.toolIds.has(part.callID)) {
4242
continue
4343
}
4444
if (part.tool !== "edit" && part.tool !== "write") {
@@ -79,7 +79,7 @@ const pruneToolOutputs = (state: SessionState, logger: Logger, messages: WithPar
7979
if (part.type !== "tool") {
8080
continue
8181
}
82-
if (!state.prune.toolIds.includes(part.callID)) {
82+
if (!state.prune.toolIds.has(part.callID)) {
8383
continue
8484
}
8585
if (part.state.status !== "completed") {
@@ -105,7 +105,7 @@ const pruneToolInputs = (state: SessionState, logger: Logger, messages: WithPart
105105
if (part.type !== "tool") {
106106
continue
107107
}
108-
if (!state.prune.toolIds.includes(part.callID)) {
108+
if (!state.prune.toolIds.has(part.callID)) {
109109
continue
110110
}
111111
if (part.state.status !== "completed") {
@@ -133,7 +133,7 @@ const pruneToolErrors = (state: SessionState, logger: Logger, messages: WithPart
133133
if (part.type !== "tool") {
134134
continue
135135
}
136-
if (!state.prune.toolIds.includes(part.callID)) {
136+
if (!state.prune.toolIds.has(part.callID)) {
137137
continue
138138
}
139139
if (part.state.status !== "error") {
@@ -158,7 +158,7 @@ const filterCompressedRanges = (
158158
logger: Logger,
159159
messages: WithParts[],
160160
): void => {
161-
if (!state.prune.messageIds?.length) {
161+
if (!state.prune.messageIds?.size) {
162162
return
163163
}
164164

@@ -193,7 +193,7 @@ const filterCompressedRanges = (
193193
}
194194

195195
// Skip messages that are in the prune list
196-
if (state.prune.messageIds.includes(msgId)) {
196+
if (state.prune.messageIds.has(msgId)) {
197197
continue
198198
}
199199

‎lib/shared-utils.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@ export const isMessageCompacted = (state: SessionState, msg: WithParts): boolean
55
if (msg.info.time.created < state.lastCompaction) {
66
return true
77
}
8-
if (state.prune.messageIds.includes(msg.info.id)) {
8+
if (state.prune.messageIds.has(msg.info.id)) {
99
return true
1010
}
1111
return false

‎lib/state/persistence.ts‎

Lines changed: 12 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -8,12 +8,18 @@ import * as fs from "fs/promises"
88
import { existsSync } from "fs"
99
import { homedir } from "os"
1010
import { join } from "path"
11-
import type { SessionState, SessionStats, Prune, CompressSummary } from "./types"
11+
import type { SessionState, SessionStats, CompressSummary } from "./types"
1212
import type { Logger } from "../logger"
1313

14+
/** Prune state as stored on disk (arrays for JSON compatibility) */
15+
export interface PersistedPrune {
16+
toolIds: string[]
17+
messageIds: string[]
18+
}
19+
1420
export interface PersistedSessionState {
1521
sessionName?: string
16-
prune: Prune
22+
prune: PersistedPrune
1723
compressSummaries: CompressSummary[]
1824
stats: SessionStats
1925
lastUpdated: string
@@ -45,7 +51,10 @@ export async function saveSessionState(
4551

4652
const state: PersistedSessionState = {
4753
sessionName: sessionName,
48-
prune: sessionState.prune,
54+
prune: {
55+
toolIds: [...sessionState.prune.toolIds],
56+
messageIds: [...sessionState.prune.messageIds],
57+
},
4958
compressSummaries: sessionState.compressSummaries,
5059
stats: sessionState.stats,
5160
lastUpdated: new Date().toISOString(),

‎lib/state/state.ts‎

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -48,8 +48,8 @@ export function createSessionState(): SessionState {
4848
sessionId: null,
4949
isSubAgent: false,
5050
prune: {
51-
toolIds: [],
52-
messageIds: [],
51+
toolIds: new Set<string>(),
52+
messageIds: new Set<string>(),
5353
},
5454
compressSummaries: [],
5555
stats: {
@@ -69,8 +69,8 @@ export function resetSessionState(state: SessionState): void {
6969
state.sessionId = null
7070
state.isSubAgent = false
7171
state.prune = {
72-
toolIds: [],
73-
messageIds: [],
72+
toolIds: new Set<string>(),
73+
messageIds: new Set<string>(),
7474
}
7575
state.compressSummaries = []
7676
state.stats = {
@@ -115,8 +115,8 @@ export async function ensureSessionInitialized(
115115
}
116116

117117
state.prune = {
118-
toolIds: persisted.prune.toolIds || [],
119-
messageIds: persisted.prune.messageIds || [],
118+
toolIds: new Set(persisted.prune.toolIds || []),
119+
messageIds: new Set(persisted.prune.messageIds || []),
120120
}
121121
state.compressSummaries = persisted.compressSummaries || []
122122
state.stats = {

‎lib/state/types.ts‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -26,8 +26,8 @@ export interface CompressSummary {
2626
}
2727

2828
export interface Prune {
29-
toolIds: string[]
30-
messageIds: string[]
29+
toolIds: Set<string>
30+
messageIds: Set<string>
3131
}
3232

3333
export interface SessionState {

0 commit comments

Comments
 (0)