@ponythewhite/base-context 1.0.0 → 1.0.1

package/dist/modes/daemon/daemon-mode.js

@@ -46,6 +46,7 @@ import { filterClientEnv, withClientEnv } from "./daemon-client-env.js";
 import { deserializeDaemonError, serializeDaemonError } from "./daemon-errors.js";
 import { bindActiveSessionState } from "./daemon-extension-binding.js";
 import { CANONICAL_SESSION_OWNERSHIP_COMPATIBILITY, createDaemonEventMeta, createDaemonReplayInfo, DAEMON_DEFAULT_CLIENT_CAPABILITIES, DAEMON_DEFAULT_SERVER_CAPABILITIES, DAEMON_PROTOCOL_INFO, DAEMON_SCHEMA_ID, DAEMON_SCHEMA_REVISION, DAEMON_SUPPORTED_CLIENT_CAPABILITIES, DAEMON_UPDATE_RESTART_FORMAT_VERSION, failure, isDaemonCommandEnvelope, isDaemonDialogExtensionUiRequest, isDaemonMutatingCommand, isLegacyDaemonInspection, isSessionPlaneDaemonCommand, salvageDaemonCommandId, success, UPDATE_RESTART_DRAIN_COMMANDS, } from "./daemon-protocol.js";
+import { DaemonRlmCapacityClient } from "./daemon-rlm-capacity.js";
 import { getDaemonRuntimeIdentity } from "./daemon-runtime-identity.js";
 import { buildRlmChildSnapshots, buildSessionList, classifySessionRosterStatus, hasLiveSessionWork, inactiveLifecycleForSession, scheduledJobRegistrations, snapshotActiveSessionSummary, summaryForActiveSession, } from "./daemon-session-list.js";
 import { DaemonSessionSummarizer } from "./daemon-session-summarizer.js";
