Skip to content

Commit 72fdb68

Browse files
spoons-and-mirrorsspoons-and-mirrors
committed
refactor
1 parent dd79375 commit 72fdb68

File tree

12 files changed

+373
-236
lines changed

12 files changed

+373
-236
lines changed

‎lib/commands/compression-targets.ts‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@ export interface CompressionTarget {
55
runId: number
66
topic: string
77
compressedTokens: number
8+
durationMs: number
89
grouped: boolean
910
blocks: CompressionBlock[]
1011
}
@@ -26,6 +27,7 @@ function buildTarget(blocks: CompressionBlock[]): CompressionTarget {
2627
runId: first.runId,
2728
topic: grouped ? first.batchTopic || first.topic : first.topic,
2829
compressedTokens: ordered.reduce((total, block) => total + block.compressedTokens, 0),
30+
durationMs: ordered.reduce((total, block) => Math.max(total, block.durationMs), 0),
2931
grouped,
3032
blocks: ordered,
3133
}

‎lib/commands/stats.ts‎

Lines changed: 9 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@ import { sendIgnoredMessage } from "../ui/notification"
99
import { formatTokenCount } from "../ui/utils"
1010
import { loadAllSessionStats, type AggregatedStats } from "../state/persistence"
1111
import { getCurrentParams } from "../token-utils"
12+
import { getActiveCompressionTargets } from "./compression-targets"
1213

