Skip to content
Open
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
2 changes: 2 additions & 0 deletions 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 @@ -68,6 +69,7 @@ const plugin: Plugin = (async (ctx) => {
) as any,
"chat.message": createChatMessageHandler(state, logger, config, hostPermissions),
"experimental.text.complete": createTextCompleteHandler(),
event: createEventHandler(state, logger),
"command.execute.before": createCommandExecuteHandler(
ctx.client,
state,
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
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
27 changes: 27 additions & 0 deletions lib/compress/state.ts
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,31 @@ export function allocateRunId(state: SessionState): number {
return next
}

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

let updates = 0
for (const block of state.prune.messages.blocksById.values()) {
const matchesCall = block.compressCallId === callId
const matchesMessage = !block.compressCallId && block.compressMessageId === messageId
if (!matchesCall && !matchesMessage) {
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 +118,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
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
}
112 changes: 111 additions & 1 deletion lib/hooks.ts
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@ import {
} from "./messages"
import { renderSystemPrompt, type PromptStore } from "./prompts"
import { buildProtectedToolsExtension } from "./prompts/extensions/system"
import { attachCompressionDuration } from "./compress/state"
import {
applyPendingManualTrigger,
handleContextCommand,
Expand All @@ -29,7 +30,7 @@ import {
} from "./commands"
import { type HostPermissionSnapshot } from "./host-permissions"
import { compressPermission, syncCompressPermissionState } from "./compress-permission"
import { checkSession, ensureSessionInitialized, syncToolCache } from "./state"
import { checkSession, ensureSessionInitialized, saveSessionState, syncToolCache } from "./state"
import { cacheSystemPromptTokens } from "./ui/utils"

const INTERNAL_AGENT_SIGNATURES = [
Expand Down Expand Up @@ -266,6 +267,115 @@ export function createTextCompleteHandler() {
}
}

export function createEventHandler(state: SessionState, logger: Logger) {
return async (input: { event: any }) => {
const eventTime =
typeof input.event?.time === "number" && Number.isFinite(input.event.time)
? input.event.time
: typeof input.event?.properties?.time === "number" &&
Number.isFinite(input.event.properties.time)
? input.event.properties.time
: undefined

if (input.event.type !== "message.part.updated") {
return
}

const part = input.event.properties?.part
if (part?.type !== "tool" || part.tool !== "compress") {
return
}

if (part.state.status === "pending") {
if (typeof part.callID !== "string" || typeof part.messageID !== "string") {
return
}

if (state.compressionStarts.has(part.callID)) {
return
}

const startedAt = eventTime ?? Date.now()
state.compressionStarts.set(part.callID, {
messageId: part.messageID,
startedAt,
})
logger.debug("Recorded compression start", {
callID: part.callID,
messageID: part.messageID,
startedAt,
})
return
}

if (part.state.status === "completed") {
if (typeof part.callID !== "string" || typeof part.messageID !== "string") {
return
}

const start = state.compressionStarts.get(part.callID)
state.compressionStarts.delete(part.callID)

const runningAt =
typeof part.state.time?.start === "number" && Number.isFinite(part.state.time.start)
? part.state.time.start
: eventTime
const pendingToRunningMs =
start && typeof runningAt === "number"
? Math.max(0, runningAt - start.startedAt)
: undefined

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

const durationMs =
typeof pendingToRunningMs === "number" ? pendingToRunningMs : runtimeMs
if (typeof durationMs !== "number") {
return
}

const updates = attachCompressionDuration(
state,
part.callID,
part.messageID,
durationMs,
)
if (updates === 0) {
return
}

logger.info("Attached compression time to blocks", {
callID: part.callID,
messageID: part.messageID,
blocks: updates,
durationMs,
})

saveSessionState(state, logger).catch((error) => {
logger.warn("Failed to persist compression time update", {
error: error instanceof Error ? error.message : String(error),
})
})
return
}

if (part.state.status === "running") {
return
}

if (typeof part.callID === "string") {
state.compressionStarts.delete(part.callID)
}
}
}

export function createChatMessageHandler(
state: SessionState,
logger: Logger,
Expand Down
3 changes: 2 additions & 1 deletion lib/state/state.ts
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
import type { SessionState, ToolParameterEntry, WithParts } from "./types"
import type { CompressionStart, SessionState, ToolParameterEntry, WithParts } from "./types"
import type { Logger } from "../logger"
import { loadSessionState, saveSessionState } from "./persistence"
import {
Expand Down Expand Up @@ -81,6 +81,7 @@ export function createSessionState(): SessionState {
pruneTokenCounter: 0,
totalPruneTokens: 0,
},
compressionStarts: new Map<string, CompressionStart>(),
toolParameters: new Map<string, ToolParameterEntry>(),
subAgentResultCache: new Map<string, string>(),
toolIdList: [],
Expand Down
Loading
Loading