@ponythewhite/base-context 1.0.0 → 1.0.2

package/dist/core/agent-session.js

@@ -26,7 +26,7 @@ import { createToolHtmlRenderer } from "./export-html/tool-renderer.js";
 import { ExtensionRunner, wrapRegisteredTools, } from "./extensions/index.js";
 import { emitSessionShutdownEvent } from "./extensions/runner.js";
 import { createGoalContextMessage, emptyGoalState, GOAL_CONTEXT_CUSTOM_TYPE, GOAL_CONTEXT_PREVIEW_LABEL, GOAL_SKILL_NAME, GOAL_STATE_CUSTOM_TYPE, goalHostResponse, goalTokenDeltaForUsage, isPersistedGoalState, normalizeGoalState, validateGoalBudget, validateGoalObjective, } from "./goals.js";
-import { captureNativeBranchRequests, captureNativeCompactionRequests, captureNativePlannerRequests, captureNativeReviewerRequests, InferenceCoordinator, } from "./inference-coordinator.js";
+import { captureNativeBranchRequests, captureNativeCompactionRequests, captureNativePlannerRequests, captureNativeReviewerRequests, InferenceCoordinator, takeNativeInferenceAuthSource, } from "./inference-coordinator.js";
 import { snapshotPathIn } from "./kernel/state-snapshot.js";
 import { convertToLlm, createCompactionOutcomeMessage, createHeartbeatPromptMessage, createRefinementOutcomeMessage, createRlmChildFailureMessage, createRlmChildTerminalNoticeMessage, createSessionSlashCommandMessage, createSessionSlashCommandResultMessage, HEARTBEAT_PROMPT_CUSTOM_TYPE, HEARTBEAT_PROMPT_PREVIEW_LABEL, IPYTHON_STATE_RESTORED_CUSTOM_TYPE, isSessionSlashCommandMessage, RLM_CHILD_FAILURE_CUSTOM_TYPE, RLM_CHILD_TERMINAL_NOTICE_CUSTOM_TYPE, } from "./messages.js";
 import { throwIfPromptAdmissionCancelled } from "./prompt-admission.js";
@@ -190,6 +190,7 @@ class StaleCompactionOwnerError extends Error {
 function createPostCompactionContinuationSettlement(resume) {
     return { ...createAgentMessageDeferred(), resume, settled: false };
 }
+import { LocalRlmSubagentCapacity, } from "./rlm-max-subagents.js";
 const KERNEL_STATE_LISTING_TIMEOUT_MS = 5000;
 const RLM_MAX_DEPTH_STATE_CUSTOM_TYPE = "rlm_max_depth_state";
 function noopRlmChildAbort() { }
@@ -349,7 +350,6 @@ export class AgentSession {
     _retryAttempt = 0;
     _retryPromise = undefined;
     _retryResolve = undefined;
-    _retryAuthFailureSources = [];
     _agentMessageClearEpoch = 0;
     _agentMessageOutcomes = new Map();
     _lateIpythonSentAgentMessages = new Map();
@@ -411,8 +411,10 @@ export class AgentSession {
     _repliedToParentSinceTask;
     _parentReplyCount = 0;
     _subagentRuntimeHost;
-    _rlmChildAdmission;
+    _rlmChildAdmissions = new Set();
+    _rlmSubagentCapacity;
     _rlmParentAdmission;
+    _rlmRootAdmission;
     _rlmResidentDisposalComplete = false;
     _releaseRlmResidentCapacity;
     _activeRlmChildRuns = new Map();
@@ -474,6 +476,12 @@ export class AgentSession {
     _serializedPlanClaim;
     _serializedExplicitRefineOptions;
     constructor(config) {
+        this._rlmRootAdmission = config.rlmRootAdmission;
+        this._rlmRootAdmission?.bind(this);
+        this._rlmSubagentCapacity =
+            config.rlmChildAdmission?.parent._rlmSubagentCapacity ??
+                config.rlmSubagentCapacity ??
+                new LocalRlmSubagentCapacity(config.settingsManager);
         this._rlmParentAdmission = config.rlmChildAdmission;
         this._rlmParentAdmission?.bind(this);
         this.agent = config.agent;
@@ -568,7 +576,19 @@ export class AgentSession {
                 maxEntries: limits.maxMessages,
                 maxSourceBytes: limits.maxSourceBytes,
             });
-            const captured = this.requests.capture(compaction);
+            let captured;
+            try {
+                captured = this.requests.capture(compaction);
+            }
+            catch (error) {
+                try {
+                    await compaction.release();
+                }
+                catch (cleanupError) {
+                    throw new AggregateError([error, cleanupError], "Canonical context capture and release failed");
+                }
+                throw error;
+            }
             try {
                 const messages = await captured.readHistory(async (view) => {
                     const sameSource = controls?.sessionId === view.source.sessionId && controls?.sessionFile === view.source.sessionFile;
@@ -836,7 +856,6 @@ export class AgentSession {
                     });
                 }
                 this._retryAttempt = 0;
-                this._retryAuthFailureSources = [];
                 this._resolveRetry();
             },
         });
@@ -3673,6 +3692,13 @@ export class AgentSession {
                     record.started = true;
             }
         }
+        if (event.type === "message_end" && event.message.role === "assistant") {
+            // Native result settlement precedes message publication. Keep auth failure
+            // ownership independent of later message replacement, cancellation, or retry cleanup.
+            const source = takeNativeInferenceAuthSource(event.message);
+            if (source && this._isConcreteProviderAuthFailure(event.message))
+                this._markProviderAuthStale(source);
+        }
         const compactionOwner = event.type === "agent_end" ? this._invocationCompactionOwner : undefined;
         const job = this._agentEventQueue.then(() => this._processAgentEvent(event, nativeMessageWrite, compactionOwner), () => this._processAgentEvent(event, nativeMessageWrite, compactionOwner));
         this._agentEventQueue = job;
@@ -3754,7 +3780,6 @@ export class AgentSession {
                 });
                 this._retryAttempt = 0;
             }
