Something went wrong. Try again.
I'm going to build my own OpenClaw, with blackjack... and bun!
Something went wrong. Try again.
20 kB · 538 lines
TypeScript
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539/** * agent-pool/session-manager.ts – Main/side session lifecycle management for AgentPool. * * Extracts chat-session creation, side-session sync, idle eviction, and shutdown * mechanics out of the top-level AgentPool coordinator while preserving the * existing map-based cache structure used by callers and tests. */
import type { AgentSession, AgentSessionRuntime, ExtensionFactory, ModelRegistry, SettingsManager, AuthStorage } from "@mariozechner/pi-coding-agent";
import { claimDeferredBranchSeed, finalizeClaimedDeferredBranchSeed, getDeferredBranchSeedFingerprint, hasDeferredBranchSeed, restoreClaimedDeferredBranchSeed, seedSessionManagerFromDeferredBranchSeed,} from "./branch-seeding.js";import { getChatBranchByChatJid } from "../db.js";import { createDefaultSession, createSessionInDir, ensureNamedSessionDir, ensureSessionDir } from "./session.js";import { forcePersistSessionFile, seedRotatedSession } from "../session-rotation.js";
/** Cached session entry stored for each chat JID. */export interface PoolEntry { runtime: AgentSessionRuntime; lastUsed: number;}
/** Dependencies required to manage AgentPool session lifecycles. */export interface AgentSessionManagerOptions { pool: Map<string, PoolEntry>; sidePool: Map<string, PoolEntry>; createSession?: (chatJid: string, sessionDir: string) => Promise<AgentSessionRuntime>; createSideSession?: (chatJid: string, sessionDir: string) => Promise<AgentSessionRuntime>; authStorage: AuthStorage; modelRegistry: ModelRegistry; settingsManager: SettingsManager; createDefaultTools: () => NonNullable<Parameters<typeof createDefaultSession>[1]["tools"]>; createCustomToolOverrides?: () => unknown[]; getSessionExtensionFactories?: (chatJid: string) => Promise<ExtensionFactory[]>; bindSession: (runtime: AgentSessionRuntime, chatJid: string) => Promise<void>; ensureBranchRegistration: (chatJid: string, session?: AgentSession | null) => void; onInfo?: (message: string, details: Record<string, unknown>) => void; onWarn?: (message: string, details: Record<string, unknown>) => void; onError?: (message: string, details: Record<string, unknown>) => void;}
/** * Handles lifecycle operations for main and auxiliary AgentSession instances. */export class AgentSessionManager { private readonly branchSeedRealizationInFlight = new Map<string, Promise<boolean>>(); private readonly createInFlight = new Map<string, Promise<AgentSessionRuntime>>(); private readonly invalidDeferredBranchSeedErrors = new Map<string, { error: Error; fingerprint: string | null }>(); private readonly runtimeDisposeInFlight = new WeakMap<AgentSessionRuntime, Promise<void>>(); private readonly prewarmInFlight = new Set<string>(); private readonly queuedPrewarms = new Set<string>(); private readonly prewarmQueue: string[] = []; private readonly prewarmCooldownByChat = new Map<string, number>(); private prewarmLoopActive = false;
constructor(private readonly options: AgentSessionManagerOptions) {}
private getBlockingInvalidSeedError(chatJid: string): Error | null { const knownInvalidSeedState = this.invalidDeferredBranchSeedErrors.get(chatJid); if (!knownInvalidSeedState) return null; if (knownInvalidSeedState.fingerprint === getDeferredBranchSeedFingerprint(chatJid)) { return knownInvalidSeedState.error; } this.invalidDeferredBranchSeedErrors.delete(chatJid); return null; }
private prunePrewarmCooldownEntries(now: number): void { for (const [chatJid, lastQueuedAt] of this.prewarmCooldownByChat) { if (now - lastQueuedAt >= 30_000) { this.prewarmCooldownByChat.delete(chatJid); } } }
async refreshRuntime(chatJid: string, runtime: AgentSessionRuntime): Promise<void> { const entry = this.options.pool.get(chatJid); if (entry) { entry.lastUsed = Date.now(); } await this.options.bindSession(runtime, chatJid); this.options.ensureBranchRegistration(chatJid, runtime.session); }
async getOrCreate(chatJid: string): Promise<AgentSessionRuntime> { const knownInvalidSeedError = this.getBlockingInvalidSeedError(chatJid); if (knownInvalidSeedError) { throw knownInvalidSeedError; }
const pendingCreate = this.createInFlight.get(chatJid); if (pendingCreate) { return await pendingCreate; }
const existing = this.options.pool.get(chatJid); if (existing) { existing.lastUsed = Date.now(); try { await this.realizeDeferredBranchSeed(chatJid, existing.runtime); return existing.runtime; } catch (error) { await this.disposeMainRuntimeAfterError(chatJid, existing.runtime, "get_or_create.realize_existing_after_error"); throw error; } }
const task = (async () => { this.options.onInfo?.("Creating new session", { operation: "get_or_create.create_main_session", chatJid, });
const chatSessionDir = ensureSessionDir(chatJid);
const extensionFactories = await this.options.getSessionExtensionFactories?.(chatJid) ?? []; const runtime = this.options.createSession ? await this.options.createSession(chatJid, chatSessionDir) : await createDefaultSession(chatJid, { authStorage: this.options.authStorage, modelRegistry: this.options.modelRegistry, settingsManager: this.options.settingsManager, tools: this.options.createDefaultTools(), customTools: this.options.createCustomToolOverrides?.() ?? [], extensionFactories, });
this.options.pool.set(chatJid, { runtime, lastUsed: Date.now() }); try { const realizedDeferredSeed = await this.realizeDeferredBranchSeed(chatJid, runtime); if (!realizedDeferredSeed) { await this.applyDefaultModel(runtime.session); } await this.refreshRuntime(chatJid, runtime); this.options.onInfo?.("Session ready", { operation: this.options.createSession ? "get_or_create.create_main_session" : "get_or_create.create_default_session", chatJid, poolSize: this.options.pool.size, }); return runtime; } catch (err) { await this.disposeMainRuntimeAfterError(chatJid, runtime, "get_or_create.dispose_after_error"); throw err; } })().finally(() => { this.createInFlight.delete(chatJid); });
this.createInFlight.set(chatJid, task); return await task; }
async getOrCreateSide(chatJid: string): Promise<AgentSessionRuntime> { const existing = this.options.sidePool.get(chatJid); if (existing) { existing.lastUsed = Date.now(); return existing.runtime; }
this.options.onInfo?.("Creating new side session", { operation: "get_or_create_side.create_session", chatJid, }); const sideSessionDir = ensureNamedSessionDir(chatJid, "btw-side");
const extensionFactories = await this.options.getSessionExtensionFactories?.(chatJid) ?? []; const runtime = this.options.createSideSession ? await this.options.createSideSession(chatJid, sideSessionDir) : await createSessionInDir(sideSessionDir, { authStorage: this.options.authStorage, modelRegistry: this.options.modelRegistry, settingsManager: this.options.settingsManager, tools: this.options.createDefaultTools(), customTools: this.options.createCustomToolOverrides?.() ?? [], extensionFactories, });
this.options.sidePool.set(chatJid, { runtime, lastUsed: Date.now() }); return runtime; }
async syncSideSessionFromMain(mainSession: AgentSession, sideRuntime: AgentSessionRuntime): Promise<void> { try { const mainContext = mainSession.sessionManager.buildSessionContext(); await sideRuntime.newSession({ setup: async (sessionManager) => { seedRotatedSession(sessionManager, mainContext, { sessionName: "BTW", model: mainContext.model, }); }, }); } catch (err) { this.options.onWarn?.("Failed to reseed side session from main context", { operation: "sync_side_session_from_main.reseed", err, }); }
const sideSession = sideRuntime.session; const mainModel = mainSession.model; const sideModel = sideSession.model; if (mainModel && (!sideModel || sideModel.provider !== mainModel.provider || sideModel.id !== mainModel.id)) { try { await sideSession.setModel(mainModel); } catch (err) { this.options.onWarn?.("Failed to sync side-session model", { operation: "sync_side_session_from_main.model", model: `${mainModel.provider}/${mainModel.id}`, err, }); } }
try { sideSession.setThinkingLevel(mainSession.thinkingLevel); } catch (err) { this.options.onWarn?.("Failed to sync side-session thinking level", { operation: "sync_side_session_from_main.thinking_level", err, }); }
try { sideSession.setActiveToolsByName(mainSession.getActiveToolNames()); } catch (err) { this.options.onWarn?.("Failed to sync side-session tool selection", { operation: "sync_side_session_from_main.tools", err, }); } }
async recreate(chatJid: string): Promise<void> { await this.disposeEntry(this.options.pool, chatJid, "recreate.dispose_main_session"); await this.disposeEntry(this.options.sidePool, chatJid, "recreate.dispose_side_session", true); }
/** * Remove any pending (not-yet-realized) prewarm for a chat. Used by * branch prune so a queued warmup cannot materialize a blank runtime * (or realize the deferred seed) for an archived chat. */ cancelPrewarm(chatJid: string): void { const normalized = String(chatJid || "").trim(); if (!normalized) return; this.queuedPrewarms.delete(normalized); const idx = this.prewarmQueue.indexOf(normalized); if (idx >= 0) this.prewarmQueue.splice(idx, 1); }
prewarm(chatJid: string, options: { priority?: boolean } = {}): boolean { const normalizedChatJid = String(chatJid || "").trim(); if (!normalizedChatJid) return false; if (this.getBlockingInvalidSeedError(normalizedChatJid)) return false; if (this.options.pool.has(normalizedChatJid) && !hasDeferredBranchSeed(normalizedChatJid)) return false; if (this.prewarmInFlight.has(normalizedChatJid) || this.queuedPrewarms.has(normalizedChatJid)) return false; const now = Date.now(); this.prunePrewarmCooldownEntries(now); const lastQueuedAt = this.prewarmCooldownByChat.get(normalizedChatJid) ?? 0; if (!options.priority && now - lastQueuedAt < 30_000) return false;
this.queuedPrewarms.add(normalizedChatJid); this.prewarmCooldownByChat.set(normalizedChatJid, now); if (options.priority) { this.prewarmQueue.unshift(normalizedChatJid); } else { this.prewarmQueue.push(normalizedChatJid); } this.drainPrewarmQueue(); return true; }
async shutdown(): Promise<void> { for (const jid of [...this.options.pool.keys()]) { await this.disposeEntry(this.options.pool, jid, "shutdown.dispose_main_session"); } for (const jid of [...this.options.sidePool.keys()]) { await this.disposeEntry(this.options.sidePool, jid, "shutdown.dispose_side_session", true); } this.options.pool.clear(); this.options.sidePool.clear(); }
private async disposeEntry( map: Map<string, PoolEntry>, chatJid: string, operation: string, side = false, ): Promise<void> { const entry = map.get(chatJid); if (!entry) return; try { await entry.runtime.dispose(); this.options.onInfo?.(side ? "Disposed side session" : "Disposed session", { operation, chatJid, }); } catch (err) { this.options.onError?.(side ? "Failed to dispose side session" : "Failed to dispose session", { operation, chatJid, err, }); } finally { map.delete(chatJid); } }
evictIdle(idleTtlMs: number): void { const now = Date.now(); for (const [jid, entry] of this.options.pool) { const session = entry.runtime.session; if (session.isStreaming || session.isBashRunning || session.isCompacting) { entry.lastUsed = now; continue; } if (now - entry.lastUsed > idleTtlMs) { this.options.onInfo?.("Evicting idle session", { operation: "evict_idle.main_session", chatJid: jid, }); Promise.resolve(entry.runtime.dispose()).catch((err) => { this.options.onWarn?.("Failed to dispose evicted session", { operation: "evict_idle.main_session", chatJid: jid, err, }); }); this.options.pool.delete(jid); } } for (const [jid, entry] of this.options.sidePool) { const session = entry.runtime.session; if (session.isStreaming || session.isBashRunning || session.isCompacting) { entry.lastUsed = now; continue; } if (now - entry.lastUsed > idleTtlMs) { this.options.onInfo?.("Evicting idle side session", { operation: "evict_idle.side_session", chatJid: jid, }); Promise.resolve(entry.runtime.dispose()).catch((err) => { this.options.onWarn?.("Failed to dispose evicted side session", { operation: "evict_idle.side_session", chatJid: jid, err, }); }); this.options.sidePool.delete(jid); } } }
private async applyDefaultModel(session: AgentSession): Promise<void> { const provider = this.options.settingsManager.getDefaultProvider(); const modelId = this.options.settingsManager.getDefaultModel(); if (!provider || !modelId) return;
await this.applyResolvedModel(session, { provider, modelId }, "apply_default_model"); }
private async realizeDeferredBranchSeed(chatJid: string, runtime: AgentSessionRuntime): Promise<boolean> { const pending = this.branchSeedRealizationInFlight.get(chatJid); if (pending) return await pending;
const task = (async () => { let seed; try { seed = claimDeferredBranchSeed(chatJid); } catch (error) { if (error instanceof Error && error.message.startsWith("Invalid deferred branch seed for ")) { this.invalidDeferredBranchSeedErrors.set(chatJid, { error, fingerprint: getDeferredBranchSeedFingerprint(chatJid), }); } throw error; } if (!seed) return false;
try { const result = await runtime.newSession({ ...(seed.parentSession ? { parentSession: seed.parentSession } : {}), setup: async (sessionManager) => { seedSessionManagerFromDeferredBranchSeed(sessionManager, seed); }, }); if (result.cancelled) { throw new Error("Deferred branch seed was cancelled."); }
const session = runtime.session; await this.applyResolvedModel(session, seed.model, "realize_deferred_branch_seed"); try { if (seed.thinkingLevel) { session.setThinkingLevel(seed.thinkingLevel); } } catch (err) { this.options.onWarn?.("Failed to restore deferred branch thinking level", { operation: "realize_deferred_branch_seed.thinking_level", chatJid, err, }); } try { // Prefer the current branch record's agent_name over the seed's // captured sessionName: if the user renamed the branch between fork // and realization, the DB has the new name and applying the stale // seed value would overwrite the rename on first warmup. const liveAgentName = getChatBranchByChatJid(chatJid)?.agent_name?.trim() || ""; const sessionNameToApply = liveAgentName || seed.sessionName?.trim() || ""; if (sessionNameToApply) { session.setSessionName(sessionNameToApply); } } catch (err) { this.options.onWarn?.("Failed to restore deferred branch session name", { operation: "realize_deferred_branch_seed.session_name", chatJid, err, }); }
forcePersistSessionFile(session); finalizeClaimedDeferredBranchSeed(chatJid); this.invalidDeferredBranchSeedErrors.delete(chatJid); return true; } catch (error) { restoreClaimedDeferredBranchSeed(chatJid); throw error; } })().finally(() => { this.branchSeedRealizationInFlight.delete(chatJid); });
this.branchSeedRealizationInFlight.set(chatJid, task); return await task; }
private async applyResolvedModel( session: AgentSession, model: { provider: string; modelId: string } | null, operation: string, ): Promise<void> { if (!model) return;
const current = session.model; if (current && current.provider === model.provider && current.id === model.modelId) return;
const sessionRegistry = (session as AgentSession & { modelRegistry?: ModelRegistry }).modelRegistry ?? this.options.modelRegistry; const resolved = sessionRegistry.find(model.provider, model.modelId); if (!resolved) return;
const setModel = (session as { setModel?: (model: typeof resolved) => Promise<void> }).setModel; if (typeof setModel !== "function") return;
try { await setModel.call(session, resolved); } catch (err) { this.options.onWarn?.("Failed to restore model", { operation, model: `${model.provider}/${model.modelId}`, err, }); } }
private drainPrewarmQueue(): void { if (this.prewarmLoopActive) return; this.prewarmLoopActive = true;
void (async () => { try { while (this.prewarmQueue.length > 0) { const chatJid = this.prewarmQueue.shift(); if (!chatJid) continue; this.queuedPrewarms.delete(chatJid); if (this.prewarmInFlight.has(chatJid)) continue; if (this.options.pool.has(chatJid) && !hasDeferredBranchSeed(chatJid)) continue;
this.prewarmInFlight.add(chatJid); try { await this.getOrCreate(chatJid); this.options.onInfo?.("Prewarmed chat session", { operation: "prewarm_session", chatJid, }); } catch (err) { this.options.onWarn?.("Failed to prewarm chat session", { operation: "prewarm_session", chatJid, err, }); } finally { this.prewarmInFlight.delete(chatJid); } } } finally { this.prewarmLoopActive = false; if (this.prewarmQueue.length > 0) { this.drainPrewarmQueue(); } } })(); }
private async disposeMainRuntimeAfterError(chatJid: string, runtime: AgentSessionRuntime, operation: string): Promise<void> { if (this.options.pool.get(chatJid)?.runtime === runtime) { this.options.pool.delete(chatJid); } const pendingDispose = this.runtimeDisposeInFlight.get(runtime); if (pendingDispose) { await pendingDispose; return; }
const task = (async () => { try { await runtime.dispose(); } catch (disposeErr) { this.options.onWarn?.("Failed to dispose session after initialization error", { operation, chatJid, err: disposeErr, }); } })(); this.runtimeDisposeInFlight.set(runtime, task); await task; }}