@mariozechner/pi-coding-agent 0.52.9 → 0.52.10

package/dist/core/agent-session.js

@@ -26,6 +26,7 @@ import { exportSessionToHtml } from "./export-html/index.js";
 import { createToolHtmlRenderer } from "./export-html/tool-renderer.js";
 import { ExtensionRunner, wrapRegisteredTools, wrapToolsWithExtensions, } from "./extensions/index.js";
 import { expandPromptTemplate } from "./prompt-templates.js";
+import { getLatestCompactionEntry } from "./session-manager.js";
 import { BUILTIN_SLASH_COMMANDS } from "./slash-commands.js";
 import { buildSystemPrompt } from "./system-prompt.js";
 import { createAllTools } from "./tools/index.js";
@@ -264,6 +265,57 @@ export class AgentSession {
             await this._extensionRunner.emit(extensionEvent);
             this._turnIndex++;
         }
+        else if (event.type === "message_start") {
+            const extensionEvent = {
+                type: "message_start",
+                message: event.message,
+            };
+            await this._extensionRunner.emit(extensionEvent);
+        }
+        else if (event.type === "message_update") {
+            const extensionEvent = {
+                type: "message_update",
+                message: event.message,
+                assistantMessageEvent: event.assistantMessageEvent,
+            };
+            await this._extensionRunner.emit(extensionEvent);
+        }
+        else if (event.type === "message_end") {
+            const extensionEvent = {
+                type: "message_end",
+                message: event.message,
+            };
+            await this._extensionRunner.emit(extensionEvent);
+        }
+        else if (event.type === "tool_execution_start") {
+            const extensionEvent = {
+                type: "tool_execution_start",
+                toolCallId: event.toolCallId,
+                toolName: event.toolName,
+                args: event.args,
+            };
+            await this._extensionRunner.emit(extensionEvent);
+        }
+        else if (event.type === "tool_execution_update") {
+            const extensionEvent = {
+                type: "tool_execution_update",
+                toolCallId: event.toolCallId,
+                toolName: event.toolName,
+                args: event.args,
+                partialResult: event.partialResult,
+            };
+            await this._extensionRunner.emit(extensionEvent);
+        }
+        else if (event.type === "tool_execution_end") {
+            const extensionEvent = {
+                type: "tool_execution_end",
+                toolCallId: event.toolCallId,
+                toolName: event.toolName,
+                result: event.result,
+                isError: event.isError,
+            };
+            await this._extensionRunner.emit(extensionEvent);
+        }
     }
     /**
      * Subscribe to agent events.
@@ -1190,8 +1242,8 @@ export class AgentSession {
         // The error shouldn't trigger another compaction since we already compacted.
         // Example: opus fails → switch to codex → compact → switch back to opus → opus error
         // is still in context but shouldn't trigger compaction again.
-        const compactionEntry = this.sessionManager.getBranch().find((e) => e.type === "compaction");
-        const errorIsFromBeforeCompaction = compactionEntry && assistantMessage.timestamp < new Date(compactionEntry.timestamp).getTime();
+        const compactionEntry = getLatestCompactionEntry(this.sessionManager.getBranch());
+        const errorIsFromBeforeCompaction = compactionEntry !== null && assistantMessage.timestamp < new Date(compactionEntry.timestamp).getTime();
         // Case 1: Overflow - LLM returned context overflow error
         if (sameModel && !errorIsFromBeforeCompaction && isContextOverflow(assistantMessage, contextWindow)) {
             // Remove the error message from agent state (it IS saved to session for history,
@@ -2186,15 +2238,38 @@ 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 percent = (estimate.tokens / contextWindow) * 100;
         return {
             tokens: estimate.tokens,
             contextWindow,
             percent,
-            usageTokens: estimate.usageTokens,
-            trailingTokens: estimate.trailingTokens,
-            lastUsageIndex: estimate.lastUsageIndex,
         };
     }
     /**