Something went wrong. Try again.
fork of prime agent with changes i want but are not upstreamable probqbly (and some fixes)
Something went wrong. Try again.
40 kB · 1075 lines
TypeScript
at main
12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091929394959697989910010110210310410510610710810911011111211311411511611711811912012112212312412512612712812913013113213313413513613713813914014114214314414514614714814915015115215315415515615715815916016116216316416516616716816917017117217317417517617717817918018118218318418518618718818919019119219319419519619719819920020120220320420520620720820921021121221321421521621721821922022122222322422522622722822923023123223323423523623723823924024124224324424524624724824925025125225325425525625725825926026126226326426526626726826927027127227327427527627727827928028128228328428528628728828929029129229329429529629729829930030130230330430530630730830931031131231331431531631731831932032132232332432532632732832933033133233333433533633733833934034134234334434534634734834935035135235335435535635735835936036136236336436536636736836937037137237337437537637737837938038138238338438538638738838939039139239339439539639739839940040140240340440540640740840941041141241341441541641741841942042142242342442542642742842943043143243343443543643743843944044144244344444544644744844945045145245345445545645745845946046146246346446546646746846947047147247347447547647747847948048148248348448548648748848949049149249349449549649749849950050150250350450550650750850951051151251351451551651751851952052152252352452552652752852953053153253353453553653753853954054154254354454554654754854955055155255355455555655755855956056156256356456556656756856957057157257357457557657757857958058158258358458558658758858959059159259359459559659759859960060160260360460560660760860961061161261361461561661761861962062162262362462562662762862963063163263363463563663763863964064164264364464564664764864965065165265365465565665765865966066166266366466566666766866967067167267367467567667767867968068168268368468568668768868969069169269369469569669769869970070170270370470570670770870971071171271371471571671771871972072172272372472572672772872973073173273373473573673773873974074174274374474574674774874975075175275375475575675775875976076176276376476576676776876977077177277377477577677777877978078178278378478578678778878979079179279379479579679779879980080180280380480580680780880981081181281381481581681781881982082182282382482582682782882983083183283383483583683783883984084184284384484584684784884985085185285385485585685785885986086186286386486586686786886987087187287387487587687787887988088188288388488588688788888989089189289389489589689789889990090190290390490590690790890991091191291391491591691791891992092192292392492592692792892993093193293393493593693793893994094194294394494594694794894995095195295395495595695795895996096196296396496596696796896997097197297397497597697797897998098198298398498598698798898999099199299399499599699799899910001001100210031004100510061007100810091010101110121013101410151016101710181019102010211022102310241025102610271028102910301031103210331034103510361037103810391040104110421043104410451046104710481049105010511052105310541055105610571058105910601061106210631064106510661067106810691070107110721073107410751076import { existsSync } from "node:fs";import { join } from "node:path";import type { AgentTool } from "@earendil-works/pi-agent-core";import type { ImageContent, TextContent } from "@earendil-works/pi-ai";import { type Static, Type } from "typebox";import { IMAGE_MIME_TYPES } from "../../utils/mime.js";import { resolveKernelBashShell } from "../../utils/shell.js";import type { ExtensionContext, ToolDefinition } from "../extensions/types.js";import { withKernelBootPermit } from "../kernel/boot-gate.js";import { KERNEL_VENV_LOCK_STALE_MS, KERNEL_VENV_LOCK_SUFFIX, type KernelBootstrapProgressHandler,} from "../kernel/bootstrap.js";import { CELL_TIMEOUT_MAX_SECONDS, DEFAULT_CELL_DEADLINE_MS, type ExecuteResult, type HostRequestHandlers, type KernelAttachment, KernelBusyAfterInterruptError, type KernelClient, type KernelContentDisplay, type KernelDiffDisplay, type KernelSentAgentMessage, type RemoteHostConfig, ReplKernelManager,} from "../kernel/index.js";import { manifestPathIn, type RestoreResult, snapshotPathIn } from "../kernel/state-snapshot.js";import type { PythonSkillRuntimeInfo } from "../skills.js";import { wrapToolDefinition } from "./tool-definition-wrapper.js";
const RLM_BOOTSTRAP_HEADER_CODE = `import asyncioimport importlib as _prime_agent_importlibimport inspect as _prime_agent_inspectimport os as _prime_agent_osimport sys as _prime_agent_sysimport time as _prime_agent_timeimport types as _prime_agent_types
_prime_agent_os.environ["NO_COLOR"] = "1"`.trim();
const RLM_BOOTSTRAP_RUNTIME_CODE = `_PRIME_AGENT_BOOTSTRAP_LOCK_SUFFIX = ${JSON.stringify(KERNEL_VENV_LOCK_SUFFIX)}_PRIME_AGENT_BOOTSTRAP_STALE_LOCK_SECONDS = ${KERNEL_VENV_LOCK_STALE_MS / 1000}_PRIME_AGENT_BOOTSTRAP_WAIT_SECONDS = float( _prime_agent_os.environ.get("PRIME_AGENT_KERNEL_BOOTSTRAP_WAIT_SECONDS", "120"))_PRIME_AGENT_BOOTSTRAP_POLL_SECONDS = 0.2_PRIME_AGENT_RLM_IMPORT_ERROR = None_PRIME_AGENT_SKILL_IMPORT_ERRORS = {}
class _PrimeAgentCallableSkillModule(_prime_agent_types.ModuleType): async def __call__(self, *args, **kwargs): result = self.run(*args, **kwargs) if _prime_agent_inspect.isawaitable(result): return await result return result
def _prime_agent_wrap_skill_module(module): run = getattr(module, "run", None) if not callable(run): return module if isinstance(module, _PrimeAgentCallableSkillModule): return module wrapped = _PrimeAgentCallableSkillModule(module.__name__) wrapped.__dict__.update(module.__dict__) try: wrapped.__signature__ = _prime_agent_inspect.signature(run) except Exception: pass doc = getattr(run, "__doc__", None) if doc: wrapped.__doc__ = doc _prime_agent_sys.modules[module.__name__] = wrapped return wrapped
class _PrimeAgentUnavailableSkill: def __init__(self, name, error): self.__name__ = name self._prime_agent_import_error = error self.__doc__ = f"Python skill {name} is unavailable: {error}"
def _rebind(self): """Retry the import: an editable install that was mid-flight may be whole now.""" try: module = _prime_agent_wrap_skill_module(_prime_agent_importlib.import_module(self.__name__)) except Exception as _prime_agent_skill_error: self._prime_agent_import_error = str(_prime_agent_skill_error) venv = _prime_agent_os.environ.get("PRIME_AGENT_KERNEL_VENV") or _prime_agent_sys.prefix raise RuntimeError( f"Python skill {self.__name__} is unavailable in this kernel. " f"Remove {venv} so prime-agent can rebuild it, or set " "PRIME_AGENT_KERNEL_PYTHON to a kernel environment with prime-agent-runtime installed. " f"Import error: {self._prime_agent_import_error}" ) from _prime_agent_skill_error globals()[self.__name__] = module _PRIME_AGENT_SKILL_IMPORT_ERRORS.pop(self.__name__, None) return module
def __getattr__(self, name): module = self._rebind() return getattr(module, name)
async def run(self, *args, **kwargs): run = getattr(self._rebind(), "run", None) if not callable(run): raise RuntimeError(f"Python skill {self.__name__} has no run() to call") result = run(*args, **kwargs) if _prime_agent_inspect.isawaitable(result): return await result return result
async def __call__(self, *args, **kwargs): return await self.run(*args, **kwargs)
def __repr__(self): return f"<unavailable Python skill {self.__name__!r}: {self._prime_agent_import_error}>"
def _prime_agent_bootstrap_lock_path(): """The host's venv build lock for the venv this kernel runs from, if it can tell.""" venv = _prime_agent_os.environ.get("PRIME_AGENT_KERNEL_VENV") or _prime_agent_sys.prefix if not venv: return None venv = _prime_agent_os.path.abspath(_prime_agent_os.path.expanduser(venv)) return _prime_agent_os.path.join( _prime_agent_os.path.dirname(venv), _prime_agent_os.path.basename(venv) + _PRIME_AGENT_BOOTSTRAP_LOCK_SUFFIX, )
def _prime_agent_pid_is_alive(pid): if pid <= 0: return False if _prime_agent_sys.platform == "win32": try: import ctypes SYNCHRONIZE = 0x00100000 PROCESS_QUERY_LIMITED_INFORMATION = 0x1000 handle = ctypes.windll.kernel32.OpenProcess( SYNCHRONIZE | PROCESS_QUERY_LIMITED_INFORMATION, False, pid ) if handle: ctypes.windll.kernel32.CloseHandle(handle) return True return False except Exception: return True try: _prime_agent_os.kill(pid, 0) return True except ProcessLookupError: return False except PermissionError: return True except OSError as err: import errno if getattr(err, "errno", None) == errno.ESRCH: return False if getattr(err, "errno", None) == errno.EPERM: return True return False
def _prime_agent_parse_lock(lock): try: with open(lock, "r", encoding="utf-8") as _f: content = _f.read() except OSError: return None trimmed = content.strip() if not trimmed: return None parts = trimmed.split(":") if ":" in trimmed else trimmed.split() first = parts[0].strip() if first.isdigit(): pid = int(first) token = parts[1].strip() if len(parts) > 1 else None return pid, token return None
def _prime_agent_build_in_progress(): """True while a live build holds the lock. A lock with an alive owner pid is live; an unidentifiable lock falls back to the age threshold. Dead locks are cleaned up.""" lock = _prime_agent_bootstrap_lock_path() if lock is None: return False owner = _prime_agent_parse_lock(lock) if owner is not None: pid, _ = owner if _prime_agent_pid_is_alive(pid): return True try: _prime_agent_os.unlink(lock) except OSError: pass return False try: age = _prime_agent_time.time() - _prime_agent_os.stat(lock).st_mtime except OSError: return False if age < _PRIME_AGENT_BOOTSTRAP_STALE_LOCK_SECONDS: return True try: _prime_agent_os.unlink(lock) except OSError: pass return False
def _prime_agent_wait_for_build(): """Wait out a concurrent venv build, so this kernel imports a complete tree, never a partial one.""" if _prime_agent_build_in_progress(): print( f"prime-agent: waiting up to {_PRIME_AGENT_BOOTSTRAP_WAIT_SECONDS:g}s " f"for the kernel venv build at {_prime_agent_bootstrap_lock_path()} to finish", flush=True, ) deadline = _prime_agent_time.monotonic() + _PRIME_AGENT_BOOTSTRAP_WAIT_SECONDS while _prime_agent_build_in_progress() and _prime_agent_time.monotonic() < deadline: _prime_agent_time.sleep(_PRIME_AGENT_BOOTSTRAP_POLL_SECONDS) return not _prime_agent_build_in_progress()
def _prime_agent_bind_runtime(): """Publish rlm, bash, and mcp into the cell namespace, and isolate optional modules. False keeps what is bound now.""" global _PRIME_AGENT_RLM_IMPORT_ERROR _prime_agent_wait_for_build() try: import rlm as _prime_agent_rlm_module import rlm.mcp as _prime_agent_mcp_module except Exception as _prime_agent_rlm_error: _PRIME_AGENT_RLM_IMPORT_ERROR = str(_prime_agent_rlm_error) return False globals().update( rlm=_prime_agent_rlm_module.rlm, bash=_prime_agent_rlm_module.bash, mcp=_prime_agent_mcp_module, ) _PRIME_AGENT_RLM_IMPORT_ERROR = None try: import remote as _prime_agent_remote_module globals()["remote"] = _prime_agent_remote_module except Exception as _prime_agent_remote_error: globals()["remote"] = _PrimeAgentUnavailableSkill("remote", str(_prime_agent_remote_error)) return True
def _prime_agent_missing_runtime_error(): message = ( "prime-agent-runtime is not installed in this kernel. " "Remove ~/.prime/agent/kernel-venv so prime-agent can rebuild it, or set " "PRIME_AGENT_KERNEL_PYTHON to a kernel environment with prime-agent-runtime installed. " f"Import error: {_PRIME_AGENT_RLM_IMPORT_ERROR}" ) if _prime_agent_build_in_progress(): message += ( f" A kernel venv build is still in progress at {_prime_agent_bootstrap_lock_path()}; " "call again once it finishes." ) return RuntimeError(message)
def _prime_agent_require_runtime(): """The module namespace once the runtime imports, else the teaching error.""" if _prime_agent_bind_runtime(): return globals() raise _prime_agent_missing_runtime_error()
_PRIME_AGENT_RUNTIME_READY = _prime_agent_bind_runtime()
if not _PRIME_AGENT_RUNTIME_READY:
class _PrimeAgentMissingRlm: """Stands in for the runtime until it imports, then rebinds itself out of the namespace."""
def _rebind(self): _prime_agent_require_runtime() return globals()["rlm"]
def __getattr__(self, name): target = self._rebind() return getattr(target, name)
async def spawn(self, prompt, **kwargs): return await self._rebind().spawn(prompt, **kwargs)
async def models(self): return await self._rebind().models()
async def create_session(self, prompt, **kwargs): return await self._rebind().create_session(prompt, **kwargs)
async def list_subagents(self): return await self._rebind().list_subagents()
async def delete_subagent(self, target): return await self._rebind().delete_subagent(target)
rlm = _PrimeAgentMissingRlm()
def bash(command, *args, **kwargs): return _prime_agent_require_runtime()["bash"](command, *args, **kwargs)`.trim();
export function buildRlmBootstrapCode(pythonSkills: readonly PythonSkillRuntimeInfo[] = []): string { const baseCode = [RLM_BOOTSTRAP_HEADER_CODE, RLM_BOOTSTRAP_RUNTIME_CODE].join("\n\n"); const importNames = [...new Set(pythonSkills.map((skill) => skill.importName))]; if (importNames.length === 0) { return baseCode; }
return `${baseCode}
for _prime_agent_skill_name in ${JSON.stringify(importNames)}: try: globals()[_prime_agent_skill_name] = _prime_agent_wrap_skill_module( _prime_agent_importlib.import_module(_prime_agent_skill_name) ) except Exception as _prime_agent_skill_error: _PRIME_AGENT_SKILL_IMPORT_ERRORS[_prime_agent_skill_name] = str(_prime_agent_skill_error) globals()[_prime_agent_skill_name] = _PrimeAgentUnavailableSkill( _prime_agent_skill_name, str(_prime_agent_skill_error), )
if "search" in globals() and hasattr(globals()["search"], "ls"): globals()["ls"] = globals()["search"].ls`.trim();}const ipythonSchema = Type.Object({ code: Type.String({ description: "Python code to execute in the persistent Python REPL. Use the target project's own environment for project imports, tests, scripts, CLIs, and dependency checks instead of direct kernel imports.", }), timeout: Type.Optional( Type.Number({ description: "Cell deadline in seconds. Omitted: this session's cell deadline, shown in the tool description. Max 3600; a larger value is clamped to it. Set it for work that is legitimately slow (a long build, a full test run).", minimum: 1, maximum: CELL_TIMEOUT_MAX_SECONDS, }), ),});
const BUSY_KERNEL_WAIT_CHOICE = "Wait and preserve state";const BUSY_KERNEL_KILL_CHOICE = "Kill kernel and restart";const BUSY_KERNEL_PROMPT = [ "Interrupted Python cell is still running", "Ctrl+C sent an interrupt, but the previous cell has not stopped yet. A new command cannot start until it finishes.", "Waiting preserves the current kernel state. Killing restarts the kernel and loses in-memory variables, imports, and running tasks.",].join("\n");const KERNEL_RESTART_NOTICE = [ "<ipython_kernel_reset>", "The Python kernel was restarted after a previous interrupted cell kept running. Variables, imports, async tasks, and open resources from before the restart are no longer available; recreate them before using them.", "</ipython_kernel_reset>",].join("\n");
const CELL_TIMEOUT_DETACH_ADVICE = 'If you were waiting on work, do not poll inside a cell: start it detached and end your turn -- h = bash("nix build .#prime-agent") returns at once, and its completion notice wakes you with the result.';
/** What one cell asked for and what it actually ran under. All in ms; seconds only at the edge. */interface ResolvedCellDeadline { /** The model's `timeout`, or the session default when it passed none. */ requestedMs?: number; /** The deadline the kernel timer arms with. */ effectiveMs?: number; /** The request was above the ceiling and ran at the ceiling instead. */ clamped: boolean;}
function resolveCellDeadline( timeoutSeconds: number | undefined, sessionDeadlineMs: number | undefined,): ResolvedCellDeadline { if (timeoutSeconds === undefined) { return sessionDeadlineMs === undefined ? { clamped: false } : { requestedMs: sessionDeadlineMs, effectiveMs: sessionDeadlineMs, clamped: false }; } const requestedMs = Math.round(timeoutSeconds * 1000); const effectiveMs = Math.min(Math.max(requestedMs, 1000), CELL_TIMEOUT_MAX_SECONDS * 1000); return { requestedMs, effectiveMs, clamped: requestedMs > effectiveMs };}
function cellClampNotice(requestedMs: number): string { return `timeout=${requestedMs / 1000}s exceeds the ${CELL_TIMEOUT_MAX_SECONDS}s ceiling; this cell used ${CELL_TIMEOUT_MAX_SECONDS}s.`;}
/** Largest deadline still worth suggesting: an order of magnitude up, capped, and only when it is actually larger. */function cellRetryDeadlineSeconds(effectiveMs: number, repeated: boolean): number | undefined { if (repeated) { return undefined; } const effectiveSeconds = effectiveMs / 1000; const suggested = Math.min(CELL_TIMEOUT_MAX_SECONDS, Math.round(effectiveSeconds * 10)); return suggested > effectiveSeconds ? suggested : undefined;}
interface CellTimeoutNoticeInput { /** The deadline this cell ran under. */ deadlineSeconds: number; /** Wall time from taking the execution slot to the result. */ elapsedSeconds: number; kernelRestarted: boolean; retryDeadlineSeconds?: number; clampNotice?: string;}
/** * The model reads `content`, never `details`, so a timeout has to be stated in the * cell's text. Detach-first on purpose: the detached handle is the one long-work * path the cell deadline cannot reach. */function buildCellTimeoutNotice(input: CellTimeoutNoticeInput): string { const lines = [ "<ipython_cell_timeout>", `Cell deadline ${input.deadlineSeconds}s hit at ${input.elapsedSeconds.toFixed(1)}s; the cell was interrupted. Nothing from before it is lost: earlier cells' variables and imports survive, and everything this cell printed or already wrote is kept.`, ]; if (input.kernelRestarted) { lines.push("The kernel was killed and restarted because the cell did not stop."); } lines.push("cell: for _ in range(60): await asyncio.sleep(5)"); lines.push( input.retryDeadlineSeconds === undefined ? CELL_TIMEOUT_DETACH_ADVICE : `${CELL_TIMEOUT_DETACH_ADVICE} If it was one genuinely long call (cargo/nix build, test suite, training), re-run it with a larger deadline: timeout=${input.retryDeadlineSeconds} (seconds, max ${CELL_TIMEOUT_MAX_SECONDS}).`, ); if (input.clampNotice) { lines.push(input.clampNotice); } lines.push("</ipython_cell_timeout>"); return lines.join("\n");}
function createAbortError(): Error { return new Error("Python execution aborted");}
function raceWithAbort<T>(promise: Promise<T>, signal: AbortSignal | undefined, onAbort?: () => void): Promise<T> { if (!signal) { return promise; } if (signal.aborted) { onAbort?.(); return Promise.reject(createAbortError()); } return new Promise<T>((resolve, reject) => { let settled = false; const cleanup = () => signal.removeEventListener("abort", abort); const abort = () => { if (settled) { return; } settled = true; cleanup(); onAbort?.(); reject(createAbortError()); }; signal.addEventListener("abort", abort, { once: true }); promise.then( (value) => { if (settled) { return; } settled = true; cleanup(); resolve(value); }, (error: unknown) => { if (settled) { return; } settled = true; cleanup(); reject(error); }, ); });}
function createLinkedAbortSignal(sources: readonly (AbortSignal | undefined)[]): { signal: AbortSignal; cleanup: () => void;} { const controller = new AbortController(); const cleanups: Array<() => void> = []; const abort = () => controller.abort(); for (const source of sources) { if (!source) { continue; } if (source.aborted) { controller.abort(); continue; } const listener = () => abort(); source.addEventListener("abort", listener, { once: true }); cleanups.push(() => source.removeEventListener("abort", listener)); } return { signal: controller.signal, cleanup: () => { for (const cleanup of cleanups) { cleanup(); } }, };}
function setWorkingMessage(ctx: ExtensionContext | undefined, message?: string): void { try { ctx?.ui.setWorkingMessage(message); } catch { // Stale UI context; cosmetic only. }}
export type IpythonToolInput = Static<typeof ipythonSchema>;
export interface IpythonToolDetails { durationMs?: number; status?: "ok" | "error" | "aborted" | "starting"; errorEname?: string; stdout?: string; stderr?: string; result?: string; /** Output that arrived without this cell's id (threads, other cells' leftovers), shown separately from stdout. */ backgroundOutput?: string; /** Diffs streamed from file edits, rendered by the cell view. */ diffs?: KernelDiffDisplay[]; /** Rich content displays emitted to the user (e.g. by the display skill). */ displays?: KernelContentDisplay[]; /** Media attachments loaded into context (e.g. by the attach-image skill). */ attachments?: KernelAttachment[]; /** Agent messages sent from this cell. */ sentAgentMessages?: KernelSentAgentMessage[]; /** True when this result came after killing and restarting a busy kernel. */ kernelRestarted?: boolean; /** True when the cell's own deadline interrupted it (the user did not abort). */ timedOut?: boolean; /** The deadline the cell ran under when it timed out, in ms. */ timeoutMs?: { requested: number; effective: number }; error?: { ename: string; evalue: string; traceback: string[]; };}
export interface IpythonToolOptions { /** Python override. Must have prime-agent-runtime installed. */ python?: string; env?: Record<string, string>; /** Command prefix prepended to every bash() command. */ commandPrefix?: string; /** Deadline for one ipython cell in ms; "off" disables it. Default: DEFAULT_CELL_DEADLINE_MS. */ cellTimeoutMs?: number | "off"; /** Shell used by bash(). */ shellPath?: string; sessionId?: string; /** Typed host request handlers for the kernel↔host bridge (rlm.run, goal.*, …). */ hostHandlers?: HostRequestHandlers; pythonSkills?: readonly PythonSkillRuntimeInfo[]; /** Per-session artifact dir where the kernel namespace snapshot is stored. Omit to disable snapshots. */ snapshotDir?: string; /** Resolves before this kernel starts — e.g. the previous provisioner's dispose, so a * /reload's old-kernel snapshot flush can't race the new kernel's restore. */ readyGate?: Promise<unknown>; /** * Fires once per kernel start when a previous session's namespace was revived * (some names restored or some failed), so the session can tell the model. */ onRestore?: (result: RestoreResult) => void; /** * Fires once per kernel start when the runtime advertises the `signatures` capability * and answers with a generated skill-surface digest, so the session can publish it in * the system prompt. Silent when the runtime does not advertise it. */ onSkillDigest?: (digest: string) => void; onLateSentAgentMessage?: (toolCallId: string, message: KernelSentAgentMessage) => void; /** Remote host configuration when running on a remote host over SSH. */ remote?: RemoteHostConfig; /** Shared provisioner owning the kernel lifecycle. When provided, the remaining options are ignored. */ provisioner?: IpythonKernelProvisioner;}
/** * Owns the lazy create+start+runtime-bootstrap of one session's Python kernel. * * Concurrent ensure() calls await the same in-flight startup, a failed startup * clears the memo so the next call retries fresh, and progress listeners can * attach mid-flight (a tool call racing a background prewarm()). */export class IpythonKernelProvisioner { private managerPromise?: Promise<KernelClient>; private startedManager?: KernelClient; private readonly startupListeners = new Set<KernelBootstrapProgressHandler>(); private lastStartupMessage?: string; private _lastRestore?: RestoreResult; private _skillDigest?: string; private readonly disposeController = new AbortController(); /** Snapshot policy of the dispose that aborted a startup, honored by startKernel's failure teardown. */ private disposeSnapshot = true;
constructor( private readonly cwd: string, private readonly options?: Omit<IpythonToolOptions, "provisioner">, ) {}
/** The kernel manager, once a startup has completed successfully. */ get manager(): KernelClient | undefined { return this.startedManager; }
/** Result of reviving a prior session's namespace on the last kernel start, if any. */ get lastRestore(): RestoreResult | undefined { return this._lastRestore; }
/** Generated skill-surface digest from the running kernel, or undefined when absent. */ get skillDigest(): string | undefined { return this._skillDigest; }
/** Start the kernel in the background. Failures are swallowed here and surface on the next ensure(). */ prewarm(): void { void this.ensure().catch(() => {}); }
/** Whether a kernel has finished starting and is currently running. */ get hasRunningKernel(): boolean { return this.startedManager?.isRunning ?? false; }
/** Remove live variables above the snapshot's per-variable size limit. */ async pruneOversizedVariables(): Promise<string[] | null> { const m = this.startedManager ?? (await this.managerPromise?.catch(() => undefined)); const result = await m?.pruneOversizedVariables(); return result ? (result.pruned ?? []) : null; }
/** Live user-defined names in the kernel namespace, or null if listing failed / no kernel. */ async listNamespaceNames(signal?: AbortSignal): Promise<string[] | null> { const m = this.startedManager ?? (await this.managerPromise?.catch(() => undefined)); return (await m?.listNamespaceNames(signal)) ?? null; }
/** Dispose the kernel owned by this provisioner, including one still starting up. */ async dispose(options?: { snapshot?: boolean }): Promise<void> { this.disposeSnapshot = options?.snapshot ?? true; // Drops a still-queued boot out of the semaphore and short-circuits an // in-flight startKernel before it spawns, so a disposed session's boot // doesn't waste a slot during a fan-out. this.disposeController.abort(); const pending = this.managerPromise; this.managerPromise = undefined; this.startedManager = undefined; if (!pending) return; try { const m = await pending; await m.shutdown({ snapshot: this.disposeSnapshot, drainHostRequests: true }); } catch { // a failed startup already cleaned up after itself } }
async kill(): Promise<void> { const pending = this.managerPromise; this.managerPromise = undefined; this.startedManager = undefined; if (!pending) return; try { const m = await pending; await m.kill(); } catch { // a failed startup already cleaned up after itself } }
ensure(onProgress?: KernelBootstrapProgressHandler, signal?: AbortSignal): Promise<KernelClient> { if (signal?.aborted) { return Promise.reject(createAbortError()); } // Only a terminally dead kernel drops the memo; a repairing manager (idle/starting) recovers itself. if (this.startedManager?.isDefunct) { this.managerPromise = undefined; this.startedManager = undefined; } let cleanupProgressListener: (() => void) | undefined; if (onProgress && !this.startedManager) { this.startupListeners.add(onProgress); cleanupProgressListener = () => { this.startupListeners.delete(onProgress); signal?.removeEventListener("abort", cleanupProgressListener!); }; signal?.addEventListener("abort", cleanupProgressListener, { once: true }); // Joining an in-flight startup: replay the current stage. if (this.managerPromise && this.lastStartupMessage) { onProgress(this.lastStartupMessage); } } if (!this.managerPromise) { const startup = this.startKernel(signal); this.managerPromise = startup; startup.then( (m) => { if (this.managerPromise === startup) { this.startedManager = m; } this.settleStartup(); }, () => { // Clear the memo so the next ensure() retries instead of // rethrowing a cached rejection forever. if (this.managerPromise === startup) { this.managerPromise = undefined; } this.settleStartup(); }, ); } return raceWithAbort(this.managerPromise, signal).finally(() => { cleanupProgressListener?.(); }); }
private settleStartup(): void { this.startupListeners.clear(); this.lastStartupMessage = undefined; }
private emitStartupProgress(message: string): void { this.lastStartupMessage = message; for (const listener of [...this.startupListeners]) { listener(message); } }
/** The import names the bootstrap publishes into the cell namespace, deduplicated. */ private skillImportNames(): string[] { return [...new Set((this.options?.pythonSkills ?? []).map((skill) => skill.importName))]; }
private async startKernel(signal?: AbortSignal): Promise<KernelClient> { const startupAbort = createLinkedAbortSignal([this.disposeController.signal, signal]); const startupSignal = startupAbort.signal; // A restart replaces the modules the old digest described, so the stale one never // outlives the kernel that published it. this._skillDigest = undefined; // Wait for a previous provisioner (e.g. on /reload) to finish disposing — and // flushing its final snapshot — before we read that snapshot back, so the two // kernels can't race over the same on-disk file. Guarded so the common // no-gate path stays synchronous (callers rely on prompt startup progress). try { if (this.options?.readyGate) { await raceWithAbort( this.options.readyGate.catch(() => {}), startupSignal, ); } const snapshotDir = this.options?.snapshotDir; // Always inject an absolute trusted shell (undefined only on win32 // without bash, where the runtime's teaching error fires instead). const shellPath = resolveKernelBashShell(this.options?.shellPath); const commandPrefix = this.options?.commandPrefix; const bootstrapCode = buildRlmBootstrapCode(this.options?.pythonSkills); const m = new ReplKernelManager({ python: this.options?.python, cwd: this.cwd, // bash() reads these to pick its shell and command prefix. env: { ...this.options?.env, ...(shellPath ? { PRIME_AGENT_BASH_SHELL: shellPath } : {}), ...(commandPrefix ? { PRIME_AGENT_BASH_COMMAND_PREFIX: commandPrefix } : {}), }, sessionId: this.options?.sessionId, hostHandlers: this.options?.hostHandlers, pythonSkills: this.options?.pythonSkills, // Only persistent sessions (which have an artifact dir) get a revivable snapshot. snapshot: snapshotDir ? { path: snapshotPathIn(snapshotDir), manifestPath: manifestPathIn(snapshotDir) } : undefined, stderrLogPath: snapshotDir ? join(snapshotDir, "kernel-stderr.log") : undefined, bootstrapCode, remote: this.options?.remote, }); let pendingRestore: RestoreResult | undefined; try { // Emitted synchronously (before the permit await) so a listener attaching // mid-flight can replay the current stage. this.emitStartupProgress("Starting Python kernel..."); // Only the process spawn + port resolve contends for OS resources under a // fan-out, and it is bounded by start()'s own timeouts — so the permit // covers only start(). Restore/bootstrap run per-kernel afterwards and are // unbounded execute()s; holding the global permit across them could pin it // forever on a wedged bootstrap and starve every other session's boot. await withKernelBootPermit(() => { // Disposed while queued for the permit — don't spawn a kernel nobody wants. if (startupSignal.aborted) throw new Error("Kernel provisioner disposed before start"); return m.start({ onBootstrapProgress: (message) => this.emitStartupProgress(message), signal: startupSignal, }); }, startupSignal); // Revive a prior session's namespace before the bootstrap, so the bootstrap // then overwrites live handles (rlm, skills) on top of anything restored. if (snapshotDir) { const snapshotExisted = existsSync(snapshotPathIn(snapshotDir)); this.emitStartupProgress("Restoring Python state..."); const restore = await raceWithAbort(m.restoreState(), startupSignal); if (snapshotExisted) { pendingRestore = restore ?? { restored: [], failed: [], path: snapshotPathIn(snapshotDir) }; } } this.emitStartupProgress("Preparing Python runtime..."); const bootstrap = await m.execute(bootstrapCode, { signal: startupSignal, }); if (bootstrap.status !== "ok") { const details = [bootstrap.stderr, bootstrap.error?.traceback.join("\n")].filter(Boolean).join("\n"); throw new Error(`Failed to initialize rlm runtime in the Python kernel:\n${details}`); } // The digest indexes the wrapped skill modules, so it is read after the // bootstrap cell that installs them. It is a kernel-start artifact: cached // here so a system-prompt rebuild stays synchronous. this.emitStartupProgress("Indexing the skill surface..."); const digest = await m.skillDigest(this.skillImportNames(), startupSignal).catch(() => null); if (digest) { this._skillDigest = digest; this.options?.onSkillDigest?.(digest); } } catch (error) { // Never leak the kernel process if startup fails after spawn — and never // surface the failure before the teardown (final snapshot flush included) // finished, or a replacement provisioner gated on this dispose could // race the still-flushing kernel over the same snapshot files. await m.shutdown({ snapshot: this.disposeSnapshot, drainHostRequests: true }).catch(() => undefined); throw error; } // Only tell the model what was revived once the kernel is actually usable — // a notice claiming restored state must never outlive a failed bootstrap. if (pendingRestore) { this._lastRestore = pendingRestore; this.options?.onRestore?.(pendingRestore); } return m; } finally { startupAbort.cleanup(); } }}
async function chooseBusyKernelAction( ctx: ExtensionContext | undefined, signal: AbortSignal | undefined,): Promise<"wait" | "kill" | "cancel"> { if (!ctx?.hasUI) { return "cancel"; } const choice = await ctx.ui.select(BUSY_KERNEL_PROMPT, [BUSY_KERNEL_WAIT_CHOICE, BUSY_KERNEL_KILL_CHOICE], { signal, }); if (choice === BUSY_KERNEL_WAIT_CHOICE) { return "wait"; } if (choice === BUSY_KERNEL_KILL_CHOICE) { return "kill"; } return "cancel";}
async function executeWithBusyKernelChoice( provisioner: IpythonKernelProvisioner, reportStartupProgress: KernelBootstrapProgressHandler, toolCallId: string, code: string, signal: AbortSignal | undefined, onStream: (chunk: string, name: "stdout" | "stderr") => void, onWorkingMessage: (message?: string) => void, onLateSentAgentMessage: ((toolCallId: string, message: KernelSentAgentMessage) => void) | undefined, cellDeadlineMs: number | undefined, ctx: ExtensionContext | undefined,): Promise<{ result: ExecuteResult; kernelRestarted: boolean }> { let kernelRestarted = false; while (true) { const m = await provisioner.ensure(reportStartupProgress, signal); try { const result = await m.execute(code, { signal, timeoutMs: cellDeadlineMs, onStream, onLateSentAgentMessage: onLateSentAgentMessage ? (message) => onLateSentAgentMessage(toolCallId, message) : undefined, }); if (result.timedOut) { // The deadline interrupted the cell, but the runtime may still be running // it. Waiting without bound here would just move the hang, and asking the // user would hang a session with no UI, so escalation is mechanical. if (!(await m.waitForInterruptedCellToSettle())) { onWorkingMessage("Restarting Python kernel..."); await provisioner.kill(); kernelRestarted = true; } } return { result, kernelRestarted }; } catch (error) { if (!(error instanceof KernelBusyAfterInterruptError) || signal?.aborted) { throw error; } const action = await chooseBusyKernelAction(ctx, signal); if (action === "wait") { onWorkingMessage("Waiting for Python kernel..."); continue; } if (action === "kill") { onWorkingMessage("Restarting Python kernel..."); await provisioner.kill(); kernelRestarted = true; continue; } throw error; } }}
/** Turn kernel image attachments into `ImageContent` blocks; non-image types are dropped. */export function imageBlocksFromAttachments(attachments: readonly KernelAttachment[] | undefined): ImageContent[] { if (!attachments) return []; return attachments .filter((a) => IMAGE_MIME_TYPES.has(a.mimeType)) .map((a) => ({ type: "image", data: a.data, mimeType: a.mimeType }));}
export function createIpythonToolDefinition( cwd: string, options?: IpythonToolOptions,): ToolDefinition<typeof ipythonSchema, IpythonToolDetails> { const provisioner = options?.provisioner ?? new IpythonKernelProvisioner(cwd, options); // The session supplies the host default; "off" opts out, and no option at all // still gets the default so a tool built elsewhere cannot hang forever. const configuredCellDeadlineMs = options?.cellTimeoutMs ?? DEFAULT_CELL_DEADLINE_MS; const cellDeadlineMs = configuredCellDeadlineMs === "off" ? undefined : configuredCellDeadlineMs; const cellDeadlineSummary = cellDeadlineMs === undefined ? ` This session sets no cell deadline (settings cellDeadlineMs is "off"); pass timeout=<seconds> (max ${CELL_TIMEOUT_MAX_SECONDS}) to bound one call.` : ` Cells stop at ${cellDeadlineMs / 1000}s by default (max ${CELL_TIMEOUT_MAX_SECONDS}s); pass timeout=<seconds> to change one call.`; let lastTimedOutCode: string | undefined;
return { name: "ipython", label: "ipython", description: "Execute Python code in a persistent Python REPL. Top-level `await` is supported. Variables, imports, and loaded data persist across calls, and are revived on a best-effort basis when a session is resumed (objects that cannot be serialized are dropped and reported). Run shell commands with `bash('cmd')` / `await bash('cmd')`. Project imports, tests, scripts, CLIs, and dependency checks should run through the target project's own environment." + cellDeadlineSummary, promptSnippet: "ipython - persistent Python REPL for code, state, and bash() orchestration", // The kernel is single-threaded — pi must not run two ipython calls in parallel within a batch. executionMode: "sequential", parameters: ipythonSchema, execute: async (toolCallId, params, signal, onUpdate, ctx) => { let hasWorkingMessage = false; const setToolWorkingMessage = (message?: string) => { setWorkingMessage(ctx, message); hasWorkingMessage = message !== undefined; }; const reportStartupProgress: KernelBootstrapProgressHandler = (message) => { setToolWorkingMessage(message); onUpdate?.({ content: [{ type: "text", text: message }], details: { status: "starting" }, }); };
const deadline = resolveCellDeadline(params.timeout, cellDeadlineMs); try { const { result: r, kernelRestarted } = await executeWithBusyKernelChoice( provisioner, reportStartupProgress, toolCallId, params.code, signal, (chunk) => { onUpdate?.({ content: [{ type: "text", text: chunk }], details: { status: "ok" }, }); }, setToolWorkingMessage, options?.onLateSentAgentMessage, deadline.effectiveMs, ctx, );
let text = r.stdout; if (r.stderr) text += (text ? "\n" : "") + r.stderr; if (r.result) text += (text ? "\n" : "") + r.result; if (r.status === "error" && r.error) { text += (text ? "\n" : "") + r.error.traceback.join("\n"); } if (r.backgroundOutput) { text += `${text ? "\n" : ""}[background output (unattributed)]\n${r.backgroundOutput}`; } if (kernelRestarted) { text = text ? `${KERNEL_RESTART_NOTICE}\n\n${text}` : KERNEL_RESTART_NOTICE; } const clampNotice = deadline.clamped && deadline.requestedMs !== undefined ? cellClampNotice(deadline.requestedMs) : undefined; let notice = clampNotice; if (r.timedOut) { const deadlineMs = deadline.effectiveMs ?? r.durationMs; const repeated = params.code === lastTimedOutCode; lastTimedOutCode = params.code; notice = buildCellTimeoutNotice({ deadlineSeconds: deadlineMs / 1000, elapsedSeconds: r.durationMs / 1000, kernelRestarted, retryDeadlineSeconds: cellRetryDeadlineSeconds(deadlineMs, repeated), clampNotice, }); } if (notice) { text = text ? `${notice}\n\n${text}` : notice; }
const imageBlocks = imageBlocksFromAttachments(r.attachments); const content: (TextContent | ImageContent)[] = [{ type: "text", text: text || "" }, ...imageBlocks];
return { content, details: { durationMs: r.durationMs, status: r.status, errorEname: r.error?.ename, stdout: r.stdout, stderr: r.stderr, result: r.result, backgroundOutput: r.backgroundOutput, diffs: r.diffs, displays: r.displays, attachments: r.attachments, sentAgentMessages: r.sentAgentMessages, kernelRestarted, timedOut: r.timedOut, timeoutMs: r.timedOut && deadline.effectiveMs !== undefined ? { requested: deadline.requestedMs ?? deadline.effectiveMs, effective: deadline.effectiveMs } : undefined, error: r.error, }, isError: r.status === "error" || r.status === "aborted", }; } finally { if (hasWorkingMessage) { setToolWorkingMessage(); } } }, };}
export function createIpythonTool(cwd: string, options?: IpythonToolOptions): AgentTool<typeof ipythonSchema> { return wrapToolDefinition(createIpythonToolDefinition(cwd, options));}