@@ -150,7 +151,9 @@ const DAEMON_COMMAND_TYPES = new Set([
     "export_jsonl",
     "set_session_name",
     "get_rlm_max_depth_status",
+    "get_rlm_max_subagents_status",
     "set_rlm_max_depth",
+    "set_rlm_max_subagents",
     "rename_saved_session",
     "delete_saved_session",
     "get_session_context",
@@ -291,6 +294,7 @@ export class AgentDaemon {
     rosterFlushPending = false;
     rosterHeartbeatTimer;
     rlmSpawnLedgerInstance;
+    rlmCapacityClient;
     rlmJournalOwner;
     rlmJournalAdmissionClosed = false;
     rlmJournalClosePromise;
@@ -711,6 +715,21 @@ export class AgentDaemon {
             client.close();
         }
     }
+    workerRlmCapacityClient() {
+        const worker = this.options.worker;
+        if (!worker)
+            return undefined;
+        const socketPath = this.supervisorSocketPathFromEnv();
+        if (!socketPath || !worker.workerInstanceId) {
+            throw new Error("RLM capacity requires the current worker incarnation and its supervisor");
+        }
+        this.rlmCapacityClient ??= new DaemonRlmCapacityClient({
+            socketPath,
+            workerToken: worker.authenticationToken,
+            workerInstanceId: worker.workerInstanceId,
+        });
+        return this.rlmCapacityClient;
+    }
     rlmSpawnLedger() {
         if (!this.rlmSpawnLedgerInstance) {
             this.rlmSpawnLedgerInstance = new RlmSpawnLedger(this.agentDir, this.rlmLedgerSessionsDir(), createRlmLedgerRegistrySeedSource(), (message) => this.log(message), { mode: "remote", mutate: (mutation) => this.sendRlmLedgerMutation(mutation) });
@@ -1199,6 +1218,63 @@ export class AgentDaemon {
             this.cronScheduler.wake();
         }
     }
+    createRuntimeWithRlmCapacity = async (options) => {
+        const authority = this.workerRlmCapacityClient();
+        if (!authority || options.sessionOptions?.rlmChildAdmission)
+            return this.options.createRuntime(options);
+        const origin = {
+            sessionFile: options.sessionManager.getSessionFile(),
+            sessionId: options.sessionManager.getSessionId(),
+        };
+        const reservation = await authority.reserveRoot(origin);
+        let admitted;
+        let created;
+        const admission = {
+            get session() {
+                return admitted;
+            },
+            bind: (session) => {
+                if (admitted && admitted !== session)
+                    throw new Error("RLM root admission already owns another session");
+                admitted = session;
+            },
+            release: () => reservation.release(),
+        };
+        try {
+            const result = await this.options.createRuntime({
+                ...options,
+                sessionOptions: {
+                    ...options.sessionOptions,
+                    rlmSubagentCapacity: authority.forSession(origin),
+                    rlmRootAdmission: admission,
+                },
+            });
+            created = result.session;
+            if (admission.session !== result.session)
+                throw new Error("Runtime factory did not bind its RLM root admission");
+            return result;
+        }
+        catch (error) {
+            try {
+                const failed = admission.session ?? created;
+                if (failed) {
+                    await failed.disposeAsync();
+                    if (!admission.session)
+                        await reservation.release();
+                }
+                else {
+                    await options.sessionManager.close();
+                    await reservation.release();
+                }
+            }
+            catch (cleanupError) {
+                if (cleanupError === error || (error instanceof AggregateError && error.errors.includes(cleanupError)))
+                    throw error;
+                throw new AggregateError([error, cleanupError], "Root runtime creation and capacity cleanup failed");
+            }
+            throw error;
+        }
+    };
     async createRuntime(command, runtimeOpenGuard) {
         const config = mergeAgentSessionRuntimeConfig(this.options.defaultSessionConfig, command.config);
         if (!config.cwd) {
@@ -1376,7 +1452,7 @@ export class AgentDaemon {
                 // Extensions capture client env (e.g. herdr pane identity) synchronously
                 // while the runtime loads them, so it must be in process.env for the
                 // duration; withClientEnv restores it after.
-                const runtime = await withClientEnv(clientEnv, () => createAgentSessionRuntime(this.options.createRuntime, {
+                const runtime = await withClientEnv(clientEnv, () => createAgentSessionRuntime(this.createRuntimeWithRlmCapacity, {
                     cwd: sessionManager.getCwd(),
                     agentDir,
                     sessionManager,
@@ -2052,7 +2128,7 @@ export class AgentDaemon {
             : structuredClone(options.requestTokenBudget);
         if (requestTokenBudget !== undefined)
             options = { ...options, requestTokenBudget };
-        const admission = options.admission ?? options.parentSession.reserveRlmChildAdmission();
+        const admission = options.admission ?? (await options.parentSession.reserveRlmChildAdmission());
         try {
             if (admission.parent !== options.parentSession ||
                 parentState.runtime.session !== admission.parent ||
@@ -2065,7 +2141,7 @@ export class AgentDaemon {
         }
         finally {
             if (!options.admission)
-                admission.settle();
+                await admission.settle();
         }
     }
     async createAdmittedRlmSubagentRuntime(parentState, options) {
@@ -2088,7 +2164,7 @@ export class AgentDaemon {
             options.admission.assertCurrent();
             factoryStarted = true;
             // Subagents inherit the parent's client env (e.g. herdr pane identity).
-            runtime = await withClientEnv(parentState.clientEnv, () => createAgentSessionRuntime(this.options.createRuntime, {
+            runtime = await withClientEnv(parentState.clientEnv, () => createAgentSessionRuntime(this.createRuntimeWithRlmCapacity, {
                 cwd: sessionManager.getCwd(),
                 agentDir: parentState.runtime.services.agentDir,
                 sessionManager,
@@ -2163,7 +2239,7 @@ export class AgentDaemon {
             try {
                 await sessionManager.close();
                 if (!factoryStarted)
-                    options.admission.confirmUnboundCleanup();
+                    await options.admission.confirmUnboundCleanup();
             }
             catch (cleanupError) {
                 if (cleanupError === error || (error instanceof AggregateError && error.errors.includes(cleanupError)))
@@ -2443,15 +2519,15 @@ export class AgentDaemon {
             this.closingSessions.has(parentState.activeSessionId)) {
             throw new RuntimeOpenCancelledError();
         }
-        const admission = parentState.runtime.session.reserveRlmChildAdmission();
         const hydration = (async () => {
+            const admission = await parentState.runtime.session.reserveRlmChildAdmission();
             try {
                 if (existing)
                     await this.closeSession(existing, "replaced");
                 return await this.rehydrateCompletedRlmSubagentOnce(parentState, entry, admission, restoreActiveSessionId, clientEnv);
             }
             finally {
-                admission.settle();
+                await admission.settle();
             }
         })();
         // Explicit opens and all lazy triggers share this path-keyed publication,
@@ -2514,7 +2590,7 @@ export class AgentDaemon {
             assertAdmission();
             admission.claimFactory();
             factoryStarted = true;
-            runtime = await withClientEnv(hydrationEnv, () => createAgentSessionRuntime(this.options.createRuntime, {
+            runtime = await withClientEnv(hydrationEnv, () => createAgentSessionRuntime(this.createRuntimeWithRlmCapacity, {
                 cwd: openedManager.getCwd(),
                 agentDir: parentState.runtime.services.agentDir,
                 sessionManager: openedManager,
@@ -2634,7 +2710,7 @@ export class AgentDaemon {
             await drain(() => sessionManager?.close());
             await drain(() => sessionLease?.release());
             if (!factoryStarted && errors.length === initialErrorCount)
-                admission.confirmUnboundCleanup();
+                await admission.confirmUnboundCleanup();
             if (errors.length === initialErrorCount)
                 throw error;
             throw new AggregateError(errors, "RLM hydration and cleanup failed");
@@ -3240,6 +3316,10 @@ export class AgentDaemon {
                     setImmediate(() => void this.shutdown(0));
                     return;
                 }
+                case "worker_get_rlm_capacity": {
+                    this.writeWorkerSuccess(client, command, this.rlmCapacityClient?.snapshot() ?? { reservations: [] });
+                    return;
+                }
                 case "worker_passivate_idle_children": {
                     const count = await this.passivateIdleChildren(command.idleEvictionMinutes, command.now, command.limit);
                     this.writeWorkerSuccess(client, command, { count });
@@ -4375,6 +4455,14 @@ export class AgentDaemon {
                 await this.setStateSessionNameForCommand(state, name);
                 return success(command.id, "set_session_name");
             }
+            case "get_rlm_max_subagents_status": {
+                const state = this.getSessionState(command.activeSessionId);
+                return success(command.id, command.type, await state.runtime.session.getRlmMaxSubagentsStatus());
+            }
+            case "set_rlm_max_subagents": {
+                const state = this.getSessionState(command.activeSessionId);
+                return success(command.id, command.type, await state.runtime.session.setRlmMaxSubagents(command.maxSubagents));
+            }
             case "get_rlm_max_depth_status": {
                 const state = this.getSessionState(command.activeSessionId);
                 return success(command.id, "get_rlm_max_depth_status", state.runtime.session.getRlmMaxDepthStatus());
@@ -5310,16 +5398,23 @@ Use prime_context read/search with selected lines or a query.`;
             // runs, in which case the draft is no longer abandoned and must be kept.
             queueMicrotask(() => {
                 if (this.sessions.has(state.activeSessionId) && this.isDiscardableDraft(state)) {
-                    void this.closeSession(state, "killed");
+                    this.discardAbandonedDraft(state);
                 }
             });
         }
     }
+    /** Best-effort discard: a teardown rejection must not stop other hosted sessions. */
+    discardAbandonedDraft(state) {
+        void this.closeSession(state, "killed").catch((error) => {
+            this.log(`failed to discard abandoned empty draft ${state.activeSessionId}: ${error instanceof Error ? (error.stack ?? error.message) : String(error)}`);
+        });
+    }
     isDiscardableDraft(state) {
         if (this.options.worker) {
             return false;
         }
-        if (state.clients.size > 0) {
+        // An attaching client is not in state.clients yet, but still owns the draft.
+        if (state.clients.size > 0 || state.pendingAttaches > 0) {
             return false;
         }
         if (state.runtime.metadata.kind === "subagent") {
@@ -5879,7 +5974,7 @@ Use prime_context read/search with selected lines or a query.`;
             // settles so it doesn't linger in the daemon.
             if ((eventType === "turn_end" || eventType === "compaction_end" || eventType === "bash_end") &&
                 this.isDiscardableDraft(state)) {
-                void this.closeSession(state, "killed");
+                this.discardAbandonedDraft(state);
             }
             if (RECOVERY_CHECKPOINT_EVENTS.has(eventType)) {
                 this.recordWorkerRecoveryState(state, eventType);