@ashx-j/lunr 0.2.27 → 0.2.28

package/dist/core/agent-session.js

@@ -1,3 +1,4 @@
+import { collectChildUsage, collectUsageRequests, totalChildUsage, totalRequestUsage, } from "./usage-accounting.js";
 /**
  * AgentSession - Core abstraction for agent lifecycle and session management.
  *
@@ -14,7 +15,7 @@
  */
 import { existsSync, mkdirSync, readFileSync, writeFileSync } from "node:fs";
 import { basename, dirname, join, resolve } from "node:path";
-import { clampThinkingLevel, cleanupSessionResources, getSupportedThinkingLevels, isContextOverflow, isRetryableAssistantError, modelsAreEqual, resetApiProviders, streamSimple, } from "@ashx-j/lunr-ai/compat";
+import { clampThinkingLevel, cleanupSessionResources, estimateContextTokens as estimateRequestContextTokens, getSupportedThinkingLevels, isContextOverflow, isRetryableAssistantError, modelsAreEqual, resetApiProviders, streamSimple, } from "@ashx-j/lunr-ai/compat";
 import { getAgentDir } from "../config.js";
 import { getThemeByName, theme } from "../modes/interactive/theme/theme.js";
 import { stripFrontmatter } from "../utils/frontmatter.js";
@@ -27,6 +28,7 @@ import { DEFAULT_THINKING_LEVEL } from "./defaults.js";
 import { ExtensionRunner, wrapRegisteredTools, } from "./extensions/index.js";
 import { emitSessionShutdownEvent } from "./extensions/runner.js";
 import { MEMORY_TOOL_NAMES } from "./memory-cap.js";
+import { convertToLlm } from "./messages.js";
 import { isUserInstructionsPath, loadSelectedUserInstructions } from "./model-instructions.js";
 import { ModelRegistry } from "./model-registry.js";
 import { AUTO_MODE_ADDENDUM, deletePermissionContext, gateToolCall, getPermissionMode, initializePermissionContext, setPermissionMode, } from "./permissions.js";
