diff --git a/src/agents/pi-embedded-runner/compact.ts b/src/agents/pi-embedded-runner/compact.ts index 2aad2431a..53b0ae8c0 100644 --- a/src/agents/pi-embedded-runner/compact.ts +++ b/src/agents/pi-embedded-runner/compact.ts @@ -1,7 +1,12 @@ import fs from "node:fs/promises"; import os from "node:os"; -import { createAgentSession, SessionManager, SettingsManager } from "@mariozechner/pi-coding-agent"; +import { + createAgentSession, + estimateTokens, + SessionManager, + SettingsManager, +} from "@mariozechner/pi-coding-agent"; import { resolveHeartbeatPrompt } from "../../auto-reply/heartbeat.js"; import type { ReasoningLevel, ThinkLevel } from "../../auto-reply/thinking.js"; @@ -370,6 +375,21 @@ export async function compactEmbeddedPiSession(params: { session.agent.replaceMessages(limited); } const result = await session.compact(params.customInstructions); + // Estimate tokens after compaction by summing token estimates for remaining messages + let tokensAfter: number | undefined; + try { + tokensAfter = 0; + for (const message of session.messages) { + tokensAfter += estimateTokens(message); + } + // Sanity check: tokensAfter should be less than tokensBefore + if (tokensAfter > result.tokensBefore) { + tokensAfter = undefined; // Don't trust the estimate + } + } catch { + // If estimation fails, leave tokensAfter undefined + tokensAfter = undefined; + } return { ok: true, compacted: true, @@ -377,6 +397,7 @@ export async function compactEmbeddedPiSession(params: { summary: result.summary, firstKeptEntryId: result.firstKeptEntryId, tokensBefore: result.tokensBefore, + tokensAfter, details: result.details, }, }; diff --git a/src/agents/pi-embedded-runner/types.ts b/src/agents/pi-embedded-runner/types.ts index 56380cd1d..6a1ee1128 100644 --- a/src/agents/pi-embedded-runner/types.ts +++ b/src/agents/pi-embedded-runner/types.ts @@ -59,6 +59,7 @@ export type EmbeddedPiCompactResult = { summary: string; firstKeptEntryId: string; tokensBefore: number; + tokensAfter?: number; details?: unknown; }; }; diff --git a/src/auto-reply/reply/commands-compact.ts b/src/auto-reply/reply/commands-compact.ts index da048ec65..7e3f0d960 100644 --- a/src/auto-reply/reply/commands-compact.ts +++ b/src/auto-reply/reply/commands-compact.ts @@ -83,18 +83,13 @@ export const handleCompactCommand: CommandHandler = async (params) => { ownerNumbers: params.command.ownerList.length > 0 ? params.command.ownerList : undefined, }); - const totalTokens = - params.sessionEntry.totalTokens ?? - (params.sessionEntry.inputTokens ?? 0) + (params.sessionEntry.outputTokens ?? 0); - const contextSummary = formatContextUsageShort( - totalTokens > 0 ? totalTokens : null, - params.contextTokens ?? params.sessionEntry.contextTokens ?? null, - ); const compactLabel = result.ok ? result.compacted - ? result.result?.tokensBefore - ? `Compacted (${formatTokenCount(result.result.tokensBefore)} before)` - : "Compacted" + ? result.result?.tokensBefore != null && result.result?.tokensAfter != null + ? `Compacted (${formatTokenCount(result.result.tokensBefore)} → ${formatTokenCount(result.result.tokensAfter)})` + : result.result?.tokensBefore + ? `Compacted (${formatTokenCount(result.result.tokensBefore)} before)` + : "Compacted" : "Compaction skipped" : "Compaction failed"; if (result.ok && result.compacted) { @@ -103,8 +98,20 @@ export const handleCompactCommand: CommandHandler = async (params) => { sessionStore: params.sessionStore, sessionKey: params.sessionKey, storePath: params.storePath, + // Update token counts after compaction + tokensAfter: result.result?.tokensAfter, }); } + // Use the post-compaction token count for context summary if available + const tokensAfterCompaction = result.result?.tokensAfter; + const totalTokens = + tokensAfterCompaction ?? + params.sessionEntry.totalTokens ?? + (params.sessionEntry.inputTokens ?? 0) + (params.sessionEntry.outputTokens ?? 0); + const contextSummary = formatContextUsageShort( + totalTokens > 0 ? totalTokens : null, + params.contextTokens ?? params.sessionEntry.contextTokens ?? null, + ); const reason = result.reason?.trim(); const line = reason ? `${compactLabel}: ${reason} • ${contextSummary}` diff --git a/src/auto-reply/reply/session-updates.ts b/src/auto-reply/reply/session-updates.ts index e5ad81d8e..970a714d0 100644 --- a/src/auto-reply/reply/session-updates.ts +++ b/src/auto-reply/reply/session-updates.ts @@ -237,23 +237,42 @@ export async function incrementCompactionCount(params: { sessionKey?: string; storePath?: string; now?: number; + /** Token count after compaction - if provided, updates session token counts */ + tokensAfter?: number; }): Promise { - const { sessionEntry, sessionStore, sessionKey, storePath, now = Date.now() } = params; + const { + sessionEntry, + sessionStore, + sessionKey, + storePath, + now = Date.now(), + tokensAfter, + } = params; if (!sessionStore || !sessionKey) return undefined; const entry = sessionStore[sessionKey] ?? sessionEntry; if (!entry) return undefined; const nextCount = (entry.compactionCount ?? 0) + 1; - sessionStore[sessionKey] = { - ...entry, + // Build update payload with compaction count and optionally updated token counts + const updates: Partial = { compactionCount: nextCount, updatedAt: now, }; + // If tokensAfter is provided, update the cached token counts to reflect post-compaction state + if (tokensAfter != null && tokensAfter > 0) { + updates.totalTokens = tokensAfter; + // Clear input/output breakdown since we only have the total estimate after compaction + updates.inputTokens = undefined; + updates.outputTokens = undefined; + } + sessionStore[sessionKey] = { + ...entry, + ...updates, + }; if (storePath) { await updateSessionStore(storePath, (store) => { store[sessionKey] = { ...store[sessionKey], - compactionCount: nextCount, - updatedAt: now, + ...updates, }; }); }