1314
export interface StatsCommandContext {
1415
client: any
@@ -22,7 +23,7 @@ function formatStatsMessage(
2223
sessionTokens: number,
2324
sessionTools: number,
2425
sessionMessages: number,
25-
sessionCompressionTimeMs: number,
26+
sessionDurationMs: number,
2627
allTime: AggregatedStats,
2728
): string {
2829
const lines: string[] = []
@@ -36,7 +37,7 @@ function formatStatsMessage(
3637
lines.push(` Tokens pruned: ~${formatTokenCount(sessionTokens)}`)
3738
lines.push(` Tools pruned: ${sessionTools}`)
3839
lines.push(` Messages pruned: ${sessionMessages}`)
39-
lines.push(` LLM summary time: ${formatCompressionTime(sessionCompressionTimeMs)}`)
40+
lines.push(` Compression time: ${formatCompressionTime(sessionDurationMs)}`)
4041
lines.push("")
4142
lines.push("All-time:")
4243
lines.push("─".repeat(60))
@@ -62,7 +63,10 @@ export async function handleStatsCommand(ctx: StatsCommandContext): Promise<void
6263

6364
// Session stats from in-memory state
6465
const sessionTokens = state.stats.totalPruneTokens
65-
const sessionCompressionTimeMs = state.stats.compressionTimeMs
66+
const sessionDurationMs = getActiveCompressionTargets(state.prune.messages).reduce(
67+
(total, target) => total + target.durationMs,
68+
0,
69+
)
6670

6771
const prunedToolIds = new Set<string>(state.prune.tools.keys())
6872
for (const block of state.prune.messages.blocksById.values()) {
@@ -88,7 +92,7 @@ export async function handleStatsCommand(ctx: StatsCommandContext): Promise<void
8892
sessionTokens,
8993
sessionTools,
9094
sessionMessages,
91-
sessionCompressionTimeMs,
95+
sessionDurationMs,
9296
allTime,
9397
)
9498

@@ -99,7 +103,7 @@ export async function handleStatsCommand(ctx: StatsCommandContext): Promise<void
99103
sessionTokens,
100104
sessionTools,
101105
sessionMessages,
102-
sessionCompressionTimeMs,
106+
sessionDurationMs,
103107
allTimeTokens: allTime.totalTokens,
104108
allTimeTools: allTime.totalTools,
105109
allTimeMessages: allTime.totalMessages,

‎lib/compress/message.ts‎

Lines changed: 74 additions & 68 deletions
Original file line numberDiff line numberDiff line change
@@ -46,86 +46,92 @@ export function createCompressMessageTool(ctx: ToolContext): ReturnType<typeof t
4646
description: runtimePrompts.compressMessage + MESSAGE_FORMAT_EXTENSION,
4747
args: buildSchema(),
4848
async execute(args, toolCtx) {
49-
const input = args as CompressMessageToolArgs
50-
validateArgs(input)
51-
const { rawMessages, searchContext } = await prepareSession(
52-
ctx,
53-
toolCtx,
54-
`Compress Message: ${input.topic}`,
55-
)
56-
const { plans, skippedIssues } = resolveMessages(
57-
input,
58-
searchContext,
59-
ctx.state,
60-
ctx.config,
61-
)
49+
try {
50+
const input = args as CompressMessageToolArgs
51+
validateArgs(input)
52+
const { rawMessages, searchContext } = await prepareSession(
53+
ctx,
54+
toolCtx,
55+
`Compress Message: ${input.topic}`,
56+
)
57+
const { plans, skippedIssues } = resolveMessages(
58+
input,
59+
searchContext,
60+
ctx.state,
61+
ctx.config,
62+
)
6263

63-
if (plans.length === 0 && skippedIssues.length > 0) {
64-
throw new Error(formatIssues(skippedIssues))
65-
}
64+
if (plans.length === 0 && skippedIssues.length > 0) {
65+
throw new Error(formatIssues(skippedIssues))
66+
}
6667

67-
const notifications: NotificationEntry[] = []
68+
const notifications: NotificationEntry[] = []
6869

69-
const preparedPlans: Array<{
70-
plan: (typeof plans)[number]
71-
summaryWithTools: string
72-
}> = []
70+
const preparedPlans: Array<{
71+
plan: (typeof plans)[number]
72+
summaryWithTools: string
73+
}> = []
7374

74-
for (const plan of plans) {
75-
const summaryWithTools = await appendProtectedTools(
76-
ctx.client,
77-
ctx.state,
78-
ctx.config.experimental.allowSubAgents,
79-
plan.entry.summary,
80-
plan.selection,
81-
searchContext,
82-
ctx.config.compress.protectedTools,
83-
ctx.config.protectedFilePatterns,
84-
)
75+
for (const plan of plans) {
76+
const summaryWithTools = await appendProtectedTools(
77+
ctx.client,
78+
ctx.state,
79+
ctx.config.experimental.allowSubAgents,
80+
plan.entry.summary,
81+
plan.selection,
82+
searchContext,
83+
ctx.config.compress.protectedTools,
84+
ctx.config.protectedFilePatterns,
85+
)
8586

86-
preparedPlans.push({
87-
plan,
88-
summaryWithTools,
89-
})
90-
}
87+
preparedPlans.push({
88+
plan,
89+
summaryWithTools,
90+
})
91+
}
9192

92-
const runId = allocateRunId(ctx.state)
93+
const runId = allocateRunId(ctx.state)
94+
const durationMs = ctx.state.compressionDurations.get(toolCtx.messageID) || 0
9395

94-
for (const { plan, summaryWithTools } of preparedPlans) {
95-
const blockId = allocateBlockId(ctx.state)
96-
const storedSummary = wrapCompressedSummary(blockId, summaryWithTools)
97-
const summaryTokens = countTokens(storedSummary)
96+
for (const { plan, summaryWithTools } of preparedPlans) {
97+
const blockId = allocateBlockId(ctx.state)
98+
const storedSummary = wrapCompressedSummary(blockId, summaryWithTools)
99+
const summaryTokens = countTokens(storedSummary)
98100

99-
applyCompressionState(
100-
ctx.state,
101-
{
102-
topic: plan.entry.topic,
103-
batchTopic: input.topic,
104-
startId: plan.entry.messageId,
105-
endId: plan.entry.messageId,
106-
mode: "message",
101+
applyCompressionState(
102+
ctx.state,
103+
{
104+
topic: plan.entry.topic,
105+
batchTopic: input.topic,
106+
startId: plan.entry.messageId,
107+
endId: plan.entry.messageId,
108+
mode: "message",
109+
runId,
110+
compressMessageId: toolCtx.messageID,
111+
summaryTokens,
112+
durationMs,
113+
},
114+
plan.selection,
115+
plan.anchorMessageId,
116+
blockId,
117+
storedSummary,
118+
[],
119+
)
120+
121+
notifications.push({
122+
blockId,
107123
runId,
108-
compressMessageId: toolCtx.messageID,
124+
summary: summaryWithTools,
109125
summaryTokens,
110-
},
111-
plan.selection,
112-
plan.anchorMessageId,
113-
blockId,
114-
storedSummary,
115-
[],
116-
)
117-
118-
notifications.push({
119-
blockId,
120-
runId,
121-
summary: summaryWithTools,
122-
summaryTokens,
123-
})
124-
}
126+
})
127+
}
125128

126-
await finalizeSession(ctx, toolCtx, rawMessages, notifications, input.topic)
129+
await finalizeSession(ctx, toolCtx, rawMessages, notifications, input.topic)
127130

128-
return formatResult(plans.length, skippedIssues)
131+
return formatResult(plans.length, skippedIssues)
132+
} finally {
133+
ctx.state.compressionDurations.delete(toolCtx.messageID)
134+
}
129135
},
130136
})
131137
}

0 commit comments

Comments
 (0)