diff --git a/packages/agent/src/agent.ts b/packages/agent/src/agent.ts index bc75d86c54..f765601adf 100644 --- a/packages/agent/src/agent.ts +++ b/packages/agent/src/agent.ts @@ -111,6 +111,7 @@ export interface AgentOptions { thinkingBudgets?: ThinkingBudgets; transport?: Transport; maxRetryDelayMs?: number; + maxTokens?: number; toolExecution?: ToolExecutionMode; } @@ -207,6 +208,8 @@ export class Agent { public transport: Transport; /** Optional cap for provider-requested retry delays. */ public maxRetryDelayMs?: number; + /** Optional per-request output token cap forwarded to the stream function. */ + public maxTokens?: number; /** Tool execution strategy for assistant messages that contain multiple tool calls. */ public toolExecution: ToolExecutionMode; @@ -228,6 +231,7 @@ export class Agent { this.thinkingBudgets = options.thinkingBudgets; this.transport = options.transport ?? "auto"; this.maxRetryDelayMs = options.maxRetryDelayMs; + this.maxTokens = options.maxTokens; this.toolExecution = options.toolExecution ?? "parallel"; } @@ -448,6 +452,7 @@ export class Agent { transport: this.transport, thinkingBudgets: this.thinkingBudgets, maxRetryDelayMs: this.maxRetryDelayMs, + maxTokens: this.maxTokens, toolExecution: this.toolExecution, beforeToolCall: this.beforeToolCall, afterToolCall: this.afterToolCall, diff --git a/packages/coding-agent/src/core/agent-session.ts b/packages/coding-agent/src/core/agent-session.ts index 325718bcfc..99f1fa47ed 100644 --- a/packages/coding-agent/src/core/agent-session.ts +++ b/packages/coding-agent/src/core/agent-session.ts @@ -135,6 +135,9 @@ import type { ResourceExtensionPaths, ResourceLoader } from "./resource-loader.j import { type CreateRlmSubagentRuntimeOptions, createRlmRunHostHandler, + createRlmSendAdvanceHostHandler, + createRlmSendCloseHostHandler, + createRlmSendCreateHostHandler, type RlmInternalRunResult, type RlmRunResult, type RlmSubagentRuntime, @@ -158,7 +161,7 @@ import { type BashOperations, createLocalBashOperations } from "./tools/bash.js" import { createAllToolDefinitions } from "./tools/index.js"; import { IpythonKernelProvisioner } from "./tools/ipython.js"; import { createToolDefinitionFromAgentTool } from "./tools/tool-definition-wrapper.js"; -import { addAssistantUsage, cloneUsage, emptyUsage } from "./usage.js"; +import { addAssistantUsage, cloneUsage, emptyUsage, subtractAssistantUsage } from "./usage.js"; export type { GoalState, GoalStatus } from "./goals.js"; export type { SessionStats } from "./session-stats.js"; @@ -720,6 +723,11 @@ export class AgentSession { private _rlmParentNodeId?: string; private _subagentRuntimeHost?: SubagentRuntimeHost; private _activeRlmChildRuns = new Map(); + /** Named persistent sub-agents created via rlm.send (kept alive across host requests). */ + private _persistentRlmChildren = new Map< + string, + { runtime: RlmSubagentRuntime; sessionDir: string; lastAttributedUsage: Usage; lastReturnedRlmUsage: RlmUsage } + >(); // Model registry for API key resolution private _modelRegistry: ModelRegistry; @@ -1815,6 +1823,11 @@ export class AgentSession { if (this._disposed) { return; } + try { + await this._closeAllPersistentRlmChildren(); + } catch { + // best effort during teardown + } try { await this._ipythonKernelProvisioner?.dispose(); } catch { @@ -3681,6 +3694,13 @@ export class AgentSession { "rlm.run": createRlmRunHostHandler(({ prompt, kwargs, cellSourceCode }) => this.runRlmChild(prompt, kwargs, cellSourceCode), ), + "rlm.send.create": createRlmSendCreateHostHandler((request) => + this._createPersistentRlmChild(request.name, request.max_tokens), + ), + "rlm.send.advance": createRlmSendAdvanceHostHandler((request) => + this._advancePersistentRlmChild(request.name, request.prompt), + ), + "rlm.send.close": createRlmSendCloseHostHandler((request) => this._closePersistentRlmChild(request.name)), }; if (this._includeGoals) { for (const type of ["goal.get", "goal.create", "goal.complete"]) { @@ -3755,8 +3775,19 @@ export class AgentSession { return undefined; } - private _createChildRlmSessionDir(): string { + private _createChildRlmSessionDir(name?: string): string { const parentDir = this._ensureRlmSessionDir() ?? this._createEphemeralRlmSessionDir(); + if (name) { + // Named persistent agent: stable, human-readable directory + const safeName = + name + .replace(/[^A-Za-z0-9._-]/g, "-") + .replace(/^-+|-+$/g, "") + .slice(0, 64) || "agent"; + const childDir = join(parentDir, `sub-${safeName}`); + mkdirSync(childDir, { recursive: true }); + return childDir; + } for (let i = 0; i < 100; i++) { const childDir = join(parentDir, `sub-${randomUUID().slice(0, 8)}`); try { @@ -3827,6 +3858,7 @@ export class AgentSession { spawnCode?: string; sessionDir: string; model: Model; + maxTokens?: number; }): CreateRlmSubagentRuntimeOptions { return { parentSession: this, @@ -3844,6 +3876,7 @@ export class AgentSession { rlmDepth: this._rlmDepth + 1, rlmMaxDepth: this._rlmMaxDepth, rlmParentNodeId: options.id, + maxTokens: options.maxTokens, }; } @@ -3894,6 +3927,7 @@ export class AgentSession { thinkingBudgets: this.settingsManager.getThinkingBudgets(), transport: this.settingsManager.getTransport(), maxRetryDelayMs: this.settingsManager.getProviderRetrySettings().maxRetryDelayMs, + maxTokens: options.maxTokens, toolExecution: this.agent.toolExecution, }); @@ -4277,6 +4311,143 @@ export class AgentSession { }; } + // ========================================================================= + // Persistent / Background Sub-Agents (rlm.send) + // ========================================================================= + + /** + * Create a persistent sub-agent session that survives across host requests. + * Called by the rlm.send.create host handler. + */ + private async _createPersistentRlmChild(name: string, maxTokens?: number): Promise<{ session_dir: string | null }> { + // If already exists, just return its session dir + const existing = this._persistentRlmChildren.get(name); + if (existing) { + return { session_dir: existing.sessionDir }; + } + + if (this._rlmDepth >= this._rlmMaxDepth) { + throw new Error( + `RLM recursion depth limit reached (RLM_DEPTH=${this._rlmDepth}, RLM_MAX_DEPTH=${this._rlmMaxDepth})`, + ); + } + + const model = this.model; + if (!model) { + throw new Error("No model selected"); + } + + const childSessionDir = this._createChildRlmSessionDir(name); + const childNodeId = `persistent-${name}`; + + const subagentOptions = this._createRlmSubagentRuntimeOptions({ + id: childNodeId, + prompt: `[persistent agent: ${name}]`, + sessionDir: childSessionDir, + model, + maxTokens, + }); + + const childRuntime = await this._createRlmSubagentRuntime(subagentOptions); + this._persistentRlmChildren.set(name, { + runtime: childRuntime, + sessionDir: childSessionDir, + lastAttributedUsage: emptyUsage(), + lastReturnedRlmUsage: emptyRlmUsage(), + }); + + return { session_dir: childSessionDir }; + } + + /** + * Send a prompt to an existing persistent sub-agent and await its response. + * Called by the rlm.send.advance host handler. + */ + private async _advancePersistentRlmChild(name: string, prompt: string): Promise { + const child = this._persistentRlmChildren.get(name); + if (!child) { + throw new Error(`No persistent sub-agent named '${name}'; create it with rlm.send.create first`); + } + + const session = child.runtime.session; + const parentAssistant = this._findLastAssistantMessage(); + + await session.prompt(prompt, { expandPromptTemplates: false, source: "extension" }); + await session.agent.waitForIdle(); + + // Guard: the child may have been closed while we were awaiting above. + if (!this._persistentRlmChildren.has(name)) { + throw new Error(`Persistent sub-agent '${name}' was closed during advance`); + } + + const answer = session.getLastAssistantText() ?? ""; + const cumulativeRlmUsage = session._usageForCurrentMessages(); + const cumulativeAssistantUsage = session._assistantUsageForCurrentMessages(); + + // Compute deltas: only attribute/return usage from THIS advance, not prior ones. + const usageDelta = emptyUsage(); + addAssistantUsage(usageDelta, cumulativeAssistantUsage); + subtractAssistantUsage(usageDelta, child.lastAttributedUsage); + child.lastAttributedUsage = cloneUsage(cumulativeAssistantUsage); + + this._attributeRlmChildUsageToParent(usageDelta, parentAssistant); + + const rlmUsageDelta: RlmUsage = { + prompt_tokens: cumulativeRlmUsage.prompt_tokens - child.lastReturnedRlmUsage.prompt_tokens, + completion_tokens: cumulativeRlmUsage.completion_tokens - child.lastReturnedRlmUsage.completion_tokens, + }; + child.lastReturnedRlmUsage = { ...cumulativeRlmUsage }; + + return { + answer, + usage: rlmUsageDelta, + turns: session._assistantTurnCount(), + session_dir: child.sessionDir, + }; + } + + /** + * Close and dispose a persistent sub-agent. + * Called by the rlm.send.close host handler. + */ + private async _closePersistentRlmChild(name: string): Promise { + const child = this._persistentRlmChildren.get(name); + if (!child) { + return; // Already closed or never existed; idempotent + } + this._persistentRlmChildren.delete(name); + // Abort any in-flight advance so its prompt/waitForIdle settles before we + // dispose the runtime. Without this, disposeAsync can tear down the child + // session while _advancePersistentRlmChild is still mid-await. + try { + await child.runtime.session.abort(); + } catch { + // best effort — the session may already be idle + } + const childNodeId = `persistent-${name}`; + const subagentOptions = this._createRlmSubagentRuntimeOptions({ + id: childNodeId, + prompt: `[persistent agent: ${name}]`, + sessionDir: child.sessionDir, + model: this.model!, + }); + await this._releaseRlmSubagentRuntime(child.runtime, subagentOptions); + } + + /** + * Close all persistent sub-agents (called during session teardown). + */ + private async _closeAllPersistentRlmChildren(): Promise { + const names = [...this._persistentRlmChildren.keys()]; + for (const name of names) { + try { + await this._closePersistentRlmChild(name); + } catch { + // Best effort during teardown + } + } + } + // ========================================================================= // Auto-Retry // ========================================================================= diff --git a/packages/coding-agent/src/core/kernel/bootstrap.ts b/packages/coding-agent/src/core/kernel/bootstrap.ts index 43d1356bbb..ac902d940d 100644 --- a/packages/coding-agent/src/core/kernel/bootstrap.ts +++ b/packages/coding-agent/src/core/kernel/bootstrap.ts @@ -11,7 +11,7 @@ import { fileURLToPath } from "node:url"; import { getPackageDir } from "../../config.js"; import type { PythonSkillRuntimeInfo } from "../skills.js"; -const BOOTSTRAP_SCHEMA = 7; +const BOOTSTRAP_SCHEMA = 8; const PYTHON_VERSION = "3.11"; const IPYKERNEL_REQUIREMENT = "ipykernel"; const RUNTIME_REQUIREMENT = "prime-agent-runtime"; @@ -51,7 +51,7 @@ const REQUIRED_HARNESS_METHODS = [ "delete_prompt_note", "record_refinement", ]; -const RUNTIME_READY_CHECK = `import inspect; import rlm; from rlm.harness import HarnessEntry; _harness_methods = ${JSON.stringify(REQUIRED_HARNESS_METHODS)}; assert hasattr(rlm, 'run'); assert callable(rlm); assert hasattr(rlm, 'rlm'); assert callable(rlm.rlm); assert callable(rlm.host_request); assert hasattr(rlm, 'harness'); assert hasattr(rlm, 'get_harness_state'); assert hasattr(rlm.rlm, 'harness'); assert hasattr(rlm.rlm, 'get_harness_state'); assert all(callable(getattr(_harness, _method, None)) for _harness in (rlm.harness, rlm.rlm.harness) for _method in _harness_methods); assert 'reference' in HarnessEntry.__dataclass_fields__; assert 'reference' in inspect.signature(rlm.harness.create_skill).parameters; assert 'reference' in inspect.signature(rlm.harness.update_skill).parameters; assert not hasattr(rlm, 'background'); assert not hasattr(rlm.rlm, 'background')`; +const RUNTIME_READY_CHECK = `import inspect; import rlm; from rlm.harness import HarnessEntry; _harness_methods = ${JSON.stringify(REQUIRED_HARNESS_METHODS)}; assert hasattr(rlm, 'run'); assert callable(rlm); assert hasattr(rlm, 'rlm'); assert callable(rlm.rlm); assert callable(rlm.host_request); assert hasattr(rlm, 'send'); assert hasattr(rlm.rlm, 'send'); assert hasattr(rlm, 'harness'); assert hasattr(rlm, 'get_harness_state'); assert hasattr(rlm.rlm, 'harness'); assert hasattr(rlm.rlm, 'get_harness_state'); assert all(callable(getattr(_harness, _method, None)) for _harness in (rlm.harness, rlm.rlm.harness) for _method in _harness_methods); assert 'reference' in HarnessEntry.__dataclass_fields__; assert 'reference' in inspect.signature(rlm.harness.create_skill).parameters; assert 'reference' in inspect.signature(rlm.harness.update_skill).parameters; assert not hasattr(rlm, 'background'); assert not hasattr(rlm.rlm, 'background')`; const BOOTSTRAP_VERSION_FILE = ".bootstrap-version"; const BOOTSTRAP_LOCK_NAME = ".bootstrap.lock"; const BOOTSTRAP_LOCK_RETRY_MS = 100; diff --git a/packages/coding-agent/src/core/prompts/rlm.ts b/packages/coding-agent/src/core/prompts/rlm.ts index 5be6a41bf3..20c28ce3e0 100644 --- a/packages/coding-agent/src/core/prompts/rlm.ts +++ b/packages/coding-agent/src/core/prompts/rlm.ts @@ -62,6 +62,9 @@ export function buildRlmPrompt(options: RlmPromptOptions): string { skillLines.push( "Each Python skill may also be available as a shell command by the same name: ` ...`. Discover its CLI usage with ` --help`.", ); + skillLines.push( + "To offload a slow skill call, run it in the background: `handle = .send(...)`, then check `handle.poll()` later (same handle API as sub-agents).", + ); if (installedSkills.includes("edit")) { skillLines.push( "For targeted existing-file edits, prefer the pre-imported async `edit` skill from IPython: `old = '''...'''; new = '''...'''; await edit(path=\"pkg/file.py\", old_str=old, new_str=new)`. Use exact old/new strings; if the text contains triple double quotes, use triple single-quoted variables or build `old`/`new` from inspected file slices.", @@ -75,9 +78,12 @@ export function buildRlmPrompt(options: RlmPromptOptions): string { if (allowRecursion) { parts.push( "", - "A callable `rlm` is already in your global namespace — call it directly with `await rlm('sub-task')` to spawn a recursive sub-agent. Returns an `RLMResult` with `.answer` (string), `.usage`, `.turns`, and `.session_dir`.", - "For parallel sub-agents, use normal Python async patterns such as `await asyncio.gather(rlm('task1'), rlm('task2'))`.", - "For sub-agent work that can run in the background, keep the task handle from `asyncio.create_task(rlm('sub-task'))` so you do not block the main execution path; use normal task callbacks, `task.done()`, or `await task` later to observe completion and read the returned `RLMResult.answer`.", + "A callable `rlm` is already in your global namespace for spawning recursive sub-agents.", + "- One-off (blocking): `await rlm('sub-task')` runs a sub-agent to completion and returns an `RLMResult` (`.answer`, `.usage`, `.turns`, `.session_dir`). Run several at once with `await asyncio.gather(rlm('a'), rlm('b'))`.", + "- Persistent (background): `handle = rlm.send('sub-task', name='helper')` (omit `name` for an auto-generated one) starts a named sub-agent and returns a handle immediately — start it in one tool call and check on it in a later one. Re-`send` the same `name` to continue that agent's conversation, so you can keep a specialist around and ask it repeatedly.", + "- `handle.poll()` reports its state without blocking: `.status` is 'running', 'finished', or 'error'; `.results` is a FIFO of finished `RLMResult`s (take the next with `.results.popleft()`); `.queued` is the editable list of not-yet-started prompts; `.error` is the exception if it failed. Use `await handle.wait()` to block for the next result.", + "- When multiple sub-agents are in flight, prefer polling all handles in a single tool call over sequential `await handle.wait()` calls — this lets you check progress and act on whichever results are ready without blocking on one agent at a time.", + "- Keep the handle in a variable to poll it from a later tool call (your IPython namespace persists across calls). Inspect the full API with `help(rlm.send)`.", ); } diff --git a/packages/coding-agent/src/core/rlm-runtime.ts b/packages/coding-agent/src/core/rlm-runtime.ts index 9cba6226b5..8ee3bc5f71 100644 --- a/packages/coding-agent/src/core/rlm-runtime.ts +++ b/packages/coding-agent/src/core/rlm-runtime.ts @@ -29,10 +29,75 @@ export interface RlmInternalRunResult extends RlmRunResult { export type RlmRunHandler = (request: RlmRunRequest) => Promise; +// --- Persistent / background sub-agent requests --- + +export interface RlmSendCreateRequest { + name: string; + max_tokens?: number; + [key: string]: unknown; +} + +export interface RlmSendCreateResult { + session_dir: string | null; +} + +export interface RlmSendAdvanceRequest { + name: string; + prompt: string; +} + +export interface RlmSendCloseRequest { + name: string; +} + +export type RlmSendCreateHandler = (request: RlmSendCreateRequest) => Promise; +export type RlmSendAdvanceHandler = (request: RlmSendAdvanceRequest) => Promise; +export type RlmSendCloseHandler = (request: RlmSendCloseRequest) => Promise; + function isRecord(value: unknown): value is Record { return typeof value === "object" && value !== null && !Array.isArray(value); } +/** Adapt an RlmSendCreateHandler into the typed "rlm.send.create" handler for the kernel host bridge. */ +export function createRlmSendCreateHostHandler(handler: RlmSendCreateHandler): HostRequestHandler { + return async (payload) => { + if (typeof payload.name !== "string") { + throw new Error("rlm.send.create name must be a string"); + } + const request: RlmSendCreateRequest = { + name: payload.name, + max_tokens: typeof payload.max_tokens === "number" ? payload.max_tokens : undefined, + }; + const result = await handler(request); + return result as unknown as Record; + }; +} + +/** Adapt an RlmSendAdvanceHandler into the typed "rlm.send.advance" handler for the kernel host bridge. */ +export function createRlmSendAdvanceHostHandler(handler: RlmSendAdvanceHandler): HostRequestHandler { + return async (payload) => { + if (typeof payload.name !== "string") { + throw new Error("rlm.send.advance name must be a string"); + } + if (typeof payload.prompt !== "string") { + throw new Error("rlm.send.advance prompt must be a string"); + } + const result = await handler({ name: payload.name, prompt: payload.prompt }); + return result as unknown as Record; + }; +} + +/** Adapt an RlmSendCloseHandler into the typed "rlm.send.close" handler for the kernel host bridge. */ +export function createRlmSendCloseHostHandler(handler: RlmSendCloseHandler): HostRequestHandler { + return async (payload) => { + if (typeof payload.name !== "string") { + throw new Error("rlm.send.close name must be a string"); + } + await handler({ name: payload.name }); + return {}; + }; +} + /** Adapt an RlmRunHandler into the typed "rlm.run" handler for the kernel host bridge. */ export function createRlmRunHostHandler(handler: RlmRunHandler): HostRequestHandler { return async (payload) => { @@ -67,6 +132,8 @@ export interface CreateRlmSubagentRuntimeOptions { rlmParentNodeId: string; /** Source of the IPython cell that spawned this subagent, for display. */ spawnCode?: string; + /** Optional per-request output token cap for the child agent. */ + maxTokens?: number; } export interface SubagentRuntimeHost { diff --git a/packages/coding-agent/src/core/tools/ipython.ts b/packages/coding-agent/src/core/tools/ipython.ts index 65bf7c5e48..3c527d2a0f 100644 --- a/packages/coding-agent/src/core/tools/ipython.ts +++ b/packages/coding-agent/src/core/tools/ipython.ts @@ -93,6 +93,12 @@ def _prime_agent_wrap_skill_module(module): if doc: wrapped.__doc__ = doc _prime_agent_sys.modules[module.__name__] = wrapped + # Add .send for background skill execution + try: + from rlm.async_runtime import attach_background as _prime_agent_attach_bg + _prime_agent_attach_bg(wrapped, wrapped.run) + except Exception: + pass return wrapped _PRIME_AGENT_SKILL_IMPORT_ERRORS = {} diff --git a/packages/coding-agent/test/agent-session-recursion.test.ts b/packages/coding-agent/test/agent-session-recursion.test.ts index ed02b4ce04..6549ccd41c 100644 --- a/packages/coding-agent/test/agent-session-recursion.test.ts +++ b/packages/coding-agent/test/agent-session-recursion.test.ts @@ -112,6 +112,22 @@ interface InspectableRlmSession { _activeRlmChildRuns: Map; } +interface PersistentRlmSession { + sessionManager: { getSessionArtifactDir(): string | undefined }; + _createPersistentRlmChild(name: string, maxTokens?: number): Promise<{ session_dir: string | null }>; + _advancePersistentRlmChild( + name: string, + prompt: string, + ): Promise<{ + answer: string; + usage: { prompt_tokens: number; completion_tokens: number }; + turns: number; + session_dir: string | null; + }>; + _closePersistentRlmChild(name: string): Promise; + _persistentRlmChildren: Map; +} + interface KernelPumpTestApi { iopub: AsyncIterable & { close(): void }; startIopubPump(): void; @@ -835,6 +851,75 @@ print(_result.answer) stderrSpy.mockRestore(); } }); + + // Persistent / background sub-agents (rlm.send): a named child that survives + // across host requests and continues across turns, vs. the one-off rlm.run above. + it("keeps a named persistent sub-agent across turns and reuses its session", async () => { + const root = createSession() as unknown as PersistentRlmSession; + + const created = await root._createPersistentRlmChild("helper"); + expect(created.session_dir).not.toBeNull(); + expect(basename(created.session_dir!)).toBe("sub-helper"); + expect(dirname(created.session_dir!)).toBe(root.sessionManager.getSessionArtifactDir()); + expect(root._persistentRlmChildren.size).toBe(1); + + const first = await root._advancePersistentRlmChild("helper", "first question"); + expect(first.answer).toBe("child answer: first question"); + expect(first.turns).toBe(1); + expect(first.session_dir).toBe(created.session_dir); + + // Re-sending the same name continues the SAME agent: a second assistant turn + // accrues on one session instead of spinning up a fresh child. + const second = await root._advancePersistentRlmChild("helper", "second question"); + expect(second.answer).toBe("child answer: second question"); + expect(second.turns).toBe(2); + expect(second.session_dir).toBe(created.session_dir); + expect(root._persistentRlmChildren.size).toBe(1); + + await root._closePersistentRlmChild("helper"); + expect(root._persistentRlmChildren.size).toBe(0); + // Closing is idempotent. + await expect(root._closePersistentRlmChild("helper")).resolves.toBeUndefined(); + }); + + it("does not double-count usage across persistent sub-agent advances", async () => { + const root = createSession() as unknown as PersistentRlmSession; + + await root._createPersistentRlmChild("helper"); + + // First advance: child produces one assistant message (default usage: input=7, output=3). + const first = await root._advancePersistentRlmChild("helper", "first question"); + expect(first.usage).toEqual({ prompt_tokens: 7, completion_tokens: 3 }); + + // Second advance: child produces another assistant message with the same per-message usage. + // The returned usage must reflect only this advance, not the cumulative total. + const second = await root._advancePersistentRlmChild("helper", "second question"); + expect(second.usage).toEqual({ prompt_tokens: 7, completion_tokens: 3 }); + + // Third advance: same — only the delta from this advance. + const third = await root._advancePersistentRlmChild("helper", "third question"); + expect(third.usage).toEqual({ prompt_tokens: 7, completion_tokens: 3 }); + }); + + it("reuses the existing session when a persistent name is created twice", async () => { + const root = createSession() as unknown as PersistentRlmSession; + const created = await root._createPersistentRlmChild("helper"); + const again = await root._createPersistentRlmChild("helper"); + expect(again.session_dir).toBe(created.session_dir); + expect(root._persistentRlmChildren.size).toBe(1); + }); + + it("rejects advancing a persistent sub-agent that was never created", async () => { + const root = createSession() as unknown as PersistentRlmSession; + await expect(root._advancePersistentRlmChild("ghost", "hello")).rejects.toThrow( + "No persistent sub-agent named 'ghost'", + ); + }); + + it("rejects creating a persistent sub-agent at the recursion depth cap", async () => { + const root = createSession({ depth: 1, maxDepth: 1 }) as unknown as PersistentRlmSession; + await expect(root._createPersistentRlmChild("helper")).rejects.toThrow("RLM recursion depth limit reached"); + }); }); interface InspectableRlmDirSession { diff --git a/packages/coding-agent/test/system-prompt.test.ts b/packages/coding-agent/test/system-prompt.test.ts index 7b81d1deef..be163bed13 100644 --- a/packages/coding-agent/test/system-prompt.test.ts +++ b/packages/coding-agent/test/system-prompt.test.ts @@ -58,6 +58,7 @@ describe("buildRlmPrompt", () => { "When available, each Python skill is an async callable by the same import name. Inspect with `help()` or `inspect.signature(.run)`.", "If a Python skill is unavailable, calling it raises a RuntimeError with the import error.", "Each Python skill may also be available as a shell command by the same name: ` ...`. Discover its CLI usage with ` --help`.", + "To offload a slow skill call, run it in the background: `handle = .send(...)`, then check `handle.poll()` later (same handle API as sub-agents).", "", "IPython is the agent's long-lived notebook: a persistent control environment for reasoning, context management, state, tool orchestration, and recursive subcalls. Use it to keep intermediate variables, inspect and transform outputs, write small helper functions, and preserve useful state across turns or compaction.", "", @@ -321,14 +322,12 @@ describe("buildSystemPrompt", () => { expect(prompt).toContain("Conversation log: /repo/.pi/sessions/session.jsonl"); expect(prompt).toContain("await rlm('sub-task')"); expect(prompt).toContain("asyncio.gather"); - expect(prompt).toContain("asyncio.create_task"); - expect(prompt).toContain("sub-agent work that can run in the background"); - expect(prompt).toContain("do not block the main execution path"); - expect(prompt).toContain("keep the task handle"); - expect(prompt).toContain("normal task callbacks"); - expect(prompt).toContain("task.done()"); - expect(prompt).toContain("await task"); - expect(prompt).toContain("RLMResult.answer"); + expect(prompt).toContain("rlm.send"); + expect(prompt).toContain("handle.poll()"); + expect(prompt).toContain("handle.wait()"); + expect(prompt).toContain("Persistent (background)"); + expect(prompt).toContain("One-off (blocking)"); + expect(prompt).toContain("RLMResult"); expect(prompt).not.toContain("simple named task dictionary"); expect(prompt).not.toContain("rlm_tasks"); expect(prompt).not.toContain("globals().setdefault"); diff --git a/prime-agent-runtime/pyproject.toml b/prime-agent-runtime/pyproject.toml index 2cd0c0ccca..48ebd51b1d 100644 --- a/prime-agent-runtime/pyproject.toml +++ b/prime-agent-runtime/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "prime-agent-runtime" -version = "0.1.0" +version = "0.2.0" description = "Kernel-side runtime shim for Prime Agent recursion" requires-python = ">=3.10" dependencies = [ diff --git a/prime-agent-runtime/src/rlm/__init__.py b/prime-agent-runtime/src/rlm/__init__.py index 43236a4d3b..5e02986857 100644 --- a/prime-agent-runtime/src/rlm/__init__.py +++ b/prime-agent-runtime/src/rlm/__init__.py @@ -4,12 +4,22 @@ import asyncio import os +import re import sys import types from dataclasses import dataclass, field from pathlib import Path from typing import Any +from .async_runtime import ( + BackgroundWorker, + FnProcessor, + Handle, + Registry, + ToolState, + attach_background, + close_all_registries, +) from .harness import HarnessEntry, HarnessState, RefinementEvent, get_harness_state try: @@ -164,12 +174,112 @@ async def run(prompt: str, **kwargs: Any) -> RLMResult: return _result_from_payload(payload) +def sanitize_name(name: str) -> str: + """Make ``name`` filesystem-safe for a session dir; non-empty and bounded.""" + safe = re.sub(r"[^A-Za-z0-9._-]", "-", name).strip("-") + return (safe or "agent")[:64] + + +# --------------------------------------------------------------------------- +# Background / persistent sub-agents via host bridge +# --------------------------------------------------------------------------- + +# Per-kernel registry for named sub-agents. Imported fresh per kernel, so +# naturally per-kernel (hierarchical). +REGISTRY = Registry() + + +class _HostRlmProcessor: + """Stateful processor: each item is a prompt sent to the TS host. + + The host manages the actual AgentSession lifecycle. This processor + creates the session lazily on first use, then sends ``rlm.send.advance`` + requests for each queued prompt and returns the ``RLMResult``. + """ + + def __init__(self, agent_name: str, max_tokens: int | None = None, **kwargs: Any): + self._agent_name = agent_name + self._max_tokens = max_tokens + self._kwargs = kwargs + self._created = False + self.session_dir: str | None = None + + async def _ensure_created(self) -> None: + if self._created: + return + creation_result = await host_request("rlm.send.create", { + "name": self._agent_name, + "max_tokens": self._max_tokens, + **self._kwargs, + }) + self.session_dir = creation_result.get("session_dir") + # Update the worker's session_dir so Handle.session_dir works + if hasattr(self, "_worker") and self._worker is not None and self.session_dir: + from pathlib import Path + self._worker.session_dir = Path(self.session_dir) + self._created = True + + async def process(self, prompt: str) -> RLMResult: + await self._ensure_created() + payload = await host_request("rlm.send.advance", { + "name": self._agent_name, + "prompt": prompt, + }) + return _result_from_payload(payload) + + async def teardown(self) -> None: + if not self._created: + return + try: + await host_request("rlm.send.close", { + "name": self._agent_name, + }) + except Exception: + pass + + +def send( + prompt: str, + name: str | None = None, + max_tokens: int | None = None, + **kwargs: Any, +) -> Handle: + """Start or continue a named, persistent background sub-agent. + + Returns a handle immediately; keep it in a variable and ``handle.poll()`` it + from a later cell. Re-sending the same ``name`` appends a turn to the same + agent (multi-turn). ``name=None`` draws a random auto-name. + ``max_tokens`` caps the sub-agent's completion-token budget. + """ + _ensure_recursion_allowed() + + if name is not None: + name = sanitize_name(name) + + # Non-positive request means "no explicit budget" + if max_tokens is not None and max_tokens <= 0: + max_tokens = None + + def worker_factory(agent_name: str) -> BackgroundWorker: + processor = _HostRlmProcessor(agent_name, max_tokens=max_tokens, **kwargs) + worker = BackgroundWorker(agent_name, processor) + processor._worker = worker # so processor can update session_dir after creation + return worker + + return REGISTRY.send(prompt, name=name, worker_factory=worker_factory) + + +async def drain_agents() -> None: + """Close every background agent in this kernel (all registries). + + Invoked by the engine's teardown cascade as a cell executed in the kernel. + """ + await close_all_registries() + + try: _harness_state = get_harness_state() except Exception: # pragma: no cover - harness state must never break `import rlm` - # Importing rlm runs inside the kernel; a failure here would take down the whole - # kernel. Fall back to a true in-memory store (no path resolution, no disk) so the - # failure cannot recur and refinement is merely degraded, not fatal. _harness_state = HarnessState(in_memory=True) @@ -183,6 +293,21 @@ async def run(self, prompt: str, **kwargs: Any) -> RLMResult: async def __call__(self, prompt: str, **kwargs: Any) -> RLMResult: return await run(prompt, **kwargs) + @staticmethod + def send( + prompt: str, + name: str | None = None, + max_tokens: int | None = None, + **kwargs: Any, + ) -> Handle: + """Start or continue a named, persistent background sub-agent. + + Returns a handle immediately; keep it in a variable and ``handle.poll()`` it + from a later cell. Re-sending the same ``name`` appends a turn to the same + agent (multi-turn). ``name=None`` draws a random auto-name. + """ + return send(prompt, name=name, max_tokens=max_tokens, **kwargs) + rlm = _RLMCallable() harness = _harness_state @@ -196,14 +321,25 @@ async def __call__(self, prompt: str, **kwargs: Any) -> RLMResult: sys.modules[__name__].__class__ = _CallableModule __all__ = [ + "BackgroundWorker", + "FnProcessor", + "Handle", "HarnessEntry", "HarnessState", "RLMResult", + "REGISTRY", + "Registry", "RefinementEvent", "TokenUsage", + "ToolState", + "attach_background", + "close_all_registries", + "drain_agents", "get_harness_state", "harness", "host_request", "rlm", "run", + "sanitize_name", + "send", ] diff --git a/prime-agent-runtime/src/rlm/async_runtime.py b/prime-agent-runtime/src/rlm/async_runtime.py new file mode 100644 index 0000000000..9ca94296c7 --- /dev/null +++ b/prime-agent-runtime/src/rlm/async_runtime.py @@ -0,0 +1,317 @@ +"""Background workers for programmatic tool calls inside the IPython kernel. + +``send(...)`` schedules a tool invocation as a background asyncio task on the +kernel's event loop and returns a :class:`Handle` the model can ``poll()`` across +cells. The kernel loop keeps these tasks progressing between cells (see +``docs/background-subagents-and-tools.md``). + +A worker owns an inbox and drains it sequentially via a :class:`Processor`: + +- the default stateless processor runs each item as an independent call; +- ``rlm`` supplies a stateful processor (a live agent) so re-sending the same + name continues a multi-turn conversation. +""" + +from __future__ import annotations + +import asyncio +import inspect +import uuid +import weakref +from collections import deque +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Awaitable, Callable, Protocol, runtime_checkable + +RUNNING = "running" +FINISHED = "finished" # idle: inbox empty, ready for more +ERROR = "error" + +NO_ITEM = object() # omitted from BackgroundWorker(item=...) -> resident worker + + +@dataclass +class ToolState: + """A ``poll()`` snapshot. + + ``results`` and ``queued`` are the worker's **live** structures, not copies: + the model drains ``results`` (``.popleft()``) and edits ``queued`` directly + (cancel / reorder / edit pending items). ``poll()`` itself never consumes. + Editing is race-free because the kernel loop is single-threaded cooperative — + a sync edit in a cell can't interleave with the drain coroutine. + """ + + status: str + results: deque + queued: list + error: BaseException | None = None + + def __repr__(self) -> str: + # Bounded: a bare ``handle.poll()`` echoes this repr into model context. + parts = [f"status={self.status!r}", f"results={len(self.results)}"] + if self.queued: + parts.append(f"queued={len(self.queued)}") + if self.error is not None: + parts.append(f"error={type(self.error).__name__}") + return f"ToolState({', '.join(parts)})" + + +@runtime_checkable +class Processor(Protocol): + """Turns one queued input into a result, plus one-shot teardown. + + ``process`` is called once per queued item, in order; ``teardown`` runs once + when the worker is closed. + """ + + async def process(self, item: Any) -> Any: ... + + async def teardown(self) -> None: ... + + +class FnProcessor: + """Default stateless processor: each item is an independent call to ``fn``. + + Items are ``(args, kwargs)`` tuples produced by :meth:`Registry.send`. + """ + + def __init__(self, fn: Callable[..., Awaitable[Any]]): + self._fn = fn + + async def process(self, item: Any) -> Any: + args, kwargs = item + return await self._fn(*args, **kwargs) + + async def teardown(self) -> None: + return None + + +class BackgroundWorker: + """A background task that runs a processor on the kernel loop. + + Two lifecycles, selected by the ``item`` constructor arg: a *resident* worker + drains an inbox sequentially and parks on ``_wake`` when idle (named rlm + agents); an *ephemeral* worker runs its single item once and ends (general + tools). A processor exception halts the worker in ``ERROR`` state, surfaced + to the model via ``poll().error``. + """ + + def __init__( + self, + name: str, + processor: Processor, + *, + session_dir: Path | None = None, + item: Any = NO_ITEM, + ): + self.name = name + self.session_dir = session_dir + self.results: deque = deque() + self._processor = processor + self._error: BaseException | None = None + self._wake = asyncio.Event() + self._progress = asyncio.Event() + self._closing = False + # A resident worker (named rlm agent) starts empty and parks for more + # sends until close(); an ephemeral worker (general tool) is created with + # its one item and ends as soon as the inbox drains — it lives only as + # long as its Handle, never registered, so it's GC'd when dropped. + self._ephemeral = item is not NO_ITEM + self.queued: list = [item] if self._ephemeral else [] + self._status = RUNNING if self._ephemeral else FINISHED + self._task: asyncio.Future = asyncio.ensure_future(self._drain()) + + @property + def status(self) -> str: + return self._status + + def state(self) -> ToolState: + return ToolState( + status=self._status, + results=self.results, + queued=self.queued, + error=self._error, + ) + + def submit(self, item: Any) -> None: + if self._closing or self._task.done(): + raise RuntimeError(f"worker {self.name!r} is closed; cannot submit") + self.queued.append(item) + self._status = RUNNING + self._wake.set() + + async def _drain(self) -> None: + while not self._closing: + if not self.queued: + self._status = FINISHED + if self._ephemeral: + return # one-shot: its item is done, let the task end + self._wake.clear() + if self.queued or self._closing: # re-check after clear + continue + await self._wake.wait() + continue + item = self.queued.pop(0) + self._status = RUNNING + try: + result = await self._processor.process(item) + except asyncio.CancelledError: + raise + except Exception as exc: # surfaced to the model via poll(); halts worker + self._error = exc + self._status = ERROR + self._progress.set() + return + self.results.append(result) + self._progress.set() + + async def wait(self) -> Any: + """Await and return the next result (consumes it); raise if the worker errors.""" + while True: + if self.results: + return self.results.popleft() + if self._status == ERROR: + assert self._error is not None + raise self._error + if self._task.done(): + raise RuntimeError(f"worker {self.name!r} ended without a result") + self._progress.clear() + if self.results or self._status == ERROR: # re-check after clear + continue + await self._progress.wait() + + async def close(self) -> None: + """Cancel the drain task and run processor teardown (idempotent). + + Settles on a terminal status so a held handle never reports a stale + ``RUNNING`` after teardown: ``ERROR`` is preserved (``poll().error`` stays + meaningful), otherwise the worker reports ``FINISHED``. + """ + self._closing = True + self._wake.set() + if not self._task.done(): + self._task.cancel() + try: + await self._task + except (asyncio.CancelledError, Exception): + pass + await self._processor.teardown() + if self._status != ERROR: + self._status = FINISHED + + +class Handle: + """Model-facing reference to a background worker.""" + + def __init__(self, worker: BackgroundWorker): + self._worker = worker + + @property + def name(self) -> str: + return self._worker.name + + @property + def session_dir(self) -> Path | None: + """Session dir of the worker, if any. ``session_dir/'messages.jsonl'`` is + the live transcript (rlm); ``None`` for tools without one.""" + return self._worker.session_dir + + def poll(self) -> ToolState: + return self._worker.state() + + async def wait(self) -> Any: + return await self._worker.wait() + + def __repr__(self) -> str: + return f"Handle(name={self._worker.name!r}, status={self._worker.status!r})" + + +ALL_REGISTRIES: "weakref.WeakSet[Registry]" = weakref.WeakSet() + + +async def close_all_registries() -> None: + """Gracefully close every live registry's workers (used on kernel teardown).""" + for registry in list(ALL_REGISTRIES): + await registry.close_all() + + +class Registry: + """Per-tool, per-kernel collection of named workers.""" + + def __init__(self): + self._workers: dict[str, BackgroundWorker] = {} + ALL_REGISTRIES.add(self) + + def get(self, name: str) -> Handle | None: + worker = self._workers.get(name) + return Handle(worker) if worker is not None else None + + def send( + self, + item: Any, + *, + name: str | None, + worker_factory: Callable[[str], BackgroundWorker], + ) -> Handle: + """Enqueue ``item`` to a worker named ``name`` (creating it if needed). + + ``worker_factory(name)`` is called only when no live worker exists for + ``name``: a continuation reuses the running worker, while a name whose + previous worker halted with an error is rebuilt fresh. Evicting the dead + worker keeps the name reusable and the registry from accumulating dead + workers. + """ + if name is None: + name = uuid.uuid4().hex + worker = self._workers.get(name) + if worker is not None and worker.status == ERROR: + # The previous worker halted and already released its resources; + # evict it and fall through to rebuild, which restarts / resumes the + # agent under the same name (B13). + del self._workers[name] + worker = None + if worker is None: + worker = worker_factory(name) + self._workers[name] = worker + worker.submit(item) + return Handle(worker) + + async def close_all(self) -> None: + """Gracefully close every worker (used on kernel/agent teardown).""" + workers = list(self._workers.values()) + self._workers.clear() + for worker in workers: + try: + await worker.close() + except Exception: + pass + + +def attach_background(module, run_callable): + """Give a wrapped callable module a stateless ``.send(*a, **kw) -> Handle``. + + Each ``send`` runs ``run_callable(*a, **kw)`` on its own ephemeral worker and + returns a handle. The worker runs the call once and ends; it is not + registered, so it and its result live only as long as the model holds the + handle. A tool that wants state across calls keeps its own cache. + """ + + def send(*args, **kwargs): + worker = BackgroundWorker( + uuid.uuid4().hex, FnProcessor(run_callable), item=(args, kwargs) + ) + return Handle(worker) + + # Surface run's call signature + docstring on .send so help(.send) and + # inspect.signature show the real arguments (send forwards them to run and + # returns a Handle to poll() / wait()). + try: + send.__signature__ = inspect.signature(run_callable) + except (TypeError, ValueError): + pass + send.__doc__ = ( + "Background-launch this tool with the same arguments as run; returns a " + f"Handle to poll() / wait().\n\n{run_callable.__doc__ or ''}" + ) + module.send = send + return module diff --git a/prime-agent-runtime/test/test_async_runtime.py b/prime-agent-runtime/test/test_async_runtime.py new file mode 100644 index 0000000000..ac7f56e172 --- /dev/null +++ b/prime-agent-runtime/test/test_async_runtime.py @@ -0,0 +1,150 @@ +from __future__ import annotations + +import asyncio +import unittest + +from rlm.async_runtime import ( + ERROR, + FINISHED, + RUNNING, + BackgroundWorker, + FnProcessor, + Registry, +) + + +async def settle(predicate, *, tries: int = 2000) -> None: + """Yield to the loop until predicate() holds. + + Bounded so a wedged worker fails the test instead of hanging it. + """ + for _ in range(tries): + if predicate(): + return + await asyncio.sleep(0) + raise AssertionError("background worker never reached the expected state") + + +class Echo: + """Stateless processor that returns each item and records what it saw.""" + + def __init__(self) -> None: + self.seen: list = [] + + async def process(self, item): + self.seen.append(item) + return f"ans:{item}" + + async def teardown(self) -> None: + return None + + +class BackgroundWorkerTest(unittest.IsolatedAsyncioTestCase): + async def test_resident_worker_continues_under_one_name(self) -> None: + registry = Registry() + processor = Echo() + factory = lambda name: BackgroundWorker(name, processor) # noqa: E731 + + first = registry.send("alpha", name="helper", worker_factory=factory) + self.assertEqual(await first.wait(), "ans:alpha") + + # Re-sending the same name continues the SAME resident worker (multi-turn) + # rather than building a fresh one. + second = registry.send("beta", name="helper", worker_factory=factory) + self.assertEqual(second.name, first.name) + self.assertEqual(await second.wait(), "ans:beta") + self.assertEqual(processor.seen, ["alpha", "beta"]) + + await registry.close_all() + + async def test_poll_is_non_consuming_and_results_are_fifo(self) -> None: + registry = Registry() + factory = lambda name: BackgroundWorker(name, Echo()) # noqa: E731 + + handle = registry.send(1, name="worker", worker_factory=factory) + registry.send(2, name="worker", worker_factory=factory) + registry.send(3, name="worker", worker_factory=factory) + + await settle(lambda: len(handle.poll().results) == 3) + # poll() is a pure read: calling it repeatedly never drains a result. + self.assertEqual(len(handle.poll().results), 3) + self.assertEqual(len(handle.poll().results), 3) + + drained = [handle.poll().results.popleft() for _ in range(3)] + self.assertEqual(drained, ["ans:1", "ans:2", "ans:3"]) + self.assertEqual(len(handle.poll().results), 0) + + await registry.close_all() + + async def test_status_moves_from_running_to_finished(self) -> None: + registry = Registry() + gate = asyncio.Event() + + class Gated: + async def process(self, item): + await gate.wait() + return item + + async def teardown(self) -> None: + return None + + handle = registry.send("x", name="worker", worker_factory=lambda n: BackgroundWorker(n, Gated())) + + # Blocked inside process(): the worker reports running until released. + await settle(lambda: handle.poll().status == RUNNING) + self.assertEqual(handle.poll().status, RUNNING) + + gate.set() + self.assertEqual(await handle.wait(), "x") + + # Inbox drained -> the worker parks as finished (idle), not running. + await settle(lambda: handle.poll().status == FINISHED) + self.assertEqual(handle.poll().status, FINISHED) + + await registry.close_all() + + async def test_error_halts_worker_and_resend_rebuilds_it(self) -> None: + registry = Registry() + + class Boom: + async def process(self, item): + raise RuntimeError(f"boom:{item}") + + async def teardown(self) -> None: + return None + + class Ok: + async def process(self, item): + return f"ok:{item}" + + async def teardown(self) -> None: + return None + + handle = registry.send("x", name="worker", worker_factory=lambda n: BackgroundWorker(n, Boom())) + with self.assertRaisesRegex(RuntimeError, "boom:x"): + await handle.wait() + + state = handle.poll() + self.assertEqual(state.status, ERROR) + self.assertIsInstance(state.error, RuntimeError) + + # Re-sending the name evicts the dead worker and rebuilds a fresh one, + # so the name stays reusable. + revived = registry.send("y", name="worker", worker_factory=lambda n: BackgroundWorker(n, Ok())) + self.assertEqual(await revived.wait(), "ok:y") + + await registry.close_all() + + async def test_ephemeral_worker_runs_its_one_item_and_finishes(self) -> None: + async def add_one(value): + return value + 1 + + worker = BackgroundWorker("once", FnProcessor(add_one), item=((41,), {})) + self.assertEqual(await worker.wait(), 42) + + await settle(lambda: worker.status == FINISHED) + self.assertEqual(worker.status, FINISHED) + + +if __name__ == "__main__": + unittest.main()