open-multi-agent-kit 0.98.4 → 0.98.5

package/dist/core/agent-session.js

@@ -103,6 +103,7 @@ import { resolveResourceGovernorSettings } from "./resource-governor-settings.js
 import { admissionObservationFacts, classificationObservationFacts, ResourceObservationJournal, settledObservationFacts, snapshotObservationFacts, soundObservationFacts, } from "./resource-observation-journal.js";
 import { decideResourceSafetyGate, RESOURCE_PRESSURE_REQUIRED_ACTION } from "./resource-safety-gate.js";
 import { appendRouterFeedbackRecord, ROUTER_FEEDBACK_LEVELS, } from "./router-feedback-collector.js";
+import { RunBudgetExceededError, RunBudgetPolicyError } from "./run-budget-policy.js";
 import { RunJournalStore } from "./run-journal-store.js";
 import { RunResourceLeaseController } from "./run-resource-lease.js";
 import { SessionBashRuntime } from "./session-bash-runtime.js";
@@ -111,6 +112,8 @@ import { SessionCompactionService } from "./session-compaction-service.js";
 import { preflightFailureCause, runtimeFailureCause, terminationMessage } from "./session-failure-cause.js";
 import { CURRENT_SESSION_VERSION, getLatestCompactionEntry } from "./session-manager.js";
 import { acquireSessionOwnerLeaseSync } from "./session-owner-lease.js";
+import { PromptExecutionBusyError, SessionPromptLifecycle } from "./session-prompt-lifecycle.js";
+import { SessionRunBudget } from "./session-run-budget.js";
 import { classifyRunTermination } from "./session-run-termination.js";
 import { assembleSessionSystemPrompt } from "./session-system-prompt.js";
 import { classifySessionTermination, } from "./session-termination.js";