@@ -355,12 +357,17 @@ export class AgentSession {
                 : undefined);
         this.agent.prepareNextTurnWithContext = async (turn, signal) => {
             const previousSnapshot = await previousPrepareNextTurnWithContext?.(turn, signal);
-            let previousContext = previousSnapshot?.context ?? turn.context;
+            let previousContext = {
+                ...(previousSnapshot?.context ?? turn.context),
+                systemPrompt: this._withSystemPromptAppend(this._systemPromptOverride ?? this._baseSystemPrompt),
+                tools: this.agent.state.tools.slice(),
+            };
             const model = this.model;
             const settings = this.settingsManager.getCompactionSettings();
             if (model?.provider === "openai-codex" &&
                 settings.enabled &&
-                shouldCompact(estimateContextTokens(previousContext.messages).tokens, model.contextWindow, settings)) {
+                shouldCompact(estimateRequestContextTokens({ ...previousContext, messages: convertToLlm(previousContext.messages) })
+                    .tokens, model.contextWindow, settings)) {
                 const previousCompactionId = getLatestCompactionEntry(this.sessionManager.getBranch())?.id;
                 await this._runAutoCompaction("threshold", true);
                 const currentCompaction = getLatestCompactionEntry(this.sessionManager.getBranch());
@@ -1981,7 +1988,7 @@ export class AgentSession {
             }
             else {
                 // Generate compaction result
-                const result = await compact(preparation, this.model, apiKey, headers, customInstructions, this._compactionAbortController.signal, this.thinkingLevel, this.agent.streamFn, env);
+                const result = await compact(preparation, this.model, apiKey, headers, customInstructions, this._compactionAbortController.signal, this.thinkingLevel, this.agent.streamFn, env, (response) => this.sessionManager.appendRequestUsage("compaction", response));
                 summary = result.summary;
                 firstKeptEntryId = result.firstKeptEntryId;
                 tokensBefore = result.tokensBefore;
@@ -2215,7 +2222,7 @@ export class AgentSession {
             }
             else {
                 // Generate compaction result
-                const compactResult = await compact(preparation, this.model, apiKey, headers, undefined, this._autoCompactionAbortController.signal, this.thinkingLevel, this.agent.streamFn, env);
+                const compactResult = await compact(preparation, this.model, apiKey, headers, undefined, this._autoCompactionAbortController.signal, this.thinkingLevel, this.agent.streamFn, env, (response) => this.sessionManager.appendRequestUsage("compaction", response));
                 summary = compactResult.summary;
                 firstKeptEntryId = compactResult.firstKeptEntryId;
                 tokensBefore = compactResult.tokensBefore;
@@ -3038,6 +3045,7 @@ export class AgentSession {
                     replaceInstructions,
                     reserveTokens: branchSummarySettings.reserveTokens,
                     streamFn: this.agent.streamFn,
+                    onUsage: (response) => this.sessionManager.appendRequestUsage("branch-summary", response),
                 });
                 if (result.aborted) {
                     return { cancelled: true, aborted: true };
@@ -3162,7 +3170,7 @@ export class AgentSession {
     /**
      * Get session statistics. Aggregates over ALL session entries (including
      * history that was compacted away), so token/cost totals reflect what was
-     * actually billed across the session.
+     * reported or estimated across requests incurred in this session.
      */
     getSessionStats() {
         let userMessages = 0;
@@ -3170,11 +3178,12 @@ export class AgentSession {
         let toolResults = 0;
         let totalMessages = 0;
         let toolCalls = 0;
-        let totalInput = 0;
-        let totalOutput = 0;
-        let totalCacheRead = 0;
-        let totalCacheWrite = 0;
-        let totalCost = 0;
+        const entries = this.sessionManager.getEntries();
+        const parentUsage = totalRequestUsage(collectUsageRequests(entries));
+        const childUsage = totalChildUsage(collectChildUsage(entries));
+        const combinedUsage = { ...parentUsage };
+        for (const key of Object.keys(combinedUsage))
+            combinedUsage[key] += childUsage[key];
         for (const entry of this.sessionManager.getEntries()) {
             if (entry.type !== "message")
                 continue;
@@ -3192,12 +3201,6 @@ export class AgentSession {
                 if (Array.isArray(assistantMsg.content)) {
                     toolCalls += assistantMsg.content.filter((c) => c.type === "toolCall").length;
                 }
-                const usage = assistantMsg.usage;
-                totalInput += usage.input;
-                totalOutput += usage.output;
-                totalCacheRead += usage.cacheRead;
-                totalCacheWrite += usage.cacheWrite;
-                totalCost += usage.cost.total;
             }
         }
         return {
@@ -3209,13 +3212,16 @@ export class AgentSession {
             toolResults,
             totalMessages,
             tokens: {
-                input: totalInput,
-                output: totalOutput,
-                cacheRead: totalCacheRead,
-                cacheWrite: totalCacheWrite,
-                total: totalInput + totalOutput + totalCacheRead + totalCacheWrite,
+                input: parentUsage.input,
+                output: parentUsage.output,
+                cacheRead: parentUsage.cacheRead,
+                cacheWrite: parentUsage.cacheWrite,
+                total: parentUsage.total,
             },
-            cost: totalCost,
+            cost: parentUsage.cost,
+            parentUsage,
+            childUsage,
+            combinedUsage,
             contextUsage: this.getContextUsage(),
         };
     }
@@ -3226,38 +3232,17 @@ export class AgentSession {
         const contextWindow = model.contextWindow ?? 0;
         if (contextWindow <= 0)
             return undefined;
-        // After compaction, the last assistant usage reflects pre-compaction context size.
-        // We can only trust usage from an assistant that responded after the latest compaction.
-        // If no such assistant exists, context token count is unknown until the next LLM response.
-        const branchEntries = this.sessionManager.getBranch();
-        const latestCompaction = getLatestCompactionEntry(branchEntries);
-        if (latestCompaction) {
-            // Check if there's a valid assistant usage after the compaction boundary
-            const compactionIndex = branchEntries.lastIndexOf(latestCompaction);
-            let hasPostCompactionUsage = false;
-            for (let i = branchEntries.length - 1; i > compactionIndex; i--) {
-                const entry = branchEntries[i];
-                if (entry.type === "message" && entry.message.role === "assistant") {
-                    const assistant = entry.message;
-                    if (assistant.stopReason !== "aborted" && assistant.stopReason !== "error") {
-                        const contextTokens = calculateContextTokens(assistant.usage);
-                        if (contextTokens > 0) {
-                            hasPostCompactionUsage = true;
-                            break;
-                        }
-                    }
-                }
-            }
-            if (!hasPostCompactionUsage) {
-                return { tokens: null, contextWindow, percent: null };
-            }
-        }
-        const estimate = estimateContextTokens(this.messages);
+        const estimate = estimateRequestContextTokens({
+            systemPrompt: this.systemPrompt,
+            tools: this.agent.state.tools,
+            messages: convertToLlm(this.messages),
+        });
         const percent = (estimate.tokens / contextWindow) * 100;
         return {
             tokens: estimate.tokens,
             contextWindow,
             percent,
+            estimated: estimate.lastUsageIndex === null || estimate.trailingTokens > 0,
         };
     }
     /**