Something went wrong. Try again.
I'm going to build my own OpenClaw, with blackjack... and bun!
Something went wrong. Try again.
TypeScript
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506/** * agent-pool.ts – Manages per-chat pi-agent sessions and orchestrates agent runs. * * The AgentPool is the central coordinator between inbound messages and the * pi-coding-agent SDK. It: * - Maintains a map of chat JID → AgentSession (lazy-initialised). * - Provides runAgent() which prompts the agent, streams events, records * token usage, detects auto-compaction needs, and returns the result. * - Handles slash commands by delegating to agent-pool/slash-command.ts. * - Forwards agent-control commands (model switch, session management, etc.) * to the agent-control module. * - Manages session lifecycle: save/restore position (for scheduled tasks), * clear sessions, reload resources. * - Integrates the attachment registry for file-delivery tools. * * Consumers: * - runtime.ts / runtime/message-loop.ts creates the AgentPool at startup * and calls runAgent() for each inbound message. * - task-scheduler.ts calls runAgent() for scheduled task execution. * - channels/web.ts uses applyControlCommand() and agent status queries. * - agent-control handlers call methods on AgentPool for session/model ops. */
import { mkdirSync } from "fs";import { join } from "path";import { type AgentSession, type AgentSessionRuntime, AuthStorage, ModelRegistry, SettingsManager, getAgentDir,} from "@mariozechner/pi-coding-agent";
import { type AgentControlCommand, type AgentControlResult } from "./agent-control/index.js";import { SESSIONS_DIR, WORKSPACE_DIR } from "./core/config.js";import { createTrackedBashOperations } from "./tools/tracked-bash.js";import { type ActiveChatAgent } from "./agent-pool/branch-manager.js";import { type AgentOutput, type AgentPoolOptions, type RunAgentOptions, type SidePromptOptions, type SidePromptResult,} from "./agent-pool/contracts.js";import { runSidePrompt as runSidePromptInternal } from "./agent-pool/side-prompt-runner.js";import { runAgentPrompt } from "./agent-pool/run-agent-orchestrator.js";import { type AvailableModelsResult } from "./agent-pool/runtime-facade.js";import { createAgentPoolServices, type AgentPoolServices } from "./agent-pool/service-factory.js";import { type PoolEntry } from "./agent-pool/session-manager.js";import { type ChatBranchRecord, type SshConfig, type SshConfigApplyTiming, type SshConfigClearResult, type SshConfigSetResult, deleteSshConfig, getSshConfig, listRecentChatJids, upsertSshConfig,} from "./db.js";import { setSshToolHandlers } from "./extensions/ssh.js";import { setProxmoxToolHandlers } from "./extensions/proxmox.js";import { setPortainerToolHandlers } from "./extensions/portainer.js";import { clearStoredProxmoxConfig, getStoredProxmoxConfig, requestStoredProxmoxApi, runStoredProxmoxWorkflow, setStoredProxmoxConfig,} from "./proxmox/handlers.js";import { clearStoredPortainerConfig, getStoredPortainerConfig, requestStoredPortainerApi, runStoredPortainerWorkflow, setStoredPortainerConfig,} from "./portainer/handlers.js";import { applyLiveSshConfig, clearLiveSshConfig, hasLiveChatSshSession, resolveSshCoreConfigFromChatConfig } from "./extensions/ssh-core.js";import { createLogger } from "./utils/logger.js";
const log = createLogger("agent-pool");
export type { AgentOutput, AgentPoolOptions, RunAgentOptions, SidePromptOptions, SidePromptResult, TurnOutput,} from "./agent-pool/contracts.js";
/** How long (ms) an idle session stays cached before being disposed. */const DEFAULT_IDLE_TTL = 15 * 60 * 1000; // 15 minutesconst DEFAULT_CLEANUP_INTERVAL = 30 * 1000; // check every 30 seconds
function parsePositiveMs(value: string | undefined, fallback: number): number { const parsed = Number.parseInt(String(value || "").trim(), 10); return Number.isFinite(parsed) && parsed > 0 ? parsed : fallback;}
const IDLE_TTL = parsePositiveMs(process.env.PICLAW_SESSION_IDLE_TTL_MS, DEFAULT_IDLE_TTL);const CLEANUP_INTERVAL = parsePositiveMs(process.env.PICLAW_SESSION_CLEANUP_INTERVAL_MS, DEFAULT_CLEANUP_INTERVAL);const DEFAULT_PROVIDER_RATE_LIMIT_MAX_RETRIES = 5;const DEFAULT_PROVIDER_RATE_LIMIT_BASE_DELAY_MS = 5000;
/** * Manages a pool of persistent AgentSession instances keyed by chat JID. * * First invocation for a JID pays the warm-up cost (resource discovery, * model initialisation). Subsequent calls reuse the live session – no * process-spawn overhead, conversation context already loaded. */export class AgentPool { private pool = new Map<string, PoolEntry>(); private sidePool = new Map<string, PoolEntry>(); private activeForkBaseLeafByChat = new Map<string, string | null>(); private cleanupTimer: ReturnType<typeof setInterval> | null = null; private shuttingDown = false;
// Shared across all sessions (expensive to create, safe to reuse) private authStorage: AuthStorage; private modelRegistry: ModelRegistry; private settingsManager = SettingsManager.create(WORKSPACE_DIR, getAgentDir()); private logsDir = join(WORKSPACE_DIR, "logs"); private createSession?: AgentPoolOptions["createSession"]; private createSideSession?: AgentPoolOptions["createSideSession"]; private bashOperations = createTrackedBashOperations(); private attachments: AgentPoolServices["attachments"]; private sessionBinder: AgentPoolServices["sessionBinder"]; private toolFactory: AgentPoolServices["toolFactory"]; private turnCoordinator: AgentPoolServices["turnCoordinator"]; private sessionManager: AgentPoolServices["sessionManager"]; private branchManager: AgentPoolServices["branchManager"]; private runtimeFacade: AgentPoolServices["runtimeFacade"]; private sideStreamSimple?: NonNullable<AgentPoolOptions["sideStreamSimple"]>;
constructor(options: AgentPoolOptions = {}) { this.createSession = options.createSession; this.createSideSession = options.createSideSession; this.authStorage = AuthStorage.create(); this.modelRegistry = options.modelRegistry ?? ModelRegistry.create(this.authStorage); this.applyRateLimitRetryDefaults(); ({ attachments: this.attachments, sessionBinder: this.sessionBinder, toolFactory: this.toolFactory, turnCoordinator: this.turnCoordinator, sessionManager: this.sessionManager, runtimeFacade: this.runtimeFacade, branchManager: this.branchManager, } = createAgentPoolServices({ pool: this.pool, sidePool: this.sidePool, activeForkBaseLeafByChat: this.activeForkBaseLeafByChat, authStorage: this.authStorage, modelRegistry: this.modelRegistry, settingsManager: this.settingsManager, workspaceDir: WORKSPACE_DIR, bashOperations: this.bashOperations, createSession: this.createSession, createSideSession: this.createSideSession, onInfo: (message, details) => log.info(message, details), onWarn: (message, details) => log.warn(message, details), onError: (message, details) => log.error(message, details), })); this.sideStreamSimple = options.sideStreamSimple; setSshToolHandlers({ get: (chatJid) => this.getSshConfig(chatJid), set: (chatJid, config) => this.setSshConfig(chatJid, config), clear: (chatJid) => this.clearSshConfig(chatJid), }); setProxmoxToolHandlers({ get: (chatJid) => getStoredProxmoxConfig(chatJid), set: (chatJid, config) => setStoredProxmoxConfig(chatJid, config), clear: (chatJid) => clearStoredProxmoxConfig(chatJid), request: (chatJid, input) => requestStoredProxmoxApi(chatJid, input), workflow: (chatJid, input) => runStoredProxmoxWorkflow(chatJid, input), }); setPortainerToolHandlers({ get: (chatJid) => getStoredPortainerConfig(chatJid), set: (chatJid, config) => setStoredPortainerConfig(chatJid, config), clear: (chatJid) => clearStoredPortainerConfig(chatJid), request: (chatJid, input) => requestStoredPortainerApi(chatJid, input), workflow: (chatJid, input) => runStoredPortainerWorkflow(chatJid, input), }); mkdirSync(SESSIONS_DIR, { recursive: true }); mkdirSync(this.logsDir, { recursive: true }); this.cleanupTimer = setInterval(() => this.sessionManager.evictIdle(IDLE_TTL), CLEANUP_INTERVAL); }
private applyRateLimitRetryDefaults(): void { const settingsManager = this.settingsManager as SettingsManager & { getRetrySettings?: () => { enabled: boolean; maxRetries: number; baseDelayMs: number; maxDelayMs: number }; }; if (typeof settingsManager.getRetrySettings !== "function") return; const originalGetRetrySettings = settingsManager.getRetrySettings.bind(settingsManager); settingsManager.getRetrySettings = () => { const settings = originalGetRetrySettings(); return { ...settings, maxRetries: Math.max(settings.maxRetries ?? 0, DEFAULT_PROVIDER_RATE_LIMIT_MAX_RETRIES), baseDelayMs: Math.max(settings.baseDelayMs ?? 0, DEFAULT_PROVIDER_RATE_LIMIT_BASE_DELAY_MS), }; }; }
setSessionBinder(binder?: (runtime: AgentSessionRuntime, chatJid: string) => Promise<void> | void): void { this.sessionBinder.setBinder(binder); }
/** Run a prompt against the persistent session for `chatJid`. */ async runAgent(prompt: string, chatJid: string, options: RunAgentOptions = {}): Promise<AgentOutput> { return runAgentPrompt(prompt, chatJid, options, { getOrCreateRuntime: (nextChatJid) => this.getOrCreateRuntime(nextChatJid), turnCoordinator: this.turnCoordinator, clearAttachments: (nextChatJid) => this.attachments.clear(nextChatJid), takeAttachments: (nextChatJid) => this.attachments.take(nextChatJid), logsDir: this.logsDir, setActiveForkBaseLeaf: (nextChatJid, leafId) => this.activeForkBaseLeafByChat.set(nextChatJid, leafId), clearActiveForkBaseLeaf: (nextChatJid) => { this.activeForkBaseLeafByChat.delete(nextChatJid); }, onInfo: (message, details) => log.info(message, details), onWarn: (message, details) => log.warn(message, details), onError: (message, details) => log.error(message, details), }); }
async applyControlCommand(chatJid: string, command: AgentControlCommand): Promise<AgentControlResult> { return this.runtimeFacade.applyControlCommand(chatJid, command); }
async getCurrentModelLabel(chatJid: string): Promise<string | null> { return this.runtimeFacade.getCurrentModelLabel(chatJid); }
async runSidePrompt(chatJid: string, prompt: string, options: SidePromptOptions = {}): Promise<SidePromptResult> { return runSidePromptInternal(chatJid, prompt, options, { getOrCreate: (nextChatJid) => this.getOrCreate(nextChatJid), getOrCreateSideRuntime: (nextChatJid) => this.getOrCreateSideRuntime(nextChatJid), syncSideSessionFromMain: (mainSession, sideRuntime) => this.syncSideSessionFromMain(mainSession, sideRuntime), modelRegistry: this.modelRegistry, sideStreamSimple: this.sideStreamSimple, onWarn: (message, details) => log.warn(message, details), }); }
/** Return available model labels and current model for a chat session. */ async getAvailableModels(chatJid: string): Promise<AvailableModelsResult> { return this.runtimeFacade.getAvailableModels(chatJid); }
/** Return the current context token usage for a chat session, or null if unknown. */ async getContextUsageForChat(chatJid: string): Promise<{ tokens: number | null; contextWindow: number; percent: number | null; } | null> { return this.runtimeFacade.getContextUsageForChat(chatJid); }
scheduleRecentChatWarmup(options: { limit?: number; excludeChatJids?: string[] } = {}): string[] { if (this.shuttingDown) return []; const targetCount = Math.max(1, Math.min(8, Math.trunc(options.limit ?? 3) || 3)); const excluded = new Set( Array.isArray(options.excludeChatJids) ? options.excludeChatJids.map((jid) => String(jid || "").trim()).filter(Boolean) : [], ); const cooldownByChat = ((this as any).__piclawRecentWarmupCooldownByChat ||= new Map<string, number>()) as Map<string, number>; const now = Date.now(); for (const [chatJid, lastQueuedAt] of cooldownByChat) { if (now - lastQueuedAt >= 30_000) { cooldownByChat.delete(chatJid); } }
const scheduled: string[] = []; const seen = new Set<string>(); let fetchLimit = Math.min(100, Math.max(targetCount * 4, targetCount));
while (scheduled.length < targetCount) { const candidates = listRecentChatJids(fetchLimit, { excludeChatJids: [...excluded, ...seen], }); for (const chatJid of candidates) { if (seen.has(chatJid)) continue; seen.add(chatJid); if (excluded.has(chatJid)) continue; if (this.pool.has(chatJid)) continue; if (now - (cooldownByChat.get(chatJid) ?? 0) < 30_000) continue; scheduled.push(chatJid); if (scheduled.length >= targetCount) break; }
if (scheduled.length >= targetCount || fetchLimit >= 100 || candidates.length < fetchLimit) { break; }
const nextFetchLimit = Math.min(100, fetchLimit * 2); if (nextFetchLimit === fetchLimit) { break; } fetchLimit = nextFetchLimit; }
// Only record a cooldown / return chats that actually entered the prewarm // queue. prewarm() may reject a candidate (already queued, in flight, or // within its per-chat cooldown) and we must not consume a slot or suppress // backfill for those. const actuallyScheduled: string[] = []; for (const chatJid of scheduled) { if (this.sessionManager.prewarm(chatJid)) { cooldownByChat.set(chatJid, now); actuallyScheduled.push(chatJid); } }
return actuallyScheduled; }
scheduleChatWarmup(chatJid: string, options: { priority?: boolean } = {}): boolean { return this.sessionManager.prewarm(chatJid, options); }
getSessionTreeForChat(chatJid: string): { leafId: string | null; nodes: unknown[]; flat?: boolean; total?: number; capped?: boolean } | null { return this.runtimeFacade.getSessionTreeForChat(chatJid); }
/** * Save the current session tree position so it can be restored later. * Used by the scheduler to isolate task execution in a side branch. */ async saveSessionPosition(chatJid: string): Promise<string | null> { return this.runtimeFacade.saveSessionPosition(chatJid); }
/** * Restore the session tree to a previously saved position. * Navigates back to the saved leaf, leaving the task's output in a side branch. */ async restoreSessionPosition(chatJid: string, leafId: string | null): Promise<void> { return this.runtimeFacade.restoreSessionPosition(chatJid, leafId); }
async disposeChatSession(chatJid: string): Promise<void> { await this.sessionManager.recreate(chatJid); }
hasProviderModels(provider: string): boolean { return this.runtimeFacade.hasProviderModels(provider); }
registerModelProvider( providerName: string, config: Parameters<ModelRegistry["registerProvider"]>[1] ): void { this.runtimeFacade.registerModelProvider(providerName, config); }
resolveModelInput(input: string): { model?: string; error?: string } { return this.runtimeFacade.resolveModelInput(input); }
isStreaming(chatJid: string): boolean { return this.runtimeFacade.isStreaming(chatJid); }
isActive(chatJid: string): boolean { return this.runtimeFacade.isActive(chatJid); }
private ensureBranchRegistration(chatJid: string, session?: AgentSession | null): ChatBranchRecord { return this.branchManager.ensureBranchRegistration(chatJid, session); }
async renameChatBranch( chatJid: string, options: { agentName?: string | null } = {}, ): Promise<ChatBranchRecord> { return this.branchManager.renameChatBranch(chatJid, options); }
async pruneChatBranch(chatJid: string): Promise<ChatBranchRecord> { return this.branchManager.pruneChatBranch(chatJid); }
async restoreChatBranch( chatJid: string, options: { agentName?: string | null } = {}, ): Promise<ChatBranchRecord> { return this.branchManager.restoreChatBranch(chatJid, options); }
async createForkedChatBranch( sourceChatJid: string, options: { agentName?: string | null } = {}, ): Promise<ChatBranchRecord> { return this.branchManager.createForkedChatBranch(sourceChatJid, options); }
listActiveChats(): ActiveChatAgent[] { return this.branchManager.listActiveChats(); }
listKnownChats( rootChatJid?: string | null, options?: { includeArchived?: boolean } ): ActiveChatAgent[] { return this.branchManager.listKnownChats(rootChatJid, options); }
findActiveChatByAgentName(agentName: string): ActiveChatAgent | null { return this.branchManager.findActiveChatByAgentName(agentName); }
findChatByAgentName(agentName: string): { chat_jid: string; agent_name: string } | null { return this.branchManager.findChatByAgentName(agentName); }
getAgentHandleForChat(chatJid: string): string { return this.branchManager.getAgentHandleForChat(chatJid); }
async queueStreamingMessage( chatJid: string, text: string, behavior: "steer" | "followUp" ): Promise<{ queued: boolean; error?: string }> { return this.runtimeFacade.queueStreamingMessage(chatJid, text, behavior); }
/** Remove one queued follow-up message (first content match) from an active session queue. */ async removeQueuedFollowupMessage(chatJid: string, queuedContent?: string): Promise<boolean> { return this.runtimeFacade.removeQueuedFollowupMessage(chatJid, queuedContent); }
/** Execute a raw slash command in the AgentSession (extension commands). */ async applySlashCommand(chatJid: string, rawText: string): Promise<AgentControlResult> { return this.runtimeFacade.applySlashCommand(chatJid, rawText); }
getSshConfig(chatJid: string): SshConfig | null { return getSshConfig(chatJid); }
async setSshConfig( chatJid: string, config: Omit<SshConfig, "chat_jid" | "created_at" | "updated_at">, ): Promise<SshConfigSetResult> { const apply_timing: SshConfigApplyTiming = hasLiveChatSshSession(chatJid) ? "immediate" : "next_session"; if (apply_timing === "immediate") { await applyLiveSshConfig(chatJid, resolveSshCoreConfigFromChatConfig(config)); } const next = upsertSshConfig({ chat_jid: chatJid, ...config }); return { config: next, apply_timing }; }
async clearSshConfig(chatJid: string): Promise<SshConfigClearResult> { const apply_timing: SshConfigApplyTiming = hasLiveChatSshSession(chatJid) ? "immediate" : "next_session"; const deleted = deleteSshConfig(chatJid); if (apply_timing === "immediate") { await clearLiveSshConfig(chatJid); } return { deleted, apply_timing }; }
/** Gracefully shut down all sessions. */ async shutdown(): Promise<void> { this.shuttingDown = true; if (this.cleanupTimer) { clearInterval(this.cleanupTimer); this.cleanupTimer = null; } await this.sessionManager.shutdown(); }
/** Return an existing session for read-only introspection, or create one if needed. */ async getSessionForIntrospection(chatJid: string): Promise<AgentSession> { return this.getOrCreate(chatJid); }
// ── internal ────────────────────────────────────────────
private async getOrCreateRuntime(chatJid: string): Promise<AgentSessionRuntime> { return this.sessionManager.getOrCreate(chatJid); }
private async getOrCreate(chatJid: string): Promise<AgentSession> { return (await this.getOrCreateRuntime(chatJid)).session; }
private async getOrCreateSideRuntime(chatJid: string): Promise<AgentSessionRuntime> { return this.sessionManager.getOrCreateSide(chatJid); }
private async syncSideSessionFromMain(mainSession: AgentSession, sideRuntime: AgentSessionRuntime): Promise<void> { return this.sessionManager.syncSideSessionFromMain(mainSession, sideRuntime); }
private evictIdle(): void { this.sessionManager.evictIdle(IDLE_TTL); }}