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
3 changes: 2 additions & 1 deletion index.ts
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,7 @@ import {
createChatMessageHandler,
createChatMessageTransformHandler,
createCommandExecuteHandler,
createEventHandler,
createSystemPromptHandler,
createTextCompleteHandler,
} from "./lib/hooks"
Expand Down Expand Up @@ -57,7 +58,6 @@ const plugin: Plugin = (async (ctx) => {
config,
prompts,
),

"experimental.chat.messages.transform": createChatMessageTransformHandler(
ctx.client,
state,
Expand All @@ -76,6 +76,7 @@ const plugin: Plugin = (async (ctx) => {
ctx.directory,
hostPermissions,
),
event: createEventHandler(state, logger),
tool: {
...(config.compress.permission !== "deny" && {
compress:
Expand Down
2 changes: 2 additions & 0 deletions lib/commands/compression-targets.ts
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@ export interface CompressionTarget {
runId: number
topic: string
compressedTokens: number
durationMs: number
grouped: boolean
blocks: CompressionBlock[]
}
Expand All @@ -26,6 +27,7 @@ function buildTarget(blocks: CompressionBlock[]): CompressionTarget {
runId: first.runId,
topic: grouped ? first.batchTopic || first.topic : first.topic,
compressedTokens: ordered.reduce((total, block) => total + block.compressedTokens, 0),
durationMs: ordered.reduce((total, block) => Math.max(total, block.durationMs), 0),
grouped,
blocks: ordered,
}
Expand Down
70 changes: 65 additions & 5 deletions lib/commands/stats.ts
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@ import { sendIgnoredMessage } from "../ui/notification"
import { formatTokenCount } from "../ui/utils"
import { loadAllSessionStats, type AggregatedStats } from "../state/persistence"
import { getCurrentParams } from "../token-utils"
import { getActiveCompressionTargets } from "./compression-targets"

export interface StatsCommandContext {
client: any
Expand All @@ -20,8 +21,10 @@ export interface StatsCommandContext {

function formatStatsMessage(
sessionTokens: number,
sessionSummaryTokens: number,
sessionTools: number,
sessionMessages: number,
sessionDurationMs: number,
allTime: AggregatedStats,
): string {
const lines: string[] = []
Expand All @@ -30,11 +33,15 @@ function formatStatsMessage(
lines.push("│ DCP Statistics │")
lines.push("╰───────────────────────────────────────────────────────────╯")
lines.push("")
lines.push("Session:")
lines.push("Compression:")
lines.push("─".repeat(60))
lines.push(` Tokens pruned: ~${formatTokenCount(sessionTokens)}`)
lines.push(` Tools pruned: ${sessionTools}`)
lines.push(` Messages pruned: ${sessionMessages}`)
lines.push(
` Tokens in|out: ~${formatTokenCount(sessionTokens)} | ~${formatTokenCount(sessionSummaryTokens)}`,
)
lines.push(` Ratio: ${formatCompressionRatio(sessionTokens, sessionSummaryTokens)}`)
lines.push(` Time: ${formatCompressionTime(sessionDurationMs)}`)
lines.push(` Messages: ${sessionMessages}`)
lines.push(` Tools: ${sessionTools}`)
lines.push("")
lines.push("All-time:")
lines.push("─".repeat(60))
Expand All @@ -46,11 +53,55 @@ function formatStatsMessage(
return lines.join("\n")
}

function formatCompressionRatio(inputTokens: number, outputTokens: number): string {
if (inputTokens <= 0) {
return "0:1"
}

if (outputTokens <= 0) {
return "∞:1"
}

const ratio = Math.max(1, Math.round(inputTokens / outputTokens))
return `${ratio}:1`
}

function formatCompressionTime(ms: number): string {
const safeMs = Math.max(0, Math.round(ms))
if (safeMs < 1000) {
return `${safeMs} ms`
}

const totalSeconds = safeMs / 1000
if (totalSeconds < 60) {
return `${totalSeconds.toFixed(1)} s`
}

const wholeSeconds = Math.floor(totalSeconds)
const hours = Math.floor(wholeSeconds / 3600)
const minutes = Math.floor((wholeSeconds % 3600) / 60)
const seconds = wholeSeconds % 60

if (hours > 0) {
return `${hours}h ${minutes}m ${seconds}s`
}

return `${minutes}m ${seconds}s`
}

export async function handleStatsCommand(ctx: StatsCommandContext): Promise<void> {
const { client, state, logger, sessionId, messages } = ctx

// Session stats from in-memory state
const sessionTokens = state.stats.totalPruneTokens
const sessionSummaryTokens = Array.from(state.prune.messages.blocksById.values()).reduce(
(total, block) => (block.active ? total + block.summaryTokens : total),
0,
)
const sessionDurationMs = getActiveCompressionTargets(state.prune.messages).reduce(
(total, target) => total + target.durationMs,
0,
)

const prunedToolIds = new Set<string>(state.prune.tools.keys())
for (const block of state.prune.messages.blocksById.values()) {
Expand All @@ -72,15 +123,24 @@ export async function handleStatsCommand(ctx: StatsCommandContext): Promise<void
// All-time stats from storage files
const allTime = await loadAllSessionStats(logger)

const message = formatStatsMessage(sessionTokens, sessionTools, sessionMessages, allTime)
const message = formatStatsMessage(
sessionTokens,
sessionSummaryTokens,
sessionTools,
sessionMessages,
sessionDurationMs,
allTime,
)

const params = getCurrentParams(state, messages, logger)
await sendIgnoredMessage(client, sessionId, message, params, logger)

logger.info("Stats command executed", {
sessionTokens,
sessionSummaryTokens,
sessionTools,
sessionMessages,
sessionDurationMs,
allTimeTokens: allTime.totalTokens,
allTimeTools: allTime.totalTools,
allTimeMessages: allTime.totalMessages,
Expand Down
5 changes: 5 additions & 0 deletions lib/compress/message.ts
Original file line number Diff line number Diff line change
Expand Up @@ -48,6 +48,10 @@ export function createCompressMessageTool(ctx: ToolContext): ReturnType<typeof t
async execute(args, toolCtx) {
const input = args as CompressMessageToolArgs
validateArgs(input)
const callId =
typeof (toolCtx as unknown as { callID?: unknown }).callID === "string"
? (toolCtx as unknown as { callID: string }).callID
: undefined

const { rawMessages, searchContext } = await prepareSession(
ctx,
Expand Down Expand Up @@ -107,6 +111,7 @@ export function createCompressMessageTool(ctx: ToolContext): ReturnType<typeof t
mode: "message",
runId,
compressMessageId: toolCtx.messageID,
compressCallId: callId,
summaryTokens,
},
plan.selection,
Expand Down
2 changes: 2 additions & 0 deletions lib/compress/pipeline.ts
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@ import { sendCompressNotification } from "../ui/notification"
import type { ToolContext } from "./types"
import { buildSearchContext, fetchSessionMessages } from "./search"
import type { SearchContext } from "./types"
import { applyPendingCompressionDurations } from "./timing"

interface RunContext {
ask(input: {
Expand Down Expand Up @@ -83,6 +84,7 @@ export async function finalizeSession(
batchTopic: string | undefined,
): Promise<void> {
ctx.state.manualMode = ctx.state.manualMode ? "active" : false
applyPendingCompressionDurations(ctx.state)
await saveSessionState(ctx.state, ctx.logger)

const params = getCurrentParams(ctx.state, rawMessages, ctx.logger)
Expand Down
5 changes: 5 additions & 0 deletions lib/compress/range.ts
Original file line number Diff line number Diff line change
Expand Up @@ -59,6 +59,10 @@ export function createCompressRangeTool(ctx: ToolContext): ReturnType<typeof too
async execute(args, toolCtx) {
const input = args as CompressRangeToolArgs
validateArgs(input)
const callId =
typeof (toolCtx as unknown as { callID?: unknown }).callID === "string"
? (toolCtx as unknown as { callID: string }).callID
: undefined

const { rawMessages, searchContext } = await prepareSession(
ctx,
Expand Down Expand Up @@ -148,6 +152,7 @@ export function createCompressRangeTool(ctx: ToolContext): ReturnType<typeof too
mode: "range",
runId,
compressMessageId: toolCtx.messageID,
compressCallId: callId,
summaryTokens,
},
preparedPlan.selection,
Expand Down
26 changes: 25 additions & 1 deletion lib/compress/state.ts
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
import type { CompressionBlock, SessionState } from "../state"
import type { CompressionBlock, PruneMessagesState, SessionState } from "../state"
import { formatBlockRef, formatMessageIdTag } from "../message-ids"
import type { AppliedCompressionResult, CompressionStateInput, SelectionResolution } from "./types"

Expand Down Expand Up @@ -26,6 +26,28 @@ export function allocateRunId(state: SessionState): number {
return next
}

export function attachCompressionDuration(
messagesState: PruneMessagesState,
callId: string,
durationMs: number,
): number {
if (typeof durationMs !== "number" || !Number.isFinite(durationMs)) {
return 0
}

let updates = 0
for (const block of messagesState.blocksById.values()) {
if (block.compressCallId !== callId) {
continue
}

block.durationMs = durationMs
updates++
}

return updates
}

export function wrapCompressedSummary(blockId: number, summary: string): string {
const header = COMPRESSED_BLOCK_HEADER
const footer = formatMessageIdTag(formatBlockRef(blockId))
Expand Down Expand Up @@ -93,13 +115,15 @@ export function applyCompressionState(
deactivatedByUser: false,
compressedTokens: 0,
summaryTokens: input.summaryTokens,
durationMs: 0,
mode: input.mode,
topic: input.topic,
batchTopic: input.batchTopic,
startId: input.startId,
endId: input.endId,
anchorMessageId,
compressMessageId: input.compressMessageId,
compressCallId: input.compressCallId,
includedBlockIds: included,
consumedBlockIds: consumed,
parentBlockIds: [],
Expand Down
66 changes: 66 additions & 0 deletions lib/compress/timing.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,66 @@
import type { SessionState } from "../state/types"
import { attachCompressionDuration } from "./state"

export interface PendingCompressionDuration {
callId: string
durationMs: number
}

export interface CompressionTimingState {
startsByCallId: Map<string, number>
pendingByCallId: Map<string, PendingCompressionDuration>
}

export function consumeCompressionStart(state: SessionState, callId: string): number | undefined {
const start = state.compressionTiming.startsByCallId.get(callId)
state.compressionTiming.startsByCallId.delete(callId)
return start
}

export function resolveCompressionDuration(
startedAt: number | undefined,
eventTime: number | undefined,
partTime: { start?: unknown; end?: unknown } | undefined,
): number | undefined {
const runningAt =
typeof partTime?.start === "number" && Number.isFinite(partTime.start)
? partTime.start
: eventTime
const pendingToRunningMs =
typeof startedAt === "number" && typeof runningAt === "number"
? Math.max(0, runningAt - startedAt)
: undefined

const toolStart = partTime?.start
const toolEnd = partTime?.end
const runtimeMs =
typeof toolStart === "number" &&
Number.isFinite(toolStart) &&
typeof toolEnd === "number" &&
Number.isFinite(toolEnd)
? Math.max(0, toolEnd - toolStart)
: undefined

return typeof pendingToRunningMs === "number" ? pendingToRunningMs : runtimeMs
}

export function applyPendingCompressionDurations(state: SessionState): number {
if (state.compressionTiming.pendingByCallId.size === 0) {
return 0
}

let updates = 0
for (const [callId, entry] of state.compressionTiming.pendingByCallId) {
const applied = attachCompressionDuration(
state.prune.messages,
entry.callId,
entry.durationMs,
)
if (applied > 0) {
updates += applied
state.compressionTiming.pendingByCallId.delete(callId)
}
}

return updates
}
1 change: 1 addition & 0 deletions lib/compress/types.ts
Original file line number Diff line number Diff line change
Expand Up @@ -103,5 +103,6 @@ export interface CompressionStateInput {
mode: CompressionMode
runId: number
compressMessageId: string
compressCallId?: string
summaryTokens: number
}
Loading
Loading