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.
12 kB · 322 lines
TypeScript
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323import type { ThinkingLevel } from "@earendil-works/pi-agent-core";import type { Api, Model, ServiceTier } from "@earendil-works/pi-ai";import type { AgentSession } 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.run`; 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";
export interface RlmSubagentRegistryEntry { rlm_child_id: string; active_session_id: string | null; session_id: string | null; session_name: string; session_dir: string; status: RlmSubagentRegistryStatus;}
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 RlmFindModelsResult { models: RlmModelMatch[];}
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>;export type RlmListSubagentsHandler = () => RlmListSubagentsResult | Promise<RlmListSubagentsResult>;export type RlmDeleteSubagentHandler = (target: string) => Promise<RlmDeleteSubagentResult>;export type RlmFindModelsHandler = (query: string, limit: number) => RlmFindModelsResult | Promise<RlmFindModelsResult>;
const RLM_SUBAGENT_SESSION_NAME_MAX_LENGTH = 64;export const DEFAULT_RLM_MODEL_SEARCH_LIMIT = 8;export const MAX_RLM_MODEL_SEARCH_LIMIT = 20;
export function normalizeRequestedRlmSubagentSessionName(value: unknown, operation = "rlm.run"): 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.run",): 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.run"): 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, }));}
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.run 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 {}; };}
/** Search a bounded authenticated model catalog without adding it to the system prompt. */export function createRlmFindModelsHostHandler(handler: RlmFindModelsHandler): HostRequestHandler { return async (payload) => { if (typeof payload.query !== "string") { throw new Error("rlm.find_models query must be a string"); } const limit = payload.limit === undefined ? DEFAULT_RLM_MODEL_SEARCH_LIMIT : payload.limit; if (!Number.isInteger(limit) || (limit as number) < 1 || (limit as number) > MAX_RLM_MODEL_SEARCH_LIMIT) { throw new Error(`rlm.find_models limit must be an integer from 1 to ${MAX_RLM_MODEL_SEARCH_LIMIT}`); } return { models: (await handler(payload.query, limit as number)).models }; };}
/** 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 }; };}
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;}
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>;}