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.
19 kB · 506 lines
TypeScript
at main
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507import type { ThinkingLevel } from "@earendil-works/pi-agent-core";import type { Api, Model, ServiceTier } from "@earendil-works/pi-ai";import type { AgentSession, RlmChildAgentStatus } from "./agent-session.js";import type { ToolDefinition } from "./extensions/index.js";import type { HostRequestHandler } from "./kernel/index.js";import { THINKING_LEVELS } from "./thinking-levels.js";
/** Request emitted by `rlm.spawn`; cellSourceCode preserves the spawning cell for display. */export interface RlmRunRequest { prompt: string; kwargs: Record<string, unknown>; cellSourceCode?: string;}
interface RlmCreateSessionRequest { prompt: string; kwargs: Record<string, unknown>;}
export interface RlmCreateSessionResult { active_session_id: string; session_id: string; name: string; session_file: string; model: string;}
export interface RlmSpawnHandle { rlm_child_id: string; name: string; session_dir: string; model: string;}
export type RlmSubagentRegistryStatus = "running" | "completed" | "error";
/** * Kernel-wire shape of a child activity snapshot: snake_case like the rest of * the registry, so `JSON.stringify` needs no key rewrite on the Python side. */export interface RlmSubagentRegistryActivity { kind: "waiting" | "writing" | "executing"; tool_name?: string;}
export interface RlmSubagentRegistryEntry { rlm_child_id: string; active_session_id: string | null; session_id: string | null; session_name: string; session_dir: string; status: RlmSubagentRegistryStatus; /** Live-state extras, present when the run or retained session is locally available. */ activity?: RlmSubagentRegistryActivity; tool_use_count?: number; duration_ms?: number; /** Compacted answer preview, hard-capped for the kernel roster. */ answer_preview?: string; replied_since_task?: boolean; /** Latest child progress note (`rlm.progress.note`), newest wins. */ progress_note?: string; /** One-line task label, hard-capped for the kernel roster. */ label?: string; /** Wall-clock ms of the last tracked child activity; seeded at admission. */ last_activity_at?: number; /** Set when a running child has had no tracked activity for the staleness threshold. */ activity_stale_ms?: number; /** Pinned model selector of the child session, when the parent knows one. */ model?: string; /** Parent-visible run accounting: assistant turns, tokens, and parent follow-ups. */ turns?: number; tokens?: number; followups?: number;}
export interface RlmListSubagentsResult { subagents: RlmSubagentRegistryEntry[];}
export interface RlmDeleteSubagentResult { subagent: RlmSubagentRegistryEntry; outcome?: "deleted" | "skipped_running";}
export interface RlmModelMatch { provider: string; id: string; name: string; selector: string;}
export interface RlmCollectResultEntry { rlm_child_id: string; session_name: string | undefined; session_dir: string; /** Raw run status: queued | running | done | error | cancelled. */ status: RlmChildAgentStatus; /** True once the run reached a terminal state (settlement resolved or rejected). */ settled: boolean; answer_preview: string | undefined; error: string | undefined; duration_ms: number | undefined; tool_use_count: number | undefined; replied_since_task: boolean | undefined;}
export interface RlmCollectResult { results: RlmCollectResultEntry[];}
export type RlmCollectHandler = (targets: string[], timeoutMs: number) => Promise<RlmCollectResult>;export type RlmRunHandler = (request: RlmRunRequest) => Promise<Record<string, unknown>>;type RlmCreateSessionHandler = (request: RlmCreateSessionRequest) => Promise<RlmCreateSessionResult>;
interface AsyncBashCompletionRequest { pid: number; command: string; exitCode: number;}
type AsyncBashCompletionHandler = (request: AsyncBashCompletionRequest) => void | Promise<void>;
interface AsyncBashConsumedRequest { pid: number; command: string;}
type AsyncBashConsumedHandler = (request: AsyncBashConsumedRequest) => void | Promise<void>;export type RlmListSubagentsHandler = () => RlmListSubagentsResult | Promise<RlmListSubagentsResult>;export type RlmDeleteSubagentHandler = (target: string) => Promise<RlmDeleteSubagentResult>;export type RlmModelsHandler = () => Record<string, string[]> | Promise<Record<string, string[]>>;
const RLM_SUBAGENT_SESSION_NAME_MAX_LENGTH = 64;const RLM_MODEL_ERROR_SUGGESTION_LIMIT = 3;
export function normalizeRequestedRlmSubagentSessionName(value: unknown, operation = "rlm.spawn"): string | undefined { if (value === undefined) { return undefined; } if (typeof value !== "string") { throw new Error(`${operation} name must be a string`); } const name = value.trim(); if (!name) { throw new Error(`${operation} name must not be empty`); } if (name.length > RLM_SUBAGENT_SESSION_NAME_MAX_LENGTH) { throw new Error(`${operation} name must be at most ${RLM_SUBAGENT_SESSION_NAME_MAX_LENGTH} characters`); } return name;}
export function normalizeRequestedRlmSubagentThinkingLevel( value: unknown, operation = "rlm.spawn",): ThinkingLevel | undefined { if (value === undefined) { return undefined; } if (typeof value !== "string") { throw new Error(`${operation} thinking must be a string`); } const level = value.trim().toLowerCase(); if (!THINKING_LEVELS.includes(level as ThinkingLevel)) { throw new Error(`${operation} thinking must be one of: ${THINKING_LEVELS.join(", ")}`); } return level as ThinkingLevel;}
export function normalizeRequestedRlmSubagentModel(value: unknown, operation = "rlm.spawn"): string | undefined { if (value === undefined) { return undefined; } if (typeof value !== "string") { throw new Error(`${operation} model must be a string`); } const model = value.trim(); if (!model) { throw new Error(`${operation} model must not be empty`); } return model;}
/** Create a readable, collision-resistant default name usable as an agent-message selector. */export function createDefaultRlmSubagentSessionName(prompt: string, childId: string): string { const promptSlug = prompt .normalize("NFKD") .replace(/[\u0300-\u036f]/g, "") .toLowerCase() .replace(/[^a-z0-9]+/g, "-") .replace(/^-+|-+$/g, ""); const idSuffix = childId .replace(/^sub-/, "") .replace(/[^A-Za-z0-9]+/g, "") .slice(-8) || "child"; const fixedLength = "subagent--".length + idSuffix.length; const promptPart = (promptSlug || "worker") .slice(0, Math.max(1, RLM_SUBAGENT_SESSION_NAME_MAX_LENGTH - fixedLength)) .replace(/-+$/g, ""); return `subagent-${promptPart || "worker"}-${idSuffix}`;}
function isRecord(value: unknown): value is Record<string, unknown> { return typeof value === "object" && value !== null && !Array.isArray(value);}
function normalizeModelSearchText(value: string): string { return value.toLowerCase().replace(/[^a-z0-9]+/g, "");}
export function findRlmModelMatches(query: string, models: Model<Api>[], limit: number): RlmModelMatch[] { const normalizedQuery = normalizeModelSearchText(query.trim()); return models .map((model) => { const selector = `${model.provider}/${model.id}`; const fields = [selector, model.id, model.name || model.id]; const normalizedFields = fields.map(normalizeModelSearchText); let score = normalizedQuery ? Number.POSITIVE_INFINITY : 0; if (normalizedQuery) { const exactIndex = normalizedFields.indexOf(normalizedQuery); const prefixIndex = normalizedFields.findIndex((field) => field.startsWith(normalizedQuery)); const partialIndex = normalizedFields.findIndex((field) => field.includes(normalizedQuery)); if (exactIndex >= 0) score = exactIndex; else if (prefixIndex >= 0) score = 3 + prefixIndex; else if (partialIndex >= 0) score = 6 + partialIndex; } return { model, selector, score }; }) .filter((candidate) => Number.isFinite(candidate.score)) .sort((a, b) => a.score - b.score || a.selector.localeCompare(b.selector)) .slice(0, limit) .map(({ model, selector }) => ({ provider: model.provider, id: model.id, name: model.name || model.id, selector, }));}
/** * Models whose full selector ends with the reference, so a bare model id like * "z-ai/glm-5.3" also matches "prime-inference/z-ai/glm-5.3". */function findRlmShortFormModelMatches(reference: string, models: Model<Api>[]): Model<Api>[] { const normalized = reference.trim().toLowerCase(); if (!normalized) return []; return models.filter((model) => `${model.provider}/${model.id}`.toLowerCase().endsWith(`/${normalized}`));}
/** * The single model a short-form reference resolves to: its unique match among * models, or the fallback model when no model matches. Stays undefined when * several models match, so an ambiguous reference is never auto-resolved. */export function findUniqueRlmShortFormModelMatch( reference: string, models: Model<Api>[], fallback?: Model<Api>,): Model<Api> | undefined { const matches = findRlmShortFormModelMatches(reference, models); if (matches.length === 1) return matches[0]; if (matches.length === 0 && fallback && findRlmShortFormModelMatches(reference, [fallback]).length === 1) { return fallback; } return undefined;}
/** * Rejection message for an unresolved model reference: states that the model is * unavailable, unauthenticated, or expired, then the expected selector form and * close matches so the user can retry with a full selector. */export function formatRlmModelUnavailableError(reference: string, target: string, models: Model<Api>[]): string { const base = `Requested ${target} model "${reference}" is unavailable, unauthenticated, or expired`; const hint = `selectors use the form "provider/model-id" (e.g. "prime-inference/z-ai/glm-5.3")`; const normalizedReference = normalizeModelSearchText(reference); const closeMatches = normalizedReference ? findRlmModelMatches(reference, models, RLM_MODEL_ERROR_SUGGESTION_LIMIT).map((match) => match.selector) : []; if (closeMatches.length === 0) { return `${base}; ${hint}`; } return `${base}; ${hint}; close matches: ${closeMatches.map((selector) => `"${selector}"`).join(", ")}`;}
export function createRlmCreateSessionHostHandler(handler: RlmCreateSessionHandler): HostRequestHandler { return async (payload) => { if (typeof payload.prompt !== "string") { throw new Error("rlm.create_session prompt must be a string"); } const kwargs = isRecord(payload.kwargs) ? payload.kwargs : {}; const result = await handler({ prompt: payload.prompt, kwargs }); return result as unknown as Record<string, unknown>; };}
/** Adapt an RlmRunHandler into the typed `rlm.run` kernel host handler. */export function createRlmRunHostHandler(handler: RlmRunHandler): HostRequestHandler { return async (payload) => { if (typeof payload.prompt !== "string") { throw new Error("rlm.spawn prompt must be a string"); } const kwargs = isRecord(payload.kwargs) ? payload.kwargs : {}; const cellSourceCode = typeof payload.cellSourceCode === "string" ? payload.cellSourceCode : undefined; const result = await handler({ prompt: payload.prompt, kwargs, cellSourceCode, }); return result as unknown as Record<string, unknown>; };}
/** Adapt detached kernel bash completions into a validated host notification. */export function createAsyncBashCompletionHostHandler(handler: AsyncBashCompletionHandler): HostRequestHandler { return async (payload) => { const { pid, command, exitCode } = payload; if (typeof pid !== "number" || !Number.isInteger(pid) || pid <= 0) { throw new Error("bash.completed pid must be a positive integer"); } if (typeof command !== "string" || !command) { throw new Error("bash.completed command must be a non-empty string"); } if (typeof exitCode !== "number" || !Number.isInteger(exitCode)) { throw new Error("bash.completed exitCode must be an integer"); } await handler({ pid, command, exitCode }); return {}; };}
/** The kernel read a finished command's result, so its completion notice is stale. */export function createAsyncBashConsumedHostHandler(handler: AsyncBashConsumedHandler): HostRequestHandler { return async (payload) => { const { pid, command } = payload; if (typeof pid !== "number" || !Number.isInteger(pid) || pid <= 0) { throw new Error("bash.consumed pid must be a positive integer"); } if (typeof command !== "string" || !command) { throw new Error("bash.consumed command must be a non-empty string"); } await handler({ pid, command }); return {}; };}
/** Return the categorized available models. */export function createRlmModelsHostHandler(handler: RlmModelsHandler): HostRequestHandler { return async () => { const result = await handler(); return result as unknown as Record<string, unknown>; };}
/** Expose the current parent session's direct RLM child registry to its kernel. */export function createRlmListSubagentsHostHandler(handler: RlmListSubagentsHandler): HostRequestHandler { return async () => { const { subagents } = await handler(); return { subagents }; };}
/** Delete one direct child selected from the current parent session's registry. */export function createRlmDeleteSubagentHostHandler(handler: RlmDeleteSubagentHandler): HostRequestHandler { return async (payload) => { if (typeof payload.target !== "string" || !payload.target.trim()) { throw new Error("rlm.delete_subagent target must be a non-empty string"); } const { subagent, outcome } = await handler(payload.target.trim()); return outcome === undefined ? { subagent } : { subagent, outcome }; };}
/** * Typed fan-in for subagent results: `rlm.collect` waits (bounded) for the * selected direct children's runs to settle and returns result envelopes. * Never steers the parent: a timeout returns the current snapshots instead of * rejecting, so the caller can poll, end the turn, or retry. */export function createRlmCollectHostHandler(handler: RlmCollectHandler): HostRequestHandler { return async (payload) => { const rawTargets = payload.targets; if (rawTargets !== undefined && rawTargets !== null && !Array.isArray(rawTargets)) { throw new Error("rlm.collect targets must be an array of child ids or names"); } const targets = (rawTargets ?? []).map((target) => { if (typeof target !== "string" || !target.trim()) { throw new Error("rlm.collect targets must be non-empty strings"); } return target.trim(); }); const rawTimeout = payload.timeout_ms; if (rawTimeout === undefined || rawTimeout === null) { const { results } = await handler(targets, 0); return { results }; } if ( typeof rawTimeout !== "number" || !Number.isSafeInteger(rawTimeout) || rawTimeout < 0 || rawTimeout > 2_147_483_647 ) { throw new Error("rlm.collect timeout_ms must be a non-negative integer up to 2147483647"); } const { results } = await handler(targets, rawTimeout); return { results }; };}
export interface RlmProgressNoteResult { accepted: boolean; /** Milliseconds until the next note can be accepted; absent when accepted. */ retry_after_ms: number | undefined;}
export type RlmProgressNoteHandler = (message: string) => RlmProgressNoteResult;
/** Hard bound for one progress note; the session handler owns the time throttle. */export const RLM_PROGRESS_NOTE_MAX_LENGTH = 512;
/** * Child progress notes: `rlm.progress.note` lets a child report short in-flight * status that its parent reads from snapshots and roster entries. Pull-based * only — notes never steer the parent or grow its message queue. */export function createRlmProgressNoteHostHandler(handler: RlmProgressNoteHandler): HostRequestHandler { return async (payload) => { const raw = payload.message; if (typeof raw !== "string" || !raw.trim()) { throw new Error("rlm.progress.note message must be a non-empty string"); } const message = raw.trim(); if (message.length > RLM_PROGRESS_NOTE_MAX_LENGTH) { throw new Error(`rlm.progress.note message must be at most ${RLM_PROGRESS_NOTE_MAX_LENGTH} characters`); } const { accepted, retry_after_ms } = handler(message); return retry_after_ms === undefined ? { accepted } : { accepted, retry_after_ms }; };}
export interface RlmSubagentRuntime { session: AgentSession;}
export interface CreateRlmSubagentRuntimeOptions { parentSession: AgentSession; id: string; prompt: string; sessionName: string; sessionDir: string; model: Model<any>; thinkingLevel: ThinkingLevel; serviceTier: ServiceTier; scopedModels: Array<{ model: Model<any>; thinkingLevel?: ThinkingLevel }>; activeToolNames: string[]; allowedToolNames?: string[]; customTools: ToolDefinition[]; includeGoals: boolean; includeCompactSkill: boolean; rlmDepth: number; rlmMaxDepth: number; rlmParentNodeId: string; /** Request ID of the parent model call whose tool call caused this spawn. */ spawnedByRequestId?: string; /** Source of the Python cell that spawned this subagent, for display. */ spawnCode?: string; /** Publish the session to the parent before a host makes the runtime addressable. */ onSessionPublished?: (session: AgentSession) => void; /** Remote host destination if this subagent's kernel should execute on a remote host over SSH. */ remoteHost?: string; /** Let the child publish (git push, force reset) without a further authorization. */ allowPush?: boolean; /** * Working directory the child session, its kernel, and its `bash()` calls start in. * The parent always sets this to its own session cwd, so a child cannot inherit an * unrelated directory and walk a multi-TB home. Optional only so an older host * implementation that reads the parent session directly still type-checks; a host * that receives it must use it. */ cwd?: string;}
export interface CreateRlmRootSessionOptions { prompt: string; sessionName?: string; cwd: string; model: Model<Api>; thinkingLevel: ThinkingLevel;}
export interface SubagentRuntimeHost { createRlmSubagentRuntime(options: CreateRlmSubagentRuntimeOptions): Promise<RlmSubagentRuntime>; createRlmRootSession?(options: CreateRlmRootSessionOptions): Promise<RlmCreateSessionResult>; /** Persist host-owned completion before the child becomes passivation-eligible. */ completeRlmSubagentRuntime?(childId: string, session: AgentSession): boolean; /** Release a host-owned child after its detached initial task settles. */ releaseRlmSubagentRuntime?: ( runtime: RlmSubagentRuntime, options: CreateRlmSubagentRuntimeOptions, status: "done" | "error" | "cancelled", ) => Promise<void>; /** Close or remove the host-owned child; session is absent when a persisted child is still passive. */ deleteRlmSubagentRuntime(childId: string, session?: AgentSession): Promise<void>; disposeRlmSubagentRuntimes?(): Promise<void>;}