@@ -217,6 +220,11 @@ export class AgentSession {
     // Event subscription state
     _unsubscribeAgent;
     _eventListeners = [];
+    _runBudget;
+    _promptLifecycle = new SessionPromptLifecycle({
+        canSettle: () => !this.isStreaming && !this.agent.hasQueuedMessages(),
+        auditsLateSettlement: () => this.agent.toolExecutionPolicy?.lateSettlement !== "ignore",
+    });
     _messageEndReplacements = new WeakMap();
     /** Tracks pending steering messages for UI display. Removed when delivered. */
     _steeringMessages = [];
@@ -313,6 +321,21 @@ export class AgentSession {
     _promptCacheLastBreakReason;
     constructor(config) {
         this.agent = config.agent;
+        this._runBudget = new SessionRunBudget(this.agent, {
+            assertIdle: () => {
+                this._promptLifecycle.assertIdle();
+                if (this.isRetrying || this.isCompacting || this._branchSummaryAbortController)
+                    throw new PromptExecutionBusyError();
+            },
+            stop: (error) => {
+                this._pendingRuntimeTerminationCause = { area: "budget", code: error.code };
+                this.agent.abort();
+                this.abortRetry();
+                this.abortCompaction();
+                this.abortBranchSummary();
+            },
+            reject: (error) => this._publishRuntimeFailure(error),
+        });
         this.sessionManager = config.sessionManager;
         this.settingsManager = config.settingsManager;
         this._scopedModels = config.scopedModels ?? [];
@@ -688,6 +711,7 @@ export class AgentSession {
         throw new Error(formatNoApiKeyFoundMessage(model.provider));
     }
     async _getCompactionRequestAuth(model) {
+        this._runBudget.assertAdmission();
         const options = { minRemainingMs: COMPACTION_MIN_TOKEN_VALIDITY_MS };
         if (isBuiltinStreamFn(this.agent.streamFn)) {
             return this._getRequiredRequestAuth(model, options);
@@ -842,6 +866,8 @@ export class AgentSession {
                 }
             }
         }
+        if (event.type === "tool_execution_late_settlement")
+            this._promptLifecycle.flush();
     };
     _willRetryAfterAgentEnd(event) {
         const settings = this.settingsManager.getRetrySettings();
@@ -1109,6 +1135,8 @@ export class AgentSession {
      * Call this when completely done with the session.
      */
     dispose() {
+        this._promptLifecycle.dispose();
+        this._runBudget.close();
         try {
             try {
                 this.abortRetry();
@@ -1235,7 +1263,7 @@ export class AgentSession {
         for (const name of desiredToolNames) {
             const tool = this._toolRegistry.get(name);
             if (tool) {
-                tools.push(tool);
+                tools.push(this._promptLifecycle.wrapTool(tool));
                 validToolNames.push(name);
             }
             else if (lockedToolNames?.includes(name)) {
@@ -1460,13 +1488,14 @@ export class AgentSession {
         // One identity per top-level run (§16.2): internal retries and
         // continuations inside this call share it, as does the resource lease.
         const promptRunId = `prompt-run-${randomUUID()}`;
-        const startedAtEpochMs = Date.now();
+        const ownedRun = this._promptLifecycle.begin(promptRunId);
         // Roadmap M2: the lease spans the whole run including internal retries
         // and continuations, so they share one admission decision (§8.3).
         const resourceLease = await this._beginResourceGovernedRun(promptRunId);
         const resourceObservations = this._resourceObservationJournals.get(promptRunId) ?? null;
         let outcome = "completed";
         try {
+            this._runBudget.assertActive();
             await this.agent.prompt(messages);
             while (await this._handlePostAgentRun()) {
                 await this.agent.continue();
@@ -1478,43 +1507,18 @@ export class AgentSession {
             throw error;
         }
         finally {
+            if (this._runBudget.failure && this._lastTermination?.kind !== "budget_exhausted")
+                this._publishRuntimeFailure(this._runBudget.failure);
             outcome = promptSettlement.resolvePromptSettlementOutcome(outcome, this._lastTermination?.kind);
-            if (resourceLease !== null) {
-                // Generation-safe: a stale release after another run acquired the
-                // lease is a no-op instead of clobbering the newer cap (§8.1).
-                this._resourceLeaseController?.release(resourceLease);
-                resourceObservations?.record("resource_lease_released_v1", { promptRunId });
-            }
             this._flushPendingBashMessages();
-            // §16.4 (M4): settle only after every in-run continuation drained.
-            this._emitPromptSettledIfReady({ promptRunId, startedAtEpochMs, outcome, resourceObservations });
-        }
-    }
-    /**
-     * Emit `prompt_settled` when the §16.4 conditions hold. Called exactly
-     * once per promptRunId (from `_runAgentPrompt`'s finally), so the same
-     * run can never emit twice; late-queued messages block emission
-     * conservatively — the continuation run settles under its own id.
-     * Settlement is a UX signal only and must never affect the run (§16.1).
-     */
-    _emitPromptSettledIfReady(input) {
-        try {
-            let state = promptSettlement.createPromptSettlementState(input.promptRunId, input.startedAtEpochMs);
-            state = promptSettlement.reducePromptSettlement(state, { kind: "terminal", outcome: input.outcome });
-            if (this.isStreaming) {
-                state = promptSettlement.reducePromptSettlement(state, { kind: "tool", delta: 1 });
-            }
-            if (this.agent.hasQueuedMessages()) {
-                state = promptSettlement.reducePromptSettlement(state, { kind: "continuation", delta: 1 });
-            }
-            const settled = promptSettlement.settlePromptIfReady(state, Date.now());
-            if (settled.event !== null) {
-                input.resourceObservations?.record("prompt_settled_v1", settledObservationFacts(settled.event));
-                this._emit(settled.event);
-            }
-        }
-        catch {
-            // Never let settlement bookkeeping break a run.
+            ownedRun.finish(outcome, (event) => {
+                if (resourceLease !== null) {
+                    this._resourceLeaseController?.release(resourceLease);
+                    resourceObservations?.record("resource_lease_released_v1", { promptRunId });
+                }
+                resourceObservations?.record("prompt_settled_v1", settledObservationFacts(event));
+                this._emit(event);
+            });
         }
     }
     /**
@@ -1526,7 +1530,7 @@ export class AgentSession {
      *   delayed. The pending probe is bounded (~cpuSampleMs) and never rejects.
      * - `adaptive`/`strict`: bounded blocking probe (default 300 ms deadline),
      *   then a generation-safe lease applies the effective tool cap for this
-     *   run. The caller's `finally` releases it exactly once.
+     *   run. Prompt settlement releases it after owned tool promises terminate.
      *
      * Never throws: any probe or policy failure leaves this run ungoverned
      * (§2.1: probe failure must not crash or block the prompt).
@@ -1721,6 +1725,7 @@ export class AgentSession {
         }
     }
     async _handlePostAgentRun() {
+        this._runBudget.assertActive();
         const msg = this._lastAssistantMessage;
         this._lastAssistantMessage = undefined;
         if (!msg) {
@@ -1755,7 +1760,15 @@ export class AgentSession {
      * @throws Error if streaming and no streamingBehavior specified
      * @throws Error if no model selected or no API key available (when not streaming)
      */
-    async prompt(text, options) {
+    prompt(text, options) {
+        if (options?.runBudget === undefined && (this.isStreaming || this.isRetrying))
+            return this._prompt(text, options);
+        return this._runBudget.execute(options?.runBudget, () => this._prompt(text, options), options?.preflightResult);
+    }
+    getRunBudgetSnapshot() {
+        return this._runBudget.snapshot();
+    }
+    async _prompt(text, options) {
         const expandPromptTemplates = options?.expandPromptTemplates ?? true;
         const preflightResult = options?.preflightResult;
         let currentText = redactSensitiveText(text);
@@ -1820,6 +1833,7 @@ export class AgentSession {
                 preflightResult?.(true);
                 return;
             }
+            this._promptLifecycle.assertIdle();
             // Flush any pending bash messages before the new prompt
             this._flushPendingBashMessages();
             // Validate model
@@ -1918,11 +1932,16 @@ export class AgentSession {
             }
             this._recordPromptCachePlan(result?.systemPrompt ? "extension-override" : "turn-plan");
             await this._checkProjectedCompaction(messages);
+            this._runBudget.assertActive();
         }
         catch (error) {
             preflightResult?.(false);
             const rawMessage = error instanceof Error ? error.message : String(error);
-            const cause = preflightFailureCause(rawMessage, Boolean(this.model));
+            const cause = error instanceof RunBudgetExceededError
+                ? { area: "budget", code: error.code }
+                : error instanceof RunBudgetPolicyError
+                    ? { area: "configuration", code: "invalid" }
+                    : preflightFailureCause(rawMessage, Boolean(this.model));
             const timestamp = new Date().toISOString();
             this._publishTermination(classifySessionTermination({
                 sessionId: this.sessionId,
@@ -2172,6 +2191,7 @@ export class AgentSession {
         this._followUpMessages = [];
         this.agent.clearAllQueues();
         this._emitQueueUpdate();
+        this._promptLifecycle.flush();
         return { steering, followUp };
     }
     /** Number of pending messages (includes both steering and follow-up) */
@@ -3938,7 +3958,7 @@ export class AgentSession {
         const failoverTo = await this._maybeFailoverFromSafetyStop(message);
         const switchedVia = failoverTo ? `failover ${failoverTo}` : rotatedTo ? `route ${rotatedTo}` : undefined;
         // Safety stops are usually immediate false positives — short delay after failover, full backoff otherwise.
-        const delayMs = computeRetryDelayMs(settings.baseDelayMs, attempt, switchedVia !== undefined);
+        const delayMs = Math.min(computeRetryDelayMs(settings.baseDelayMs, attempt, switchedVia !== undefined), this._runBudget.remainingMs ?? Infinity);
         const errorMessage = switchedVia
             ? `${message.errorMessage || "content/safety stop"} → ${switchedVia}`
             : message.errorMessage || "Unknown error";