@dyyz1993/pi-coding-agent 0.70.3 → 0.70.5

package/dist/core/agent-session.js

@@ -24,13 +24,14 @@ import { theme } from "../modes/interactive/theme/theme.js";
 import { stripFrontmatter } from "../utils/frontmatter.js";
 import { sleep } from "../utils/sleep.js";
 import { executeBashWithOperations } from "./bash-executor.js";
-import { calculateContextTokens, collectEntriesForBranchSummary, compact, estimateContextTokens, generateBranchSummary, prepareCompaction, shouldCompact, } from "./compaction/index.js";
+import { calculateContextTokens, collectEntriesForBranchSummary, compact, estimateContextTokens, estimateTokens, generateBranchSummary, prepareCompaction, shouldCompact, } from "./compaction/index.js";
 import { DEFAULT_THINKING_LEVEL } from "./defaults.js";
 import { exportSessionToHtml } from "./export-html/index.js";
 import { createToolHtmlRenderer } from "./export-html/tool-renderer.js";
 import { ExtensionRunner, wrapRegisteredTools, } from "./extensions/index.js";
 import { emitSessionShutdownEvent } from "./extensions/runner.js";
 import { handleLargeInput } from "./large-input.js";
+import { resolveModelAlias } from "./model-resolver.js";
 import { expandPromptTemplate } from "./prompt-templates.js";
 import { CURRENT_SESSION_VERSION, getLatestCompactionEntry } from "./session-manager.js";
 import { createSyntheticSourceInfo } from "./source-info.js";
@@ -116,6 +117,7 @@ export class AgentSession {
     _backgroundTasks = new Set();
     // Model registry for API key resolution
     _modelRegistry;
+    _tierModels = {};
     // Tool registry for extension getTools/setTools
     _toolRegistry = new Map();
     _toolDefinitions = new Map();
@@ -523,6 +525,12 @@ export class AgentSession {
     get model() {
         return this.agent.state.model;
     }
+    getTierModels() {
+        return this._tierModels;
+    }
+    setTierModels(mapping) {
+        this._tierModels = { ...mapping };
+    }
     /** Current thinking level */
     get thinkingLevel() {
         return this.agent.state.thinkingLevel;
@@ -1589,6 +1597,11 @@ export class AgentSession {
                 tokensBefore,
                 details,
             };
+            let tokensAfter = 0;
+            for (const msg of sessionContext.messages) {
+                tokensAfter += estimateTokens(msg);
+            }
+            result.tokensAfter = tokensAfter;
             this._emit({ type: "compaction_end", reason, result, aborted: false, willRetry });
             if (willRetry) {
                 const messages = this.agent.state.messages;
@@ -1884,8 +1897,36 @@ export class AgentSession {
             },
         });
     }
+    async _resolveOptionalModel(modelSpec) {
+        if (!modelSpec)
+            return this.model;
+        const aliasResolved = resolveModelAlias(modelSpec, this._tierModels);
+        const candidate = aliasResolved ?? modelSpec;
+        let provider;
+        let modelId;
+        if (candidate.includes("/")) {
+            const slashIdx = candidate.indexOf("/");
+            provider = candidate.substring(0, slashIdx);
+            modelId = candidate.substring(slashIdx + 1);
+        }
+        else {
+            modelId = candidate;
+        }
+        let resolved;
+        if (provider) {
+            resolved = this._modelRegistry.find(provider, modelId);
+        }
+        if (!resolved) {
+            const available = this._modelRegistry.getAvailable();
+            resolved = available.find((m) => m.id === modelId);
+        }
+        if (resolved && this._modelRegistry.hasConfiguredAuth(resolved)) {
+            return resolved;
+        }
+        return this.model;
+    }
     async callLLM(options) {
-        const model = this.model;
+        const model = await this._resolveOptionalModel(options.model);
         if (!model)
             throw new Error("No model selected");
         const auth = await this._modelRegistry.getApiKeyAndHeaders(model);
@@ -2052,7 +2093,7 @@ export class AgentSession {
         throw lastError ?? new Error("callLLMStructured failed");
     }
     async forkAgent(promptText, options) {
-        const model = this.model;
+        const model = await this._resolveOptionalModel(options?.model);
         if (!model)
             throw new Error("No model selected");
         const auth = await this._modelRegistry.getApiKeyAndHeaders(model);
@@ -2800,15 +2841,10 @@ 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") {
@@ -2816,18 +2852,28 @@ export class AgentSession {
                     if (assistant.stopReason !== "aborted" && assistant.stopReason !== "error") {
                         const contextTokens = calculateContextTokens(assistant.usage);
                         if (contextTokens > 0) {
-                            hasPostCompactionUsage = true;
+                            const percent = (contextTokens / contextWindow) * 100;
+                            return { tokens: contextTokens, contextWindow, percent };
                         }
-                        break;
                     }
+                    break;
                 }
             }
-            if (!hasPostCompactionUsage) {
-                return { tokens: null, contextWindow, percent: null };
-            }
         }
         const estimate = estimateContextTokens(this.messages);
-        const percent = (estimate.tokens / contextWindow) * 100;
+        if (latestCompaction && estimate.lastUsageIndex !== null) {
+            const usageMsg = this.messages[estimate.lastUsageIndex];
+            if (usageMsg?.role === "assistant" &&
+                usageMsg.timestamp <= new Date(latestCompaction.timestamp).getTime()) {
+                let estimated = 0;
+                for (const message of this.messages) {
+                    estimated += estimateTokens(message);
+                }
+                const fallbackPercent = contextWindow > 0 ? (estimated / contextWindow) * 100 : 0;
+                return { tokens: estimated, contextWindow, percent: fallbackPercent };
+            }
+        }
+        const percent = contextWindow > 0 ? (estimate.tokens / contextWindow) * 100 : 0;
         return {
             tokens: estimate.tokens,
             contextWindow,