chainlesschain 0.165.3 → 0.165.4

package/src/gateways/ws/ws-agent-handler.js

@@ -62,6 +62,13 @@ import {
   readVerifiedMessages,
   settleWsTurnClaim,
 } from "../../harness/jsonl-session-store.js";
+import { sessionBudgetAdmissionError } from "../../lib/session-budget-production-root.js";
+import {
+  beginSessionBudgetUsage,
+  markSessionBudgetUsageUnknown,
+  recordSessionBudgetUsage,
+  rejectSessionBudgetUsageUnknown,
+} from "../../lib/session-budget-usage.js";
 
 const CANONICAL_WS_TURN_QUEUES = new Map();
 const CANONICAL_WS_CLAIM_CAS_ATTEMPTS = 4;
@@ -155,6 +162,7 @@ export class WSAgentHandler {
     compactionChatFn,
     canonicalSessionStore,
     sessionHostLease = null,
+    sessionBudgetRoot = null,
   }) {
     this.session = session;
     this.interaction = interaction;
@@ -163,6 +171,8 @@ export class WSAgentHandler {
     this._abortController = null;
     this._activeRequestId = null;
     this._sessionHostLease = sessionHostLease;
+    this._sessionBudgetRoot = sessionBudgetRoot;
+    this._sessionBudget = sessionBudgetRoot?.budget || null;
     this._onSessionHostLeaseAbort = null;
     if (sessionHostLease?.signal) {
       this._onSessionHostLeaseAbort = () => {
@@ -181,6 +191,26 @@ export class WSAgentHandler {
       );
       if (sessionHostLease.signal.aborted) this._onSessionHostLeaseAbort();
     }
+    this._sessionBudgetSignal = sessionBudgetRoot?.options?.signal || null;
+    this._onSessionBudgetAbort = null;
+    if (this._sessionBudgetSignal) {
+      this._onSessionBudgetAbort = () => {
+        const reason = sessionBudgetAdmissionError(
+          this._sessionBudget?.reason?.(),
+          "WebSocket session",
+        );
+        if (this._abortController && !this._abortController.signal.aborted) {
+          this._abortController.abort(reason);
+        }
+        this.interaction?.rejectAllPending?.(reason);
+      };
+      this._sessionBudgetSignal.addEventListener(
+        "abort",
+        this._onSessionBudgetAbort,
+        { once: true },
+      );
+      if (this._sessionBudgetSignal.aborted) this._onSessionBudgetAbort();
+    }
     // P0-2: monotonic sequence for crash-safe side-effect ledger op ids. The
     // session id + a per-op nonce keep ids unique across turns and processes.
     this._sideEffectSeq = 0;
@@ -251,6 +281,58 @@ export class WSAgentHandler {
       : Object.freeze(hostSystemPrefix);
   }
 
+  attachInteraction(interaction) {
+    if (!interaction || interaction === this.interaction) return false;
+    const previous = this.interaction;
+    this.interaction = interaction;
+    this.session.interaction = interaction;
+    previous?.rejectAllPending?.(
+      createAbortError("WebSocket session transport was reattached"),
+    );
+    return true;
+  }
+
+  assertSessionBudgetAdmission(operation = "WebSocket operation") {
+    if (!this._sessionBudget) return true;
+    if (
+      this._sessionBudget.signal?.aborted === true ||
+      this._sessionBudgetSignal?.aborted === true
+    ) {
+      throw sessionBudgetAdmissionError(
+        this._sessionBudget.reason?.(),
+        operation,
+      );
+    }
+    return true;
+  }
+
+  _turnSignal(turnSignal) {
+    const budgetSignal = this._sessionBudgetSignal;
+    if (!budgetSignal) return turnSignal;
+    return AbortSignal.any([turnSignal, budgetSignal]);
+  }
+
+  _recordSessionBudgetUsage(event, operation = "WebSocket usage settlement") {
+    return recordSessionBudgetUsage(this._sessionBudget, event, operation);
+  }
+
+  _beginSessionBudgetUsage(event, operation = "WebSocket provider call") {
+    return beginSessionBudgetUsage(this._sessionBudget, event, operation);
+  }
+
+  _markSessionBudgetUsageUnknown(event) {
+    return markSessionBudgetUsageUnknown(this._sessionBudget, event);
+  }
+
+  _consumeSessionBudgetTurn(id, operation) {
+    if (!this._sessionBudget?.consumeTurn) return null;
+    const admission = this._sessionBudget.consumeTurn({ id });
+    if (!admission?.ok) {
+      throw sessionBudgetAdmissionError(admission?.reason, operation);
+    }
+    return admission;
+  }
+
   _applyCanonicalMessages(messages) {
     const canonicalMessages = Array.isArray(messages) ? messages : [];
     this.session.messages = [
@@ -383,7 +465,12 @@ export class WSAgentHandler {
       model: this.session.model || null,
       source: "semantic-compaction",
     });
+    this._consumeSessionBudgetTurn(
+      `turn:${call.callId}`,
+      "WebSocket compaction",
+    );
     this._persistCanonicalUsageBoundary(call, "started");
+    this._beginSessionBudgetUsage(call, "WebSocket compaction provider call");
     return call;
   }
 
@@ -402,19 +489,36 @@ export class WSAgentHandler {
         },
         "unknown",
       );
+      if (this._markSessionBudgetUsageUnknown(call)) {
+        rejectSessionBudgetUsageUnknown(
+          call,
+          "WebSocket compaction provider call",
+        );
+      }
       return;
     }
     if (result?.usageEvent?.usage) {
-      this._persistCanonicalTokenUsage({
+      const usageEvent = {
         ...call,
         usage: result.usageEvent.usage,
-      });
+      };
+      this._persistCanonicalTokenUsage(usageEvent);
+      this._recordSessionBudgetUsage(
+        usageEvent,
+        "WebSocket compaction usage settlement",
+      );
       return;
     }
     this._persistCanonicalUsageBoundary(
       { ...call, code: "provider_usage_missing" },
       "unknown",
     );
+    if (this._markSessionBudgetUsageUnknown(call)) {
+      rejectSessionBudgetUsageUnknown(
+        call,
+        "WebSocket compaction provider call",
+      );
+    }
   }
 
   _settleCanonicalCompactionFailure(call) {
@@ -423,6 +527,7 @@ export class WSAgentHandler {
       { ...call, code: "provider_call_failed" },
       "unknown",
     );
+    this._markSessionBudgetUsageUnknown(call);
   }
 
   _latchCompactionSettlement({
@@ -511,7 +616,7 @@ export class WSAgentHandler {
     return null;
   }
 
-  async _compactCanonicalHistoryBeforeClaim(requestId) {
+  async _compactCanonicalHistoryBeforeClaim(requestId, signal = null) {
     const { session } = this;
     if (
       session.autoCompact === false ||
@@ -525,6 +630,8 @@ export class WSAgentHandler {
     const expectedMessages = [...session.messages];
     const strictUsageTelemetry =
       this._resolveCanonicalUsageProtocol() === "call-ledger";
+    const meterCompaction =
+      strictUsageTelemetry || Boolean(this._sessionBudget);
     let compactionUsageCall = null;
     let compactionUsageSettled = false;
     let result;
@@ -535,7 +642,7 @@ export class WSAgentHandler {
         model: session.model,
         baseUrl: session.baseUrl,
         apiKey: session.apiKey,
-        signal: this._sessionHostLease?.signal,
+        signal: signal || this._sessionBudgetSignal || this._sessionHostLease?.signal,
         onlyIfNeeded: true,
         preserveCompletedExchange: true,
         llmQuery: this._compactionLlmQuery,
@@ -544,7 +651,7 @@ export class WSAgentHandler {
           cwd: session.projectRoot,
           sessionId: session.id,
         },
-        ...(strictUsageTelemetry
+        ...(meterCompaction
           ? {
               onProviderCallStart: () => {
                 const call = this._beginCanonicalCompactionUsage();
@@ -1012,6 +1119,7 @@ export class WSAgentHandler {
 
     this._processing = true;
     const abortController = new AbortController();
+    const turnSignal = this._turnSignal(abortController.signal);
     this._abortController = abortController;
     this._activeRequestId = requestId || null;
     let sideEffectLedger = null;
@@ -1022,6 +1130,7 @@ export class WSAgentHandler {
 
     try {
       const { session } = this;
+      this.assertSessionBudgetAdmission("WebSocket turn");
       this._sessionHostLease?.assert?.();
 
       if (session.canonicalJsonlSession === true) {
@@ -1044,7 +1153,10 @@ export class WSAgentHandler {
             this._resolveCanonicalUsageProtocol() === "call-ledger";
           if (strictUsageTelemetry) this._assertUsageLedgerWritable();
           const compacted =
-            await this._compactCanonicalHistoryBeforeClaim(requestId);
+            await this._compactCanonicalHistoryBeforeClaim(
+              requestId,
+              turnSignal,
+            );
           if (!compacted) return;
         }
         canonicalTurn ||= this._claimCanonicalTurn(userMessage, requestId);
@@ -1143,7 +1255,8 @@ export class WSAgentHandler {
         shellPolicyOverrides: session.shellPolicyOverrides || null,
         slotFiller,
         interaction: this.interaction,
-        signal: abortController.signal,
+        signal: turnSignal,
+        sessionBudget: this._sessionBudget,
         // P0 authority: null unless opted in (byte-identical default — agent
         // core already defaults `options.approvalGate || null`).
         approvalGate: await this._ensureApprovalGate(),
@@ -1374,18 +1487,28 @@ export class WSAgentHandler {
                 this._persistCanonicalUsageBoundary(event, "unknown");
               }
               this._emitCompactionUsageUnknown(requestId, event);
+              if (this._markSessionBudgetUsageUnknown(event)) {
+                rejectSessionBudgetUsageUnknown(
+                  event,
+                  "WebSocket semantic compaction",
+                );
+              }
               break;
 
             case "model-usage-started":
               if (strictUsageTelemetry && event.ledgerPersisted !== true) {
                 this._persistCanonicalUsageBoundary(event, "started");
               }
+              this._beginSessionBudgetUsage(event);
               break;
 
             case "model-usage-unknown":
               if (strictUsageTelemetry && event.ledgerPersisted !== true) {
                 this._persistCanonicalUsageBoundary(event, "unknown");
               }
+              if (this._markSessionBudgetUsageUnknown(event)) {
+                rejectSessionBudgetUsageUnknown(event, "WebSocket provider call");
+              }
               break;
 
             case "compaction":
@@ -1396,9 +1519,14 @@ export class WSAgentHandler {
               break;
 
             case "token-usage":
-              if (strictUsageTelemetry && event.ledgerPersisted !== true) {
+              if (
+                session.canonicalJsonlSession === true &&
+                (strictUsageTelemetry || this._sessionBudget) &&
+                event.ledgerPersisted !== true
+              ) {
                 this._persistCanonicalTokenUsage(event);
               }
+              this._recordSessionBudgetUsage(event);
               this.interaction.emit("token-usage", {
                 requestId,
                 provider: event.provider,
@@ -1473,13 +1601,29 @@ export class WSAgentHandler {
       session.lastActivity = new Date().toISOString();
     } catch (err) {
       let surfacedError = err;
+      const sessionBudgetFailed =
+        err?.code === "CC_SESSION_BUDGET_EXHAUSTED" ||
+        err?.code === "CC_SESSION_BUDGET_USAGE_UNKNOWN" ||
+        this._sessionBudget?.signal?.aborted === true;
+      if (
+        sessionBudgetFailed &&
+        err?.code !== "CC_SESSION_BUDGET_EXHAUSTED" &&
+        err?.code !== "CC_SESSION_BUDGET_USAGE_UNKNOWN"
+      ) {
+        surfacedError = sessionBudgetAdmissionError(
+          this._sessionBudget?.reason?.(),
+          "WebSocket turn",
+        );
+      }
       if (
         canonicalTurn?.claimAcquired &&
         !canonicalTurn.settled &&
         !canonicalTurn.settlementAttempted
       ) {
-        const failureCode = isAbortError(err)
-          ? "CC_WS_TURN_INTERRUPTED"
+        const failureCode = sessionBudgetFailed
+          ? surfacedError?.code || "CC_SESSION_BUDGET_EXHAUSTED"
+          : isAbortError(err)
+            ? "CC_WS_TURN_INTERRUPTED"
           : err?.code === "CC_WS_EMPTY_ASSISTANT_RESPONSE"
             ? "CC_WS_EMPTY_ASSISTANT_RESPONSE"
             : "CC_WS_TURN_FAILED";
@@ -1521,7 +1665,10 @@ export class WSAgentHandler {
       ) {
         persistSideEffectLedger(this.session.id, sideEffectLedger);
       }
-      if (isAbortError(err) || abortController.signal.aborted) {
+      if (
+        !sessionBudgetFailed &&
+        (isAbortError(err) || abortController.signal.aborted)
+      ) {
         return;
       }
 
@@ -1609,8 +1756,34 @@ export class WSAgentHandler {
       );
     }
     this._onSessionHostLeaseAbort = null;
-    this._sessionHostLease?.release?.();
+    if (this._sessionBudgetSignal && this._onSessionBudgetAbort) {
+      this._sessionBudgetSignal.removeEventListener(
+        "abort",
+        this._onSessionBudgetAbort,
+      );
+    }
+    this._onSessionBudgetAbort = null;
+    this._sessionBudgetSignal = null;
+    const cleanupErrors = [];
+    try {
+      this._sessionBudgetRoot?.close?.();
+    } catch (error) {
+      cleanupErrors.push(error);
+    }
+    this._sessionBudgetRoot = null;
+    this._sessionBudget = null;
+    try {
+      this._sessionHostLease?.release?.();
+    } catch (error) {
+      cleanupErrors.push(error);
+    }
     this._sessionHostLease = null;
+    if (cleanupErrors.length > 0) {
+      throw new AggregateError(
+        cleanupErrors,
+        "WebSocket session authority cleanup failed",
+      );
+    }
   }
 
   /**
@@ -1727,6 +1900,8 @@ export class WSAgentHandler {
             break;
           }
         }
+        const meterCompaction =
+          strictUsageTelemetry || Boolean(this._sessionBudget);
 
         const before = session.messages.length;
         const expectedMessages = [...session.messages];
@@ -1735,13 +1910,15 @@ export class WSAgentHandler {
         let compactionUsageCall = null;
         let compactionUsageSettled = false;
         try {
+          this.assertSessionBudgetAdmission("WebSocket compaction");
           this._sessionHostLease?.assert?.();
           result = await compactConversationWithProvider(session.messages, {
             provider: session.provider,
             model: session.model,
             baseUrl: session.baseUrl,
             apiKey: session.apiKey,
-            signal: this._sessionHostLease?.signal,
+            signal:
+              this._sessionBudgetSignal || this._sessionHostLease?.signal,
             force: true,
             preserveCompletedExchange: true,
             llmQuery: this._compactionLlmQuery,
@@ -1750,7 +1927,7 @@ export class WSAgentHandler {
               cwd: session.projectRoot,
               sessionId: session.id,
             },
-            ...(strictUsageTelemetry
+            ...(meterCompaction
               ? {
                   onProviderCallStart: () => {
                     const call = this._beginCanonicalCompactionUsage();

package/src/gateways/ws/ws-server.js

@@ -415,6 +415,7 @@ export class ChainlessChainWSServer extends EventEmitter {
 
   /** Stop the server and clean up */
   async stop() {
+    const cleanupErrors = [];
     if (this._heartbeatTimer) {
       clearInterval(this._heartbeatTimer);
       this._heartbeatTimer = null;
@@ -423,13 +424,17 @@ export class ChainlessChainWSServer extends EventEmitter {
     // Close all session handlers
     for (const [sessionId, handler] of this.sessionHandlers) {
       if (handler && handler.destroy) {
-        handler.destroy();
+        try {
+          handler.destroy();
+        } catch (error) {
+          cleanupErrors.push(error);
+        }
       }
       if (this.sessionManager) {
         try {
           this.sessionManager.closeSession(sessionId);
-        } catch (_err) {
-          // Non-critical
+        } catch (error) {
+          cleanupErrors.push(error);
         }
       }
     }
@@ -477,6 +482,12 @@ export class ChainlessChainWSServer extends EventEmitter {
     }
 
     this.emit("stopped");
+    if (cleanupErrors.length > 0) {
+      throw new AggregateError(
+        cleanupErrors,
+        "WebSocket server session authority cleanup failed",
+      );
+    }
   }
 
   /** @private */

package/src/gateways/ws/ws-session-gateway.js

@@ -32,6 +32,7 @@ import {
   getSession as dbGetSession,
   listSessions as dbListSessions,
   updateSession as dbUpdateSession,
+  deleteSession as dbDeleteSession,
 } from "../../lib/session-manager.js";
 import { buildSystemPrompt } from "../../runtime/agent-core.js";
 import { SubAgentRegistry } from "../../lib/sub-agent-registry.js";
@@ -136,6 +137,7 @@ function bindSessionHooksV2Workspace(
  * @property {string|null} rulesContent
  * @property {string[]} enabledToolNames
  * @property {object|null} hostManagedToolPolicy
+ * @property {object|null} sessionBudgetRoot
  * @property {Array<object>} externalToolDefinitions
  * @property {object} externalToolDescriptors
  * @property {object} externalToolExecutors
@@ -398,6 +400,8 @@ export class WSSessionManager {
    * @param {string} [options.apiKey]
    * @param {string} [options.baseUrl]
    * @param {object} [options.hostManagedToolPolicy]
+   * @param {object} [options.sessionBudgetRoot]
+   * @param {boolean} [options.requireDurable=false]
    * @returns {{ sessionId: string }}
    */
   createSession(options = {}) {
@@ -407,6 +411,14 @@ export class WSSessionManager {
       );
     }
     const sessionId = this._generateId();
+    const requireDurable = options.requireDurable === true;
+    if (requireDurable && !this.db) {
+      const error = new Error(
+        "Durable WebSocket session creation requires a configured database",
+      );
+      error.code = "CC_WS_DURABLE_SESSION_REQUIRED";
+      throw error;
+    }
     const type = options.type || "agent";
     const baseProjectRoot = options.projectRoot || this.defaultProjectRoot;
     const hostAuthorizedBaseRoot =
@@ -479,6 +491,7 @@ export class WSSessionManager {
     const messages = [{ role: "system", content: systemPrompt }];
 
     // Persist to DB
+    let databaseCreated = false;
     if (this.db) {
       try {
         dbCreateSession(this.db, {
@@ -488,8 +501,17 @@ export class WSSessionManager {
           model: model || "",
           messages,
         });
-      } catch (_err) {
-        // Non-critical
+        databaseCreated = true;
+      } catch (error) {
+        if (requireDurable) {
+          this._discardSessionResources({
+            planManager,
+            worktree,
+            baseProjectRoot,
+          });
+          error.code ||= "CC_WS_SESSION_PERSISTENCE_FAILED";
+          throw error;
+        }
       }
     }
 
@@ -505,6 +527,7 @@ export class WSSessionManager {
       mcpClient: this.mcpClient,
       enabledToolNames,
       hostManagedToolPolicy: options.hostManagedToolPolicy || null,
+      sessionBudgetRoot: options.sessionBudgetRoot || null,
       externalToolDefinitions: externalTools.definitions,
       externalToolDescriptors: externalTools.descriptors,
       externalToolExecutors: externalTools.executors,
@@ -540,13 +563,33 @@ export class WSSessionManager {
       },
     );
 
-    if (this.db) {
+    if (this.db && databaseCreated) {
       try {
         dbUpdateSession(this.db, sessionId, {
           metadata: this._serializeSessionMetadata(session),
         });
-      } catch (_err) {
-        // Non-critical
+      } catch (error) {
+        if (requireDurable) {
+          const cleanupErrors = [];
+          try {
+            this._discardSessionResources(session);
+          } catch (cleanupError) {
+            cleanupErrors.push(cleanupError);
+          }
+          try {
+            dbDeleteSession(this.db, sessionId);
+          } catch (cleanupError) {
+            cleanupErrors.push(cleanupError);
+          }
+          error.code ||= "CC_WS_SESSION_PERSISTENCE_FAILED";
+          if (cleanupErrors.length > 0) {
+            throw new AggregateError(
+              [error, ...cleanupErrors],
+              "Durable WebSocket session creation and rollback both failed",
+            );
+          }
+          throw error;
+        }
       }
     }
 
@@ -556,6 +599,57 @@ export class WSSessionManager {
     return { sessionId };
   }
 
+  /**
+   * Remove a newly-created session that was never exposed to a client.
+   * Unlike closeSession(), this compensates the database insert instead of
+   * persisting a misleading closed session after bootstrap failed.
+   */
+  rollbackSessionCreation(sessionId) {
+    const session = this.sessions.get(sessionId);
+    const cleanupErrors = [];
+    if (session) {
+      try {
+        this._discardSessionResources(session);
+      } catch (error) {
+        cleanupErrors.push(error);
+      }
+      this.sessions.delete(sessionId);
+    }
+    if (this.db) {
+      try {
+        dbDeleteSession(this.db, sessionId);
+      } catch (error) {
+        cleanupErrors.push(error);
+      }
+    }
+    if (cleanupErrors.length > 0) {
+      throw new AggregateError(
+        cleanupErrors,
+        `WebSocket session creation rollback failed: ${sessionId}`,
+      );
+    }
+    return Boolean(session);
+  }
+
+  _discardSessionResources(session) {
+    if (!session) return;
+    if (
+      session._releaseHooksV2WorkspaceBindingOnClose === true &&
+      session.hooksV2WorkspaceBindingId
+    ) {
+      releaseRegisteredHostHooksV2Workspace(session.hooksV2WorkspaceBindingId);
+    }
+    if (typeof session._planPersistenceCleanup === "function") {
+      session._planPersistenceCleanup();
+    }
+    session.planManager?.removeAllListeners?.();
+    if (session.worktree?.path && session.baseProjectRoot) {
+      removeWorktree(session.baseProjectRoot, session.worktree.path, {
+        deleteBranch: true,
+      });
+    }
+  }
+
   /**
    * Resume an existing session from DB.
    *
@@ -645,6 +739,7 @@ export class WSSessionManager {
           metadata.enabledToolNames,
         ),
         hostManagedToolPolicy: metadata.hostManagedToolPolicy || null,
+        sessionBudgetRoot: metadata.sessionBudgetRoot || null,
         externalToolDefinitions: externalTools.definitions,
         externalToolDescriptors: externalTools.descriptors,
         externalToolExecutors: externalTools.executors,
@@ -1545,6 +1640,7 @@ export class WSSessionManager {
       baseProjectRoot: session.baseProjectRoot || session.projectRoot || null,
       baseUrl: session.baseUrl || null,
       hostManagedToolPolicy: session.hostManagedToolPolicy || null,
+      sessionBudgetRoot: session.sessionBudgetRoot || null,
       enabledToolNames: session.enabledToolNames || [],
       canonicalHostSystemPrefix: Array.isArray(
         session._canonicalHostSystemPrefix,