-            this._retryAuthFailureSources = [];
             this._resolveRetry();
             await this._emitExtensionEvent(event);
             this._emit(event);
@@ -3849,9 +3874,6 @@ export class AgentSession {
                 if (assistantMsg.stopReason !== "error") {
                     this._overflowRecovery = "idle";
                 }
-                if (this._isConcreteProviderAuthFailure(assistantMsg)) {
-                    this._captureRetryAuthFailureSource(assistantMsg);
-                }
                 // Reset retry counter immediately on successful assistant response
                 // This prevents accumulation across multiple LLM calls within a turn
                 if (assistantMsg.stopReason !== "error" &&
@@ -3863,7 +3885,6 @@ export class AgentSession {
                         attempt: this._retryAttempt,
                     });
                     this._retryAttempt = 0;
-                    this._retryAuthFailureSources = [];
                 }
                 if ((await this._accountGoalUsageForAssistantMessage(assistantMsg)) && !this._invocationOutputRefused) {
                     const message = createGoalContextMessage(this._goalState, "budget_limit");
@@ -3886,12 +3907,6 @@ export class AgentSession {
                 this._resolveRetry();
                 return;
             }
-            if (this._isConcreteProviderAuthFailure(msg)) {
-                this._markProviderAuthStaleForRetryFailure(msg, {
-                    markAuthStaleOnFailure: true,
-                    authSourceTokens: this._retryAuthFailureSources,
-                });
-            }
             const compactionWillRetry = await this._checkCompaction(msg, true, true, compactionOwner);
             if (compactionWillRetry && this._retryAttempt > 0) {
                 return;
@@ -4087,7 +4102,8 @@ export class AgentSession {
      * the latest state reaches disk instead of racing process exit.
      */
     async disposeAsync(options) {
-        this._rlmChildAdmission?.cancel("Parent session disposed");
+        for (const admission of this._rlmChildAdmissions)
+            admission.cancel("Parent session disposed");
         if (this._disposeAsyncPromise)
             return this._disposeAsyncPromise;
         this.closeAutoRefineAdmission();
@@ -4114,7 +4130,8 @@ export class AgentSession {
             if (errors.length > 1)
                 throw new AggregateError(errors, "Session disposal failed");
             this._rlmResidentDisposalComplete = true;
-            this._releaseRlmResidentCapacity?.();
+            await this._releaseRlmResidentCapacity?.();
+            await this._rlmRootAdmission?.release();
         })();
         return this._disposeAsyncPromise;
     }
@@ -4211,7 +4228,7 @@ export class AgentSession {
             throw new AggregateError(distinct, "Session refinement drain failed");
     }
     async _disposeAsyncOnce(kernelSnapshot) {
-        const admission = this._rlmChildAdmission;
+        const admissions = [...this._rlmChildAdmissions];
         const errors = [];
         const drain = async (operation) => {
             try {
@@ -4232,10 +4249,11 @@ export class AgentSession {
             // Join it before disposing any runtime resources it may still be constructing.
             if (this._initialization)
                 await Promise.allSettled([this._initialization]);
-            const admittedChild = admission?.session;
-            // Includes a native child constructed but not yet published by its factory.
-            if (admittedChild)
-                await drain(() => admittedChild.disposeAsync());
+            // Includes native children constructed but not yet published by their factories.
+            for (const admission of admissions) {
+                if (admission.session)
+                    await drain(() => admission.session?.disposeAsync());
+            }
             for (const run of [...this._activeRlmChildRuns.values()]) {
                 const childSession = run.session;
                 if (!childSession)
@@ -4267,9 +4285,10 @@ export class AgentSession {
                 for (const task of [...this._rlmRunTasks])
                     await drain(() => task);
             }
-            await drain(() => admission?.settlement);
-            if (admission?.session && admission.session !== admittedChild) {
-                await drain(() => admission.session?.disposeAsync());
+            for (const admission of admissions) {
+                await drain(() => admission.settlement);
+                if (admission.session)
+                    await drain(() => admission.session?.disposeAsync());
             }
             await drain(() => this.agent.waitForIdle());
             await drain(() => this._compactionOperation);
@@ -6938,9 +6957,9 @@ export class AgentSession {
         const current = () => !owner || this._isCompactionOwnerCurrent(owner);
         while (current() && (settlement === undefined || this._postCompactionContinuationSettlement === settlement)) {
             if (this._actionStore.queuedActions().length > 0) {
-                if (this._sessionInputPumpSuspended ||
-                    this._queuedWorkPauses.size > 0 ||
-                    (settlement && this._sessionInputAdmissionPauses.size > 0)) {
+                // A blocked pump must wait for a checkpoint change, not spin in
+                // microtasks and starve the IO that clears its busy state.
+                if (this._isBusyForSessionInput("pump") || (settlement && this._sessionInputAdmissionPauses.size > 0)) {
                     let wake = () => { };
                     const changed = new Promise((resolve) => {
                         wake = resolve;
@@ -7135,7 +7154,7 @@ export class AgentSession {
             throw new Error(`No API key for ${model.provider}/${model.id}`);
         }
         if (!(await this._modelRegistry.canUseModel(model))) {
-            throw new Error(`Model "${model.provider}/${model.id}" is not available for the current Prime team.`);
+            throw new Error(`Model "${model.provider}/${model.id}" is not available with the configured authentication.`);
         }
         if (this.sessionManager.getSessionId() !== sourceId || this.sessionManager.getSessionFile() !== sourceFile) {
             throw new Error("Session source changed during model selection");
@@ -9850,25 +9869,27 @@ export class AgentSession {
     get hasRlmParentAdmission() {
         return this._releaseRlmResidentCapacity !== undefined;
     }
-    /** Reserve before setup awaits. The concrete parent lifetime owns this one slot. */
-    reserveRlmChildAdmission() {
+    /** Reserve before setup awaits. Capacity remains owned until real cleanup completes. */
+    async reserveRlmChildAdmission() {
         if (this._disposed || this._disposing || this._disposeAsyncPromise) {
             throw new Error("Cannot spawn a subagent after its parent was disposed");
         }
-        if (this._rlmChildAdmission) {
-            throw new Error("RLM resident child limit reached (one child per parent); dispose or passivate the existing child first");
-        }
         const controller = new AbortController();
         const settlement = createAgentMessageDeferred();
         let child;
         let settled = false;
         let factoryClaimed = false;
         let unboundCleanupComplete = true;
-        const release = () => {
+        let capacity;
+        let capacityRelease;
+        const release = async () => {
             if (!settled || (child ? !child._rlmResidentDisposalComplete : !unboundCleanupComplete))
                 return;
-            if (this._rlmChildAdmission === admission)
-                this._rlmChildAdmission = undefined;
+            if (capacity) {
+                capacityRelease ??= capacity.release();
+                await capacityRelease;
+            }
+            this._rlmChildAdmissions.delete(admission);
             if (child?._releaseRlmResidentCapacity === release)
                 child._releaseRlmResidentCapacity = undefined;
             if (child?._rlmParentAdmission === admission)
@@ -9888,20 +9909,21 @@ export class AgentSession {
                 if (this._disposed ||
                     this._disposing ||
                     this._disposeAsyncPromise ||
-                    this._rlmChildAdmission !== admission) {
+                    !this._rlmChildAdmissions.has(admission)) {
                     throw new Error("RLM child admission is no longer current");
                 }
             },
             bind: (session) => {
-                if (this._rlmChildAdmission !== admission)
+                if (!this._rlmChildAdmissions.has(admission))
                     throw new Error("RLM child admission is no longer current");
                 if ((child && child !== session) ||
                     (session._releaseRlmResidentCapacity && session._releaseRlmResidentCapacity !== release)) {
                     throw new Error("RLM child already belongs to another resident admission");
                 }
                 child = session;
+                session._rlmParentAdmission = admission;
+                session._rlmSubagentCapacity = this._rlmSubagentCapacity;
                 session._releaseRlmResidentCapacity = release;
-                release();
             },
             beginSetup: () => {
                 if (settled || child)
@@ -9914,14 +9936,18 @@ export class AgentSession {
                 admission.beginSetup();
                 factoryClaimed = true;
             },
-            confirmUnboundCleanup: () => {
+            confirmUnboundCleanup: async () => {
                 unboundCleanupComplete = true;
-                release();
+                await release();
             },
-            settle: () => {
+            settle: async () => {
                 settled = true;
-                settlement.resolve();
-                release();
+                try {
+                    await release();
+                }
+                finally {
+                    settlement.resolve();
+                }
             },
             cancel: (reason) => {
                 if (settled || controller.signal.aborted)
@@ -9931,8 +9957,16 @@ export class AgentSession {
                 return true;
             },
         };
-        this._rlmChildAdmission = admission;
-        return admission;
+        this._rlmChildAdmissions.add(admission);
+        try {
+            capacity = await this._rlmSubagentCapacity.reserve();
+            admission.assertCurrent();
+            return admission;
+        }
+        catch (error) {
+            await admission.settle();
+            throw error;
+        }
     }
     _createRlmSubagentRuntimeOptions(options) {
         const requestTokenBudget = this.requests.getRequestTokenBudgetOptions();
@@ -10046,7 +10080,7 @@ export class AgentSession {
                     await failedChild.disposeAsync();
                 else {
                     await childSessionManager.close();
-                    options.admission.confirmUnboundCleanup();
+                    await options.admission.confirmUnboundCleanup();
                 }
             }
             catch (cleanupError) {
@@ -10066,7 +10100,8 @@ export class AgentSession {
         this._scheduleGoalContinuationAfterRlmWork();
     }
     _cancelActiveRlmChildRuns(reason) {
-        this._rlmChildAdmission?.cancel(reason);
+        for (const admission of this._rlmChildAdmissions)
+            admission.cancel(reason);
         for (const run of this._activeRlmChildRuns.values()) {
             this._cancelRlmChildRun(run, reason);
         }
@@ -10488,15 +10523,15 @@ export class AgentSession {
             return false;
         }
         if (!session._rlmResidentDisposalComplete) {
-            if (this._rlmChildAdmission)
-                this._rlmChildAdmission.bind(session);
+            if (session._rlmParentAdmission?.parent === this)
+                session._rlmParentAdmission.bind(session);
             else {
-                const admission = this.reserveRlmChildAdmission();
+                const admission = await this.reserveRlmChildAdmission();
                 try {
                     admission.bind(session);
                 }
                 finally {
-                    admission.settle();
+                    await admission.settle();
                 }
             }
         }
@@ -10624,8 +10659,9 @@ export class AgentSession {
     /** True when any direct or nested subagent is still running or queued. */
     hasRunningRlmChildren() {
         for (const session of this._rlmSubtreeSessions()) {
-            if (session._rlmChildAdmission?.pending)
-                return true;
+            for (const admission of session._rlmChildAdmissions)
+                if (admission.pending)
+                    return true;
             for (const run of session._activeRlmChildRuns.values()) {
                 if (run.status === "running" || run.status === "queued") {
                     return true;
@@ -10773,8 +10809,9 @@ export class AgentSession {
     cancelRunningRlmDescendants(reason = "Cancelled by user") {
         let cancelled = false;
         for (const session of this._rlmSubtreeSessions()) {
-            if (session._rlmChildAdmission?.cancel(reason))
-                cancelled = true;
+            for (const admission of session._rlmChildAdmissions)
+                if (admission.cancel(reason))
+                    cancelled = true;
             for (const run of session._activeRlmChildRuns.values()) {
                 if (session._cancelRlmChildRun(run, reason))
                     cancelled = true;
@@ -11002,7 +11039,7 @@ export class AgentSession {
                     admission.bind(childRuntime.session);
                 }
                 finally {
-                    admission.settle();
+                    await admission.settle();
                 }
                 admission.assertCurrent();
                 const child = childRuntime.session;
@@ -11289,7 +11326,7 @@ export class AgentSession {
         };
     }
     async runRlmChild(prompt, kwargs = {}, spawnCode) {
-        const admission = this.reserveRlmChildAdmission();
+        const admission = await this.reserveRlmChildAdmission();
         const start = this._startRlmChildRun(prompt, admission, kwargs, spawnCode);
         // Join pre-runtime name/model work during disposal as well as the detached task.
         const pending = start.then(() => undefined, () => undefined);
@@ -11298,7 +11335,7 @@ export class AgentSession {
             return await start;
         }
         catch (error) {
-            admission.settle();
+            await admission.settle();
             throw error;
         }
         finally {
@@ -11348,56 +11385,15 @@ export class AgentSession {
         return (/\b(?:401|403)\b/.test(message.errorMessage) &&
             /auth|unauthori[sz]ed|forbidden|api.?key|token|credential/i.test(message.errorMessage));
     }
-    _captureRetryAuthFailureSource(message) {
-        const token = this._modelRegistry.getCurrentProviderAuthSourceToken(message.provider);
-        if (!token) {
-            return undefined;
-        }
-        if (!this._retryAuthFailureSources.some((existing) => existing.provider === token.provider &&
-            existing.source === token.source &&
-            existing.identityFingerprint === token.identityFingerprint &&
-            existing.valueFingerprint === token.valueFingerprint)) {
-            this._retryAuthFailureSources.push(token);
-        }
-        return token;
-    }
-    _markProviderAuthStale(message, authSourceTokens) {
-        if (authSourceTokens && authSourceTokens.length > 0) {
-            let marked = false;
-            for (const token of authSourceTokens) {
-                marked = this._modelRegistry.markProviderAuthSourceStale(token) || marked;
-            }
-            if (marked) {
-                this._emit({
-                    type: "auth_stale",
-                    provider: message.provider,
-                    sourceTokens: authSourceTokens,
-                });
-            }
-            return marked;
-        }
-        const marked = this._modelRegistry.markProviderAuthStale(message.provider);
-        if (marked) {
-            this._emit({ type: "auth_stale", provider: message.provider });
+    _markProviderAuthStale(source) {
+        if (this._modelRegistry.markProviderAuthSourceStale(source)) {
+            this._emit({ type: "auth_stale", provider: source.provider, sourceTokens: [source] });
         }
-        return marked;
-    }
-    _markProviderAuthStaleForRetryFailure(message, options) {
-        const authSourceTokens = this._retryAuthFailureSources.length > 0 ? this._retryAuthFailureSources : options?.authSourceTokens;
-        if ((authSourceTokens?.length ?? 0) > 0 || options?.markAuthStaleOnFailure) {
-            const marked = this._markProviderAuthStale(message, authSourceTokens);
-            if (marked && message.errorMessage) {
-                message.errorMessage = addLoginGuidanceToAuthError(message.errorMessage);
-            }
-            return marked;
-        }
-        return false;
     }
     _finishActiveRetryWithFailure(message) {
         if (this._retryAttempt === 0) {
             return;
         }
-        this._markProviderAuthStaleForRetryFailure(message);
         this._emit({
             type: "auto_retry_end",
             success: false,
@@ -11405,7 +11401,6 @@ export class AgentSession {
             finalError: message.errorMessage,
         });
         this._retryAttempt = 0;
-        this._retryAuthFailureSources = [];
     }
     async _handleRetryableError(message, signal) {
         const settings = this.settingsManager.getRetrySettings();
@@ -11468,7 +11463,6 @@ export class AgentSession {
             });
             this._retryAttempt = 0;
         }
-        this._retryAuthFailureSources = [];
         this._resolveRetry();
     }
     async waitForRetry() {
@@ -11695,6 +11689,12 @@ export class AgentSession {
             this.agent.state.messages.push(bashMessage);
         }
     }
+    getRlmMaxSubagentsStatus() {
+        return this._rlmSubagentCapacity.getStatus();
+    }
+    setRlmMaxSubagents(maxSubagents) {
+        return this._rlmSubagentCapacity.setMaxSubagents(maxSubagents);
+    }
     getRlmMaxDepthStatus() {
         return { maxDepth: this._rlmMaxDepth, source: this._rlmMaxDepthSource };
     }