diff --git a/.gitignore b/.gitignore index 666433666..d8b22e91e 100644 --- a/.gitignore +++ b/.gitignore @@ -40,3 +40,6 @@ plans/ collect.sh .kernel-venv/ .worktrees/ + +__pycache__/ +*.pyc diff --git a/packages/coding-agent/docs/kernel-and-rlm-recursion.md b/packages/coding-agent/docs/kernel-and-rlm-recursion.md index 16348277a..30034ac44 100644 --- a/packages/coding-agent/docs/kernel-and-rlm-recursion.md +++ b/packages/coding-agent/docs/kernel-and-rlm-recursion.md @@ -4,6 +4,8 @@ This document explains the IPython kernel transport and the recursive `rlm` sub- The important design constraint is that the Python `rlm` package in the kernel is only a shim. It preserves the model-facing API from `rlm-harness`, but it does not run a child agent loop in Python. Child agents are run by the TypeScript host through the same `AgentSession` machinery as the parent. +The same comm transport doubles as a generic host bridge: kernel-side Python (skills such as the bundled `goal` skill) can call `rlm.host_request("", {...})` and the TypeScript host dispatches the typed request to a registered handler. `rlm.run` is just the first registered request type. + ## High-Level Shape ```text @@ -34,7 +36,7 @@ the flow becomes: ```text Python rlm shim | - | Jupyter comm target "rlm.run" + | Jupyter comm target "host.request", payload type "rlm.run" v KernelManager comm handler | @@ -56,7 +58,8 @@ Python RLMResult ```text packages/coding-agent/src/core/kernel/index.ts ZeroMQ connection setup, Jupyter wire encoding, execute_request handling, - comm_open / comm_msg handling for rlm.run. + comm_open / comm_msg handling for the "host.request" bridge, dispatching + typed requests (rlm.run, goal.*) to registered host handlers. packages/coding-agent/src/core/tools/ipython.ts Tool wrapper around KernelManager. Starts the kernel and bootstraps the @@ -66,17 +69,23 @@ packages/coding-agent/src/core/agent-session.ts Parent session owns recursion settings and implements runRlmChild(). packages/coding-agent/src/core/rlm-runtime.ts - TypeScript request/result types for the rlm.run bridge. + TypeScript request/result types for rlm.run and the adapter that registers + it as a typed host request handler. prime-agent-runtime/src/rlm/__init__.py Python shim installed into ~/.prime/agent/kernel-venv. Exposes rlm, rlm.run(), - RLMResult, TokenUsage, and the session-backed harness state helper. + host_request(), RLMResult, TokenUsage, and the session-backed harness state + helper. prime-agent-runtime/src/rlm/harness.py Session-backed JSON store for reset-free harness refinement notes: prompt notes, memory entries, reusable skill descriptions, subagent specs, and refinement events. +packages/coding-agent/skills/goal + Bundled Python skill exposing goal.get / goal.create / goal.complete in the + kernel. A thin wrapper over host_request; all goal state stays in the host. + scripts/setup-kernel-venv.sh Thin wrapper around the automatic kernel bootstrap. ``` @@ -370,6 +379,7 @@ It exports: ```python rlm run(prompt: str, **kwargs) +host_request(request_type: str, payload: dict | None = None) RLMResult TokenUsage ``` @@ -414,18 +424,18 @@ the shim: ```text creates asyncio Future -creates Jupyter Comm target "rlm.run" +creates Jupyter Comm target "host.request" registers on_msg callback opens the comm with: { - "type": "run", + "type": "rlm.run", "prompt": "subtask", "kwargs": {} } awaits the Future ``` -The callback converts host payloads into `RLMResult` or raises `RuntimeError`. +The callback strips the reply's `status` field and converts the payload into `RLMResult`, or raises `RuntimeError` when the host reports an error. The comm-open/await machinery is the public `host_request()`; `run()` is a typed wrapper around it, and other kernel-side skills reuse `host_request()` for their own request types. The shim accepts `**kwargs` for API compatibility with `rlm-harness` and includes them in the comm payload. V1 does not apply any kwargs on the TypeScript side. Unsupported kwargs are rejected with a clear error instead of being ignored. @@ -474,27 +484,29 @@ The comm handler tracks: ```text commTargets: comm_id -> target_name -handledRlmCommIds: comm ids already started +handledHostRequestCommIds: comm ids already started ``` -For target `rlm.run`, it starts work asynchronously: +For target `host.request`, it starts work asynchronously and dispatches on the payload `type`: ```text -comm_open target "rlm.run" +comm_open target "host.request" | v -startRlmRunFromComm(comm_id, data) +startHostRequestFromComm(comm_id, data) | v -handleRlmRunRequest(data) +handleHostRequest(data) | v -options.rlmRunHandler({ prompt, kwargs }) +options.hostHandlers[data.type](data) | v sendCommMessage(comm_id, { status: "ok", ...result }) ``` +Unknown types fail with `host request type "" is not available in this session`. + Errors are returned as: ```json @@ -508,10 +520,15 @@ The send path uses `control` if available, then falls back to `shell`. ## Child AgentSession Execution -The root `AgentSession` passes this handler into the root kernel: +The root `AgentSession` registers typed handlers into the root kernel: ```typescript -rlmRunHandler: ({ prompt, kwargs }) => this.runRlmChild(prompt, kwargs) +hostHandlers: { + "rlm.run": createRlmRunHostHandler(({ prompt, kwargs }) => this.runRlmChild(prompt, kwargs)), + "goal.get": ..., // only when goals are enabled + "goal.create": ..., + "goal.complete": ..., +} ``` `runRlmChild()` is the TypeScript owner of actual recursion. @@ -585,6 +602,39 @@ The answer is the child's final assistant text. This matches the RLM-1 training `RLMResult.usage` reports the child session's token usage. The parent session's displayed usage and cost do not currently fold child usage into the parent assistant-message totals. Child cost is independent in v1 and should be read from the child result or child session until parent aggregation is added. +## Goal Skill Over the Host Bridge + +Goals are not harness tools. The model's only built-in tool is `ipython`; the goal feature is exposed as the bundled Python skill `goal`, which the kernel bootstrap imports into the user namespace like any other Python skill: + +```python +await goal.get() +await goal.create("ship the release notes", token_budget=200000) +await goal.complete() +``` + +Each call is a thin wrapper over `rlm.host_request("goal.", {...})`. All goal state — status transitions, token/wall-clock accounting, persistence, continuation re-prompting — lives in the TypeScript `AgentSession`: + +```text +model runs `await goal.complete()` inside an ipython cell + | + | comm target "host.request", payload { "type": "goal.complete" } + v +KernelManager dispatches to hostHandlers["goal.complete"] + | + v +AgentSession.handleGoalHostRequest("goal.complete") + | flips goal state to complete, marks the run for usage accounting, + | emits goal_update to connected clients + v +comm reply { status: "ok", goal: {...}, completion_budget_report: ... } +``` + +Replies use snake_case keys (`tokens_used`, `remaining_tokens`, `completion_budget_report`) since they are a Python-facing API. + +The re-prompting loop is unchanged and entirely host-side: while a goal is active the session injects `goal_context` continuation messages after each turn, steers a budget-limit context when the token budget is exhausted, and stops continuing once `goal.complete` arrives over the bridge (or the user pauses/clears the goal). Because completion arrives mid-cell rather than as a detectable tool call, the session records that completion was requested during the run so the completing turn's tokens still count toward the goal budget. + +When goals are disabled for a session (`includeGoals: false`), the goal skill is filtered out of the kernel and the system prompt, and the `goal.*` handlers are not registered — calls fail with the unknown-type error above. + ## Session Directory Layout For a persisted root session: diff --git a/packages/coding-agent/skills/goal/SKILL.md b/packages/coding-agent/skills/goal/SKILL.md new file mode 100644 index 000000000..55b475bb9 --- /dev/null +++ b/packages/coding-agent/skills/goal/SKILL.md @@ -0,0 +1,44 @@ +--- +name: goal +description: Manage the persistent thread goal from IPython. Use to read goal status and budget usage, to start a goal when the user explicitly asks for one, or to mark the active goal complete once its objective is fully achieved. +--- + +# Goal + +The thread goal is a persistent objective the harness keeps re-prompting you to +pursue across turns until it is complete. Goal state (status, token budget, +usage accounting) lives in the host; this skill is the kernel-side interface to +it. Call it directly from IPython: + +```python +await goal.get() +await goal.create("ship the release notes", token_budget=200000) +await goal.complete() +``` + +## API + +- `await goal.get()` — current goal as a dict: `goal` (or `None` when no goal + is set), `remaining_tokens`, and `completion_budget_report`. The `goal` dict + carries `objective`, `status`, `token_budget`, `tokens_used`, + `time_used_seconds`, and timestamps. +- `await goal.create(objective, token_budget=None)` — start a new active goal. + Fails while a goal is still pending (active, paused, or budget-limited); a + completed or errored goal is replaced by the new one. Only create a goal when + the user or system/developer instructions explicitly ask for a persistent + long-running goal; do not infer goals from ordinary tasks. Set `token_budget` + only when an explicit token budget is requested. +- `await goal.complete()` — mark the existing goal achieved. Use only when the + objective has actually been achieved and no required work remains; do not + call it merely because the budget is nearly exhausted or because you are + stopping work. When the result includes a `completion_budget_report`, report + that final usage to the user. + +## Rules + +- Goal status transitions other than completion (pause, resume, clear, + budget-limiting) are controlled by the user and the host; there is no API for + them here. +- When an active goal is actually complete, call `await goal.complete()`; do + not merely say it is done — the harness keeps continuing the goal until the + completion call arrives. diff --git a/packages/coding-agent/skills/goal/pyproject.toml b/packages/coding-agent/skills/goal/pyproject.toml new file mode 100644 index 000000000..5b76222b8 --- /dev/null +++ b/packages/coding-agent/skills/goal/pyproject.toml @@ -0,0 +1,17 @@ +# Kernel-side package for the bundled goal skill. It only talks to the host +# through rlm.host_request; prime-agent-runtime is always installed in the +# kernel venv before skills, so it is intentionally not declared as a +# dependency (it is not published on PyPI). +[project] +name = "goal" +version = "0.1.0" +description = "Prime Agent goal skill: persistent thread goal control over the host bridge" +requires-python = ">=3.10" +dependencies = [] + +[build-system] +requires = ["hatchling"] +build-backend = "hatchling.build" + +[tool.hatch.build.targets.wheel] +packages = ["src/goal"] diff --git a/packages/coding-agent/skills/goal/src/goal/__init__.py b/packages/coding-agent/skills/goal/src/goal/__init__.py new file mode 100644 index 000000000..891e0bfbd --- /dev/null +++ b/packages/coding-agent/skills/goal/src/goal/__init__.py @@ -0,0 +1,52 @@ +"""Prime Agent goal skill: manage the persistent thread goal from the kernel. + +All goal state lives in the TypeScript host; these functions are thin typed +wrappers over the generic host bridge (`rlm.host_request`). They only work +inside the Prime Agent IPython kernel. +""" + +from __future__ import annotations + +from typing import Any + +from rlm import host_request + + +async def get() -> dict[str, Any]: + """Read the current thread goal. + + Returns a dict with `goal` (None when no goal is set), `remaining_tokens`, + and `completion_budget_report`. The `goal` dict carries the objective, + status, token budget, and token/elapsed-time usage. + """ + return await host_request("goal.get") + + +async def create(objective: str, token_budget: int | None = None) -> dict[str, Any]: + """Start a new active thread goal. + + Fails while a goal is still pending (active, paused, or budget-limited); + a completed or errored goal is replaced. Only create a goal when the user + or system/developer instructions explicitly ask for a persistent + long-running goal. Set `token_budget` only when an explicit token budget is + requested. + """ + if not isinstance(objective, str): + raise TypeError(f"objective must be str, got {type(objective).__name__}") + if token_budget is not None and not isinstance(token_budget, int): + raise TypeError(f"token_budget must be int or None, got {type(token_budget).__name__}") + payload: dict[str, Any] = {"objective": objective} + if token_budget is not None: + payload["token_budget"] = token_budget + return await host_request("goal.create", payload) + + +async def complete() -> dict[str, Any]: + """Mark the existing thread goal achieved. + + Use only when the objective has actually been achieved and no required + work remains — not because the budget is nearly exhausted or because you + are stopping work. Pause, resume, and budget-limit transitions are + controlled by the user and the host. + """ + return await host_request("goal.complete") diff --git a/packages/coding-agent/src/core/agent-session-runtime.ts b/packages/coding-agent/src/core/agent-session-runtime.ts index bac9f0b4a..1bd4e27a0 100644 --- a/packages/coding-agent/src/core/agent-session-runtime.ts +++ b/packages/coding-agent/src/core/agent-session-runtime.ts @@ -263,8 +263,7 @@ export class AgentSessionRuntime implements SubagentRuntimeHost { initialActiveToolNames: options.activeToolNames, allowedToolNames: options.allowedToolNames, customTools: options.customTools, - includeGoalTools: options.includeGoalTools, - autoActivateGoalTools: options.autoActivateGoalTools, + includeGoals: options.includeGoals, rlmDepth: options.rlmDepth, rlmMaxDepth: options.rlmMaxDepth, rlmSessionDir: options.sessionDir, diff --git a/packages/coding-agent/src/core/agent-session-services.ts b/packages/coding-agent/src/core/agent-session-services.ts index 16ab56bfe..46ed2458f 100644 --- a/packages/coding-agent/src/core/agent-session-services.ts +++ b/packages/coding-agent/src/core/agent-session-services.ts @@ -49,8 +49,7 @@ export interface AgentSessionCreationOptions { customTools?: ToolDefinition[]; initialActiveToolNames?: string[]; allowedToolNames?: string[]; - includeGoalTools?: boolean; - autoActivateGoalTools?: boolean; + includeGoals?: boolean; rlmDepth?: number; rlmMaxDepth?: number; rlmSessionDir?: string; @@ -209,8 +208,7 @@ export async function createAgentSessionFromServices( customTools: options.customTools, initialActiveToolNames: options.initialActiveToolNames, allowedToolNames: options.allowedToolNames, - includeGoalTools: options.includeGoalTools, - autoActivateGoalTools: options.autoActivateGoalTools, + includeGoals: options.includeGoals, rlmDepth: options.rlmDepth, rlmMaxDepth: options.rlmMaxDepth, rlmSessionDir: options.rlmSessionDir, diff --git a/packages/coding-agent/src/core/agent-session.ts b/packages/coding-agent/src/core/agent-session.ts index 08ea5cc25..1d020d1f8 100644 --- a/packages/coding-agent/src/core/agent-session.ts +++ b/packages/coding-agent/src/core/agent-session.ts @@ -99,20 +99,21 @@ import { import { emitSessionShutdownEvent } from "./extensions/runner.js"; import { createGoalContextMessage, - createGoalToolDefinitions, emptyGoalState, GOAL_CONTEXT_CUSTOM_TYPE, + GOAL_SKILL_NAME, GOAL_STATE_CUSTOM_TYPE, - GOAL_TOOL_NAMES, + type GoalHostResponse, type GoalState, type GoalStatus, + goalHostResponse, goalTokenDeltaForUsage, isPersistedGoalState, normalizeGoalState, - UPDATE_GOAL_TOOL_NAME, validateGoalBudget, validateGoalObjective, } from "./goals.js"; +import type { HostRequestHandlers } from "./kernel/index.js"; import type { BashExecutionMessage, CustomMessage } from "./messages.js"; import type { ModelRegistry } from "./model-registry.js"; import { expandPromptTemplate, type PromptTemplate } from "./prompt-templates.js"; @@ -129,13 +130,14 @@ import { saveHarnessState, } from "./refinement/index.js"; import type { ResourceExtensionPaths, ResourceLoader } from "./resource-loader.js"; -import type { - CreateRlmSubagentRuntimeOptions, - RlmInternalRunResult, - RlmRunResult, - RlmSubagentRuntime, - RlmUsage, - SubagentRuntimeHost, +import { + type CreateRlmSubagentRuntimeOptions, + createRlmRunHostHandler, + type RlmInternalRunResult, + type RlmRunResult, + type RlmSubagentRuntime, + type RlmUsage, + type SubagentRuntimeHost, } from "./rlm-runtime.js"; import type { BranchSummaryEntry, CompactionEntry, SessionMessageEntry } from "./session-manager.js"; import { @@ -146,7 +148,7 @@ import { } from "./session-manager.js"; import type { SessionStats } from "./session-stats.js"; import type { SettingsManager } from "./settings-manager.js"; -import { getPythonSkillRuntimeInfo } from "./skills.js"; +import { getPythonSkillRuntimeInfo, type Skill } from "./skills.js"; import type { SlashCommandInfo } from "./slash-commands.js"; import { createSyntheticSourceInfo, type SourceInfo } from "./source-info.js"; import { type BuildSystemPromptOptions, buildSystemPrompt } from "./system-prompt.js"; @@ -290,10 +292,12 @@ export interface AgentSessionConfig { initialActiveToolNames?: string[]; /** Optional allowlist of tool names. When provided, only these tool names are exposed. */ allowedToolNames?: string[]; - /** Whether the built-in long-running goal tools are exposed. Default: true. */ - includeGoalTools?: boolean; - /** Whether goal tools are automatically active without an explicit goal. Default: true. */ - autoActivateGoalTools?: boolean; + /** + * Whether the built-in long-running goals feature is available: the bundled + * goal skill in the IPython kernel, its goal.* host handlers, and /goal. + * Default: true. + */ + includeGoals?: boolean; /** * Override base tools (useful for custom runtimes). * @@ -679,8 +683,7 @@ export class AgentSession { private _extensionRunnerRef?: { current?: ExtensionRunner }; private _initialActiveToolNames?: string[]; private _allowedToolNames?: Set; - private _includeGoalTools: boolean; - private _autoActivateGoalTools: boolean; + private _includeGoals: boolean; private _baseToolsOverride?: Record; private _sessionStartEvent: SessionStartEvent; private _extensionUIContext?: ExtensionUIContext; @@ -723,8 +726,7 @@ export class AgentSession { this._extensionRunnerRef = config.extensionRunnerRef; this._initialActiveToolNames = config.initialActiveToolNames; this._allowedToolNames = config.allowedToolNames ? new Set(config.allowedToolNames) : undefined; - this._includeGoalTools = config.includeGoalTools ?? true; - this._autoActivateGoalTools = config.autoActivateGoalTools ?? true; + this._includeGoals = config.includeGoals ?? true; this._baseToolsOverride = config.baseToolsOverride; this._sessionStartEvent = config.sessionStartEvent ?? { type: "session_start", reason: "startup" }; this._rlmDepth = config.rlmDepth ?? parseDepth(process.env.RLM_DEPTH, 0, "RLM_DEPTH"); @@ -1125,38 +1127,30 @@ export class AgentSession { } } - private _ensureGoalToolsActive(context?: AgentContext): void { - const goalTools: AgentTool[] = []; - for (const toolName of GOAL_TOOL_NAMES) { - const tool = this._toolRegistry.get(toolName); - if (!tool) { - throw new Error("Goal tools are disabled. Enable goal tools before using /goal."); - } - goalTools.push(tool); + /** + * Goals are pursued through the IPython goal skill, so the only tool the + * model needs is ipython. Force-activate it (including into a live + * continuation context) so the model can always reach `goal.complete()`. + */ + private _ensureGoalRuntimeActive(context?: AgentContext): void { + if (!this._includeGoals) { + throw new Error("Goals are disabled. Enable goals before using /goal."); } - if (!this._includeGoalTools) { - throw new Error("Goal tools are disabled. Enable goal tools before using /goal."); + const ipythonTool = this._toolRegistry.get("ipython"); + if (!ipythonTool) { + throw new Error("Goals require the ipython tool, which is not available in this session."); } const activeToolNames = new Set(this.getActiveToolNames()); - let changed = false; - for (const toolName of GOAL_TOOL_NAMES) { - if (!activeToolNames.has(toolName)) { - activeToolNames.add(toolName); - changed = true; - } - } - if (changed) { + if (!activeToolNames.has("ipython")) { + activeToolNames.add("ipython"); this.setActiveToolsByName([...activeToolNames]); } if (context) { const contextTools = [...(context.tools ?? [])]; - const contextToolNames = new Set(contextTools.map((tool) => tool.name)); - for (const tool of goalTools) { - if (!contextToolNames.has(tool.name)) { - contextTools.push(tool); - } + if (!contextTools.some((tool) => tool.name === "ipython")) { + contextTools.push(ipythonTool); + context.tools = contextTools; } - context.tools = contextTools; } } @@ -1167,7 +1161,7 @@ export class AgentSession { if (!this._goalState.objective) { return; } - this._ensureGoalToolsActive(); + this._ensureGoalRuntimeActive(); const message = createGoalContextMessage(this._goalState, kind, images); if (this.isStreaming) { if (kind === "budget_limit") { @@ -1213,7 +1207,7 @@ export class AgentSession { if (!this.isStreaming) { await this._validateCanStartAgentRun(); } - this._ensureGoalToolsActive(); + this._ensureGoalRuntimeActive(); this._clearQueuedGoalContexts(); this._startGoal(command.objective, command.tokenBudget); await this._runOrQueueGoalContext(previousWasActive ? "objective_updated" : "continuation", images); @@ -1230,10 +1224,12 @@ export class AgentSession { if (this._goalAccountedAssistantMessages.has(message)) { return false; } - const updateGoalRequested = message.content.some( - (content) => content.type === "toolCall" && content.name === UPDATE_GOAL_TOOL_NAME, - ); - if (this._goalState.status !== "active" && !(this._goalState.status === "complete" && updateGoalRequested)) { + // Usage is attributed at the assistant message's message_end, which fires + // before that turn's ipython cell runs. goal.complete() only arrives later + // over the kernel host bridge, so the completing turn is always accounted + // while the goal is still active. Only count turns spent pursuing the goal; + // post-completion turns (e.g. a closing summary) must not be attributed. + if (this._goalState.status !== "active") { return false; } this._goalAccountedAssistantMessages.add(message); @@ -1243,8 +1239,7 @@ export class AgentSession { ...goal, tokensUsed: goal.tokensUsed + tokenDelta, }; - const budgetReached = - !updateGoalRequested && nextGoal.tokenBudget !== undefined && nextGoal.tokensUsed >= nextGoal.tokenBudget; + const budgetReached = nextGoal.tokenBudget !== undefined && nextGoal.tokensUsed >= nextGoal.tokenBudget; if (!budgetReached) { this._setGoalState(nextGoal); return false; @@ -1298,20 +1293,63 @@ export class AgentSession { return true; } - private createGoalFromTool(objective: string, tokenBudget: number | undefined): GoalState { - if (this._goalState.status !== "idle") { - throw new Error( - "cannot create a new goal because this thread already has a goal; use update_goal only when the existing goal is complete", - ); + /** + * Handle a goal.* request from the IPython kernel host bridge (the bundled + * goal skill). All goal state stays host-side; the kernel only sees the + * serialized snake_case response. + */ + handleGoalHostRequest(type: string, payload: Record = {}): GoalHostResponse { + if (!this._includeGoals) { + throw new Error("goals are disabled in this session"); + } + switch (type) { + case "goal.get": + return goalHostResponse(this.goalState, false); + case "goal.create": { + if (typeof payload.objective !== "string") { + throw new Error("goal.create objective must be a string"); + } + if (payload.token_budget !== undefined && typeof payload.token_budget !== "number") { + throw new Error("goal.create token_budget must be an integer when provided"); + } + return goalHostResponse(this._createGoalFromHost(payload.objective, payload.token_budget), false); + } + case "goal.complete": + return goalHostResponse(this._completeGoalFromHost(), true); + default: + throw new Error(`unknown goal request type "${type}"`); + } + } + + private _createGoalFromHost(objective: string, tokenBudget: number | undefined): GoalState { + switch (this._goalState.status) { + case "active": + throw new Error( + "cannot create a new goal because this thread already has an active goal; run `await goal.complete()` when it is achieved, or ask the user to clear it with /goal clear", + ); + case "paused": + throw new Error( + "cannot create a new goal because a paused goal exists; ask the user to resume it with /goal resume or clear it with /goal clear", + ); + case "budget_limited": + throw new Error( + "cannot create a new goal because a budget-limited goal exists; ask the user to resume it with /goal resume or clear it with /goal clear", + ); + default: + // idle, or a terminal record (complete / error): nothing pending, start fresh. + return this._startGoal(objective, tokenBudget); } - return this._startGoal(objective, tokenBudget); } - private completeGoalFromTool(): GoalState { + private _completeGoalFromHost(): GoalState { if (!this._goalState.objective || this._goalState.status === "idle") { - throw new Error("cannot update goal because this thread has no goal"); + throw new Error("cannot complete goal because this thread has no goal"); } const goal = this._goalWithAccountedWallClock(); + // A turn can cross the budget and complete the goal at once: accounting + // runs at message_end, before the completing ipython cell executes, so a + // budget-limit context may already be steered. It is stale now — drop it. + this._clearQueuedGoalContexts(); this._setGoalState({ ...goal, active: false, @@ -1333,7 +1371,7 @@ export class AgentSession { return []; } try { - this._ensureGoalToolsActive(context.context); + this._ensureGoalRuntimeActive(context.context); const nextGoal = { ...this._goalState, continuationsUsed: this._goalState.continuationsUsed + 1, @@ -1861,7 +1899,7 @@ export class AgentSession { const loaderAppendSystemPrompt = this._resourceLoader.getAppendSystemPrompt(); const appendSystemPrompt = loaderAppendSystemPrompt.length > 0 ? loaderAppendSystemPrompt.join("\n\n") : undefined; - const loadedSkills = this._resourceLoader.getSkills().skills; + const loadedSkills = this._modelVisibleSkills(); const loadedContextFiles = this._resourceLoader.getAgentsFiles().agentsFiles; this._baseSystemPromptOptions = { @@ -3359,7 +3397,7 @@ export class AgentSession { }): void { const shellCommandPrefix = this.settingsManager.getShellCommandPrefix(); const shellPath = this.settingsManager.getShellPath(); - const pythonSkills = getPythonSkillRuntimeInfo(this._resourceLoader.getSkills().skills); + const pythonSkills = getPythonSkillRuntimeInfo(this._modelVisibleSkills()); let configuredBaseToolDefinitions: Record; if (this._baseToolsOverride) { configuredBaseToolDefinitions = Object.fromEntries( @@ -3375,7 +3413,7 @@ export class AgentSession { this._ipythonKernelProvisioner = new IpythonKernelProvisioner(this._cwd, { env: this._rlmKernelEnv(), sessionId: this.sessionId, - rlmRunHandler: ({ prompt, kwargs }) => this.runRlmChild(prompt, kwargs), + hostHandlers: this._createKernelHostHandlers(), pythonSkills, }); configuredBaseToolDefinitions = createAllToolDefinitions(this._cwd, { @@ -3383,20 +3421,9 @@ export class AgentSession { bash: { commandPrefix: shellCommandPrefix, shellPath }, }); } - const goalToolDefinitions = this._includeGoalTools - ? createGoalToolDefinitions({ - getGoalState: () => this.goalState, - createGoalFromTool: (objective, tokenBudget) => this.createGoalFromTool(objective, tokenBudget), - completeGoalFromTool: () => this.completeGoalFromTool(), - }) - : []; - const baseToolDefinitions = { - ...configuredBaseToolDefinitions, - ...Object.fromEntries(goalToolDefinitions.map((definition) => [definition.name, definition])), - }; this._baseToolDefinitions = new Map( - Object.entries(baseToolDefinitions).map(([name, tool]) => [name, tool as ToolDefinition]), + Object.entries(configuredBaseToolDefinitions).map(([name, tool]) => [name, tool as ToolDefinition]), ); const extensionsResult = this._resourceLoader.getExtensions(); @@ -3420,12 +3447,10 @@ export class AgentSession { this._applyExtensionBindings(this._extensionRunner); const defaultActiveToolNames = this._baseToolsOverride ? Object.keys(this._baseToolsOverride) : ["ipython"]; - if (this._includeGoalTools && this._autoActivateGoalTools) { - defaultActiveToolNames.push(...GOAL_TOOL_NAMES); - } const baseActiveToolNames = [...(options.activeToolNames ?? defaultActiveToolNames)]; - if (this._goalState.status === "active" && this._includeGoalTools) { - baseActiveToolNames.push(...GOAL_TOOL_NAMES); + if (this._goalState.status === "active" && this._includeGoals) { + // An active goal needs ipython so the model can reach the goal skill. + baseActiveToolNames.push("ipython"); } this._refreshToolRegistry({ activeToolNames: [...new Set(baseActiveToolNames)], @@ -3437,6 +3462,31 @@ export class AgentSession { } } + /** + * Skills exposed to the model (system prompt + kernel). The bundled goal + * skill is withheld when goals are disabled for this session. + */ + private _modelVisibleSkills(): Skill[] { + const skills = this._resourceLoader.getSkills().skills; + if (this._includeGoals) { + return skills; + } + return skills.filter((skill) => skill.name !== GOAL_SKILL_NAME); + } + + /** Typed handlers for host requests arriving from the IPython kernel comm bridge. */ + private _createKernelHostHandlers(): HostRequestHandlers { + const handlers: HostRequestHandlers = { + "rlm.run": createRlmRunHostHandler(({ prompt, kwargs }) => this.runRlmChild(prompt, kwargs)), + }; + if (this._includeGoals) { + for (const type of ["goal.get", "goal.create", "goal.complete"]) { + handlers[type] = async (payload) => this.handleGoalHostRequest(type, payload); + } + } + return handlers; + } + async reload(): Promise { const previousFlagValues = this._extensionRunner.getFlagValues(); await emitSessionShutdownEvent(this._extensionRunner, { type: "session_shutdown", reason: "reload" }); @@ -3564,8 +3614,7 @@ export class AgentSession { activeToolNames: this.getActiveToolNames(), allowedToolNames: this._allowedToolNames ? [...this._allowedToolNames] : undefined, customTools: [...this._customTools], - includeGoalTools: this._includeGoalTools, - autoActivateGoalTools: this._autoActivateGoalTools, + includeGoals: this._includeGoals, rlmDepth: this._rlmDepth + 1, rlmMaxDepth: this._rlmMaxDepth, rlmParentNodeId: options.id, @@ -3633,8 +3682,7 @@ export class AgentSession { modelRegistry: this._modelRegistry, initialActiveToolNames: options.activeToolNames, allowedToolNames: options.allowedToolNames, - includeGoalTools: options.includeGoalTools, - autoActivateGoalTools: options.autoActivateGoalTools, + includeGoals: options.includeGoals, rlmDepth: options.rlmDepth, rlmMaxDepth: options.rlmMaxDepth, rlmSessionDir: options.sessionDir, diff --git a/packages/coding-agent/src/core/goals.ts b/packages/coding-agent/src/core/goals.ts index 51c8c092a..52728e705 100644 --- a/packages/coding-agent/src/core/goals.ts +++ b/packages/coding-agent/src/core/goals.ts @@ -1,14 +1,9 @@ import type { ImageContent, TextContent } from "@earendil-works/pi-ai"; -import { type Static, Type } from "typebox"; -import type { ToolDefinition } from "./extensions/types.js"; import type { CustomMessage } from "./messages.js"; export const GOAL_STATE_CUSTOM_TYPE = "thread_goal_state"; export const GOAL_CONTEXT_CUSTOM_TYPE = "goal_context"; -export const GET_GOAL_TOOL_NAME = "get_goal"; -export const CREATE_GOAL_TOOL_NAME = "create_goal"; -export const UPDATE_GOAL_TOOL_NAME = "update_goal"; -export const GOAL_TOOL_NAMES = [GET_GOAL_TOOL_NAME, CREATE_GOAL_TOOL_NAME, UPDATE_GOAL_TOOL_NAME] as const; +export const GOAL_SKILL_NAME = "goal"; export const MAX_THREAD_GOAL_OBJECTIVE_CHARS = 4000; export type GoalStatus = "idle" | "active" | "paused" | "budget_limited" | "complete" | "error"; @@ -29,22 +24,24 @@ export interface GoalState { lastError?: string; } -export interface SerializedGoal { - goalId?: string; +/** Goal payload returned to the kernel-side goal skill. Keys are Python-conventional snake_case. */ +export type SerializedGoal = { + goal_id?: string; objective: string; - status: Exclude | "error"; - tokenBudget?: number; - tokensUsed: number; - timeUsedSeconds: number; - createdAt?: number; - updatedAt?: number; -} - -export interface GoalToolResponse { + status: Exclude; + token_budget?: number; + tokens_used: number; + time_used_seconds: number; + created_at?: number; + updated_at?: number; +}; + +/** Reply payload for goal.* host requests from the IPython kernel. */ +export type GoalHostResponse = { goal: SerializedGoal | null; - remainingTokens: number | null; - completionBudgetReport: string | null; -} + remaining_tokens: number | null; + completion_budget_report: string | null; +}; export interface GoalContextDetails { kind: GoalContextKind; @@ -54,38 +51,6 @@ export interface GoalContextDetails { continuationsUsed: number; } -export interface GoalToolController { - getGoalState(): GoalState; - createGoalFromTool(objective: string, tokenBudget: number | undefined): GoalState; - completeGoalFromTool(): GoalState; -} - -const getGoalSchema = Type.Object({}, { additionalProperties: false }); -const createGoalSchema = Type.Object( - { - objective: Type.String({ - description: - "Required. The concrete objective to start pursuing. This starts a new active goal only when no goal is currently defined; if a goal already exists, this tool fails.", - }), - token_budget: Type.Optional( - Type.Integer({ description: "Optional positive token budget for the new active goal." }), - ), - }, - { additionalProperties: false }, -); -const updateGoalSchema = Type.Object( - { - status: Type.Literal("complete", { - description: "Required. Set to complete only when the objective is achieved and no required work remains.", - }), - }, - { additionalProperties: false }, -); - -type GetGoalArgs = Static; -type CreateGoalArgs = Static; -type UpdateGoalArgs = Static; - export function emptyGoalState(): GoalState { return { active: false, @@ -156,99 +121,35 @@ export function isPersistedGoalState(value: unknown): value is GoalState { ); } -export function goalToolResponse(goal: GoalState, includeCompletionReport: boolean): GoalToolResponse { +export function goalHostResponse(goal: GoalState, includeCompletionReport: boolean): GoalHostResponse { if (goal.status === "idle" || !goal.objective) { return { goal: null, - remainingTokens: null, - completionBudgetReport: null, + remaining_tokens: null, + completion_budget_report: null, }; } const remainingTokens = goal.tokenBudget === undefined ? null : Math.max(0, goal.tokenBudget - goal.tokensUsed); const serializedGoal: SerializedGoal = { - goalId: goal.goalId, + goal_id: goal.goalId, objective: goal.objective, status: goal.status, - tokenBudget: goal.tokenBudget, - tokensUsed: goal.tokensUsed, - timeUsedSeconds: goal.timeUsedSeconds, - createdAt: goal.createdAt, - updatedAt: goal.updatedAt, + token_budget: goal.tokenBudget, + tokens_used: goal.tokensUsed, + time_used_seconds: goal.timeUsedSeconds, + created_at: goal.createdAt, + updated_at: goal.updatedAt, }; return { goal: serializedGoal, - remainingTokens, - completionBudgetReport: + remaining_tokens: remainingTokens, + completion_budget_report: includeCompletionReport && goal.status === "complete" ? completionBudgetReport(goal) : null, }; } -export function createGoalToolDefinitions(controller: GoalToolController): ToolDefinition[] { - return [ - { - name: GET_GOAL_TOOL_NAME, - label: "Get Goal", - description: - "Get the current goal for this thread, including status, budgets, token and elapsed-time usage, and remaining token budget.", - parameters: getGoalSchema, - execute: async (_toolCallId: string, _params: GetGoalArgs) => { - const response = goalToolResponse(controller.getGoalState(), false); - return { - content: [{ type: "text", text: JSON.stringify(response, null, 2) }], - details: response, - }; - }, - }, - { - name: CREATE_GOAL_TOOL_NAME, - label: "Create Goal", - description: `Create a goal only when explicitly requested by the user or system/developer instructions; do not infer goals from ordinary tasks. -Set token_budget only when an explicit token budget is requested. Fails if a goal exists; use ${UPDATE_GOAL_TOOL_NAME} only for status.`, - promptGuidelines: [ - `Use ${CREATE_GOAL_TOOL_NAME} only when the user or higher-priority instructions explicitly ask you to set a persistent long-running goal.`, - ], - parameters: createGoalSchema, - execute: async (_toolCallId: string, params: CreateGoalArgs) => { - const goal = controller.createGoalFromTool(params.objective, params.token_budget); - const response = goalToolResponse(goal, false); - return { - content: [{ type: "text", text: JSON.stringify(response, null, 2) }], - details: response, - }; - }, - }, - { - name: UPDATE_GOAL_TOOL_NAME, - label: "Update Goal", - description: `Update the existing goal. -Use this tool only to mark the goal achieved. -Set status to complete only when the objective has actually been achieved and no required work remains. -Do not mark a goal complete merely because its budget is nearly exhausted or because you are stopping work. -You cannot use this tool to pause, resume, or budget-limit a goal; those status changes are controlled by the user or system. -When marking a budgeted goal achieved with status complete, report the final token usage from the tool result to the user.`, - promptGuidelines: [ - `When an active goal is actually complete, call ${UPDATE_GOAL_TOOL_NAME} with status "complete"; do not merely say it is done.`, - ], - parameters: updateGoalSchema, - execute: async (_toolCallId: string, params: UpdateGoalArgs) => { - if (params.status !== "complete") { - throw new Error( - "update_goal can only mark the existing goal complete; pause, resume, and budget-limited status changes are controlled by the user or system", - ); - } - const goal = controller.completeGoalFromTool(); - const response = goalToolResponse(goal, true); - return { - content: [{ type: "text", text: JSON.stringify(response, null, 2) }], - details: response, - }; - }, - }, - ]; -} - export function createGoalContextMessage( goal: GoalState, kind: GoalContextKind, @@ -322,9 +223,9 @@ Goal state: The goal persists across turns. Ending one turn does not reduce or redefine the objective. If the goal is not complete yet, make concrete progress toward the full objective. -Before marking the goal complete, audit the current state against every requirement in the objective. Do not rely on intent, partial progress, memory of earlier work, or a plausible final answer as proof of completion. If the objective is achieved, call update_goal with status "complete" so usage accounting is preserved. +Before marking the goal complete, audit the current state against every requirement in the objective. Do not rely on intent, partial progress, memory of earlier work, or a plausible final answer as proof of completion. If the objective is achieved, run \`await goal.complete()\` in ipython so usage accounting is preserved. -Do not call update_goal unless the goal is complete. Do not mark a goal complete merely because the budget is nearly exhausted or because you are stopping work.`; +Do not call \`goal.complete()\` unless the goal is complete. Do not mark a goal complete merely because the budget is nearly exhausted or because you are stopping work.`; } function budgetLimitPrompt(goal: GoalState): string { @@ -345,7 +246,7 @@ Goal state: The system has marked the goal budget_limited. Do not start new substantive work. Wrap up this turn soon with progress made, remaining work, blockers, and a concrete next step. -Do not call update_goal unless the goal is actually complete.`; +Do not run \`await goal.complete()\` unless the goal is actually complete.`; } function objectiveUpdatedPrompt(goal: GoalState): string { @@ -366,7 +267,7 @@ Goal state: - token budget: ${budget} - remaining tokens: ${remaining} -Adjust the current turn to pursue the updated objective. Do not call update_goal unless the updated goal is actually complete.`; +Adjust the current turn to pursue the updated objective. Do not run \`await goal.complete()\` unless the updated goal is actually complete.`; } function completionBudgetReport(goal: GoalState): string | null { diff --git a/packages/coding-agent/src/core/kernel/bootstrap.ts b/packages/coding-agent/src/core/kernel/bootstrap.ts index b2dab9a40..76bf3b2f3 100644 --- a/packages/coding-agent/src/core/kernel/bootstrap.ts +++ b/packages/coding-agent/src/core/kernel/bootstrap.ts @@ -48,7 +48,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 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, '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; @@ -577,7 +577,9 @@ async function ensureKernelPythonUncached( const missing: string[] = []; if (!(await hasIpykernel(python))) missing.push("ipykernel"); if (!(await hasPrimeAgentRuntime(python))) { - missing.push("a current prime-agent-runtime with callable rlm.run and explicit harness CRUD methods"); + missing.push( + "a current prime-agent-runtime with callable rlm.run, rlm.host_request, and explicit harness CRUD methods", + ); } if (missing.length === 0) { const missingExtraImports = await missingRlmExtraImportLabels(python); diff --git a/packages/coding-agent/src/core/kernel/index.ts b/packages/coding-agent/src/core/kernel/index.ts index 2dbdb4f74..8d1a74689 100644 --- a/packages/coding-agent/src/core/kernel/index.ts +++ b/packages/coding-agent/src/core/kernel/index.ts @@ -8,7 +8,6 @@ import { setTimeout as sleep } from "node:timers/promises"; import { registerSessionResourceCleanup } from "@earendil-works/pi-ai"; import { v4 as uuid } from "uuid"; import { Dealer, Subscriber } from "zeromq"; -import type { RlmRunHandler, RlmRunResult } from "../rlm-runtime.js"; import { ensureKernelPython, type KernelBootstrapProgressHandler, type KernelPythonSkill } from "./bootstrap.js"; const DELIM = Buffer.from(""); @@ -18,7 +17,19 @@ const READY_TIMEOUT_MS = 5000; // Loopback PUB/SUB subscription propagation is usually sub-ms, but keep a small guard before first execute. const IOPUB_SUBSCRIBE_DELAY_MS = 50; const DEFAULT_MAX_OUTPUT_CHARS = 65536; -const RLM_DISPOSE_TIMEOUT_MS = 5000; +const HOST_REQUEST_DISPOSE_TIMEOUT_MS = 5000; + +/** Comm target the kernel-side `rlm.host_request` shim opens for typed host requests. */ +export const HOST_COMM_TARGET = "host.request"; + +/** + * Handles one typed request from Python code running in the kernel. + * The returned record is sent back verbatim as the comm reply payload. + */ +export type HostRequestHandler = (payload: Record) => Promise>; + +/** Host request handlers keyed by request type (e.g. "rlm.run", "goal.complete"). */ +export type HostRequestHandlers = Record; export interface KernelManagerOptions { /** Python interpreter that has `ipykernel` available. Defaults to the auto-bootstrapped kernel. */ @@ -26,7 +37,7 @@ export interface KernelManagerOptions { cwd?: string; env?: Record; sessionId?: string; - rlmRunHandler?: RlmRunHandler; + hostHandlers?: HostRequestHandlers; pythonSkills?: readonly KernelPythonSkill[]; /** Default: "prime-agent". */ username?: string; @@ -54,8 +65,6 @@ export interface ExecuteResult { durationMs: number; } -type RlmCommResult = RlmRunResult; - interface ConnectionInfo { ip: string; transport: "tcp"; @@ -287,12 +296,12 @@ function installSignalHandlersOnce(): void { export class KernelManager { private readonly options: Pick< KernelManagerOptions, - "python" | "cwd" | "env" | "sessionId" | "rlmRunHandler" | "pythonSkills" + "python" | "cwd" | "env" | "sessionId" | "hostHandlers" | "pythonSkills" > & Required>; private readonly session = uuid(); private readonly commTargets = new Map(); - private readonly handledRlmCommIds = new Set(); + private readonly handledHostRequestCommIds = new Set(); private kernel?: ChildProcess; private shell?: Dealer; private iopub?: Subscriber; @@ -304,7 +313,7 @@ export class KernelManager { /** Serializes execute() calls — Jupyter shell channel is request/reply. */ private executionQueue: Promise = Promise.resolve(); private activeExecution?: ActiveExecution; - private readonly inFlightRlmRuns = new Set>(); + private readonly inFlightHostRequests = new Set>(); private state: "idle" | "starting" | "running" | "shutdown" = "idle"; /** Memoized so concurrent callers all await the same in-flight startup. */ private startPromise?: Promise; @@ -315,7 +324,7 @@ export class KernelManager { cwd: options.cwd, env: options.env, sessionId: options.sessionId, - rlmRunHandler: options.rlmRunHandler, + hostHandlers: options.hostHandlers, pythonSkills: options.pythonSkills, username: options.username ?? "prime-agent", }; @@ -707,7 +716,7 @@ export class KernelManager { if (msgType === "comm_close") { this.commTargets.delete(commId); - this.handledRlmCommIds.delete(commId); + this.handledHostRequestCommIds.delete(commId); return; } @@ -717,69 +726,64 @@ export class KernelManager { return; } this.commTargets.set(commId, targetName); - if (targetName === "rlm.run") { - this.startRlmRunFromComm(commId, content.data); + if (targetName === HOST_COMM_TARGET) { + this.startHostRequestFromComm(commId, content.data); } return; } const targetName = this.commTargets.get(commId); - if (msgType === "comm_msg" && targetName === "rlm.run") { - this.startRlmRunFromComm(commId, content.data); + if (msgType === "comm_msg" && targetName === HOST_COMM_TARGET) { + this.startHostRequestFromComm(commId, content.data); } } - private startRlmRunFromComm(commId: string, data: unknown): void { - if (this.handledRlmCommIds.has(commId)) { + private startHostRequestFromComm(commId: string, data: unknown): void { + if (this.handledHostRequestCommIds.has(commId)) { return; } - this.handledRlmCommIds.add(commId); + this.handledHostRequestCommIds.add(commId); const task = (async () => { try { - const result = await this.handleRlmCommRequest(data); + const result = await this.handleHostRequest(data); try { await this.sendCommMessage(commId, { status: "ok", ...result }); } catch (replyError) { this.appendKernelDiagnostic( - `failed to send rlm.run ok reply for comm ${commId}: ${errorMessage(replyError)}`, + `failed to send host request ok reply for comm ${commId}: ${errorMessage(replyError)}`, ); } } catch (error) { - this.appendKernelDiagnostic(`rlm.run failed for comm ${commId}: ${errorMessage(error)}`); + this.appendKernelDiagnostic(`host request failed for comm ${commId}: ${errorMessage(error)}`); try { await this.sendCommMessage(commId, { status: "error", error: errorMessage(error) }); } catch (replyError) { this.appendKernelDiagnostic( - `failed to send rlm.run error reply for comm ${commId}: ${errorMessage(replyError)}`, + `failed to send host request error reply for comm ${commId}: ${errorMessage(replyError)}`, ); } } })(); - this.inFlightRlmRuns.add(task); + this.inFlightHostRequests.add(task); void task.finally(() => { - this.inFlightRlmRuns.delete(task); + this.inFlightHostRequests.delete(task); }); } - private async handleRlmCommRequest(data: unknown): Promise { + private async handleHostRequest(data: unknown): Promise> { if (!isRecord(data)) { - throw new Error("rlm.run comm payload must be an object"); + throw new Error("host request payload must be an object"); } - - if (data.type === "run") { - const handler = this.options.rlmRunHandler; - if (!handler) { - throw new Error("rlm.run is not available in this session"); - } - if (typeof data.prompt !== "string") { - throw new Error("rlm.run prompt must be a string"); - } - const kwargs = isRecord(data.kwargs) ? data.kwargs : {}; - return handler({ prompt: data.prompt, kwargs }); + if (typeof data.type !== "string" || data.type.length === 0) { + throw new Error("host request payload must have a string type"); } - throw new Error("rlm.run comm payload must have a supported type"); + const handler = this.options.hostHandlers?.[data.type]; + if (!handler) { + throw new Error(`host request type "${data.type}" is not available in this session`); + } + return handler(data); } private async sendCommMessage(commId: string, data: Record): Promise { @@ -820,7 +824,7 @@ export class KernelManager { this.startPromise = undefined; } - private async waitForRlmRunsToSettle(tasks: Promise[], timeoutMs: number): Promise { + private async waitForHostRequestsToSettle(tasks: Promise[], timeoutMs: number): Promise { let timeout: ReturnType | undefined; const timeoutPromise = new Promise<"timeout">((resolve) => { timeout = globalThis.setTimeout(() => resolve("timeout"), timeoutMs); @@ -835,7 +839,7 @@ export class KernelManager { } if (result === "timeout") { this.appendKernelDiagnostic( - `timed out waiting ${timeoutMs}ms for ${tasks.length} rlm.run task(s) during dispose`, + `timed out waiting ${timeoutMs}ms for ${tasks.length} host request task(s) during dispose`, ); } } @@ -878,12 +882,12 @@ export class KernelManager { } } - /** Graceful cleanup. Waits briefly for in-flight rlm.run handlers before closing sockets. */ + /** Graceful cleanup. Waits briefly for in-flight host request handlers before closing sockets. */ dispose(): Promise { this.state = "shutdown"; liveKernels.delete(this); - const inFlightRlmRuns = [...this.inFlightRlmRuns]; - if (inFlightRlmRuns.length === 0) { + const inFlightHostRequests = [...this.inFlightHostRequests]; + if (inFlightHostRequests.length === 0) { this.cleanupResources(); return Promise.resolve(); } @@ -891,7 +895,7 @@ export class KernelManager { return (async () => { // TODO: plumb AbortSignal through AgentSession.prompt so disposal can cancel long-running child loops. try { - await this.waitForRlmRunsToSettle(inFlightRlmRuns, RLM_DISPOSE_TIMEOUT_MS); + await this.waitForHostRequestsToSettle(inFlightHostRequests, HOST_REQUEST_DISPOSE_TIMEOUT_MS); } finally { this.cleanupResources(); } diff --git a/packages/coding-agent/src/core/rlm-runtime.ts b/packages/coding-agent/src/core/rlm-runtime.ts index 01d1ddc10..71530481a 100644 --- a/packages/coding-agent/src/core/rlm-runtime.ts +++ b/packages/coding-agent/src/core/rlm-runtime.ts @@ -2,6 +2,7 @@ import type { ThinkingLevel } from "@earendil-works/pi-agent-core"; import type { Model, Usage } 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"; export interface RlmUsage { prompt_tokens: number; @@ -26,6 +27,22 @@ export interface RlmInternalRunResult extends RlmRunResult { export type RlmRunHandler = (request: RlmRunRequest) => Promise; +function isRecord(value: unknown): value is Record { + return typeof value === "object" && value !== null && !Array.isArray(value); +} + +/** Adapt an RlmRunHandler into the typed "rlm.run" handler for the kernel host bridge. */ +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 result = await handler({ prompt: payload.prompt, kwargs }); + return result as unknown as Record; + }; +} + export interface RlmSubagentRuntime { session: AgentSession; } @@ -41,8 +58,7 @@ export interface CreateRlmSubagentRuntimeOptions { activeToolNames: string[]; allowedToolNames?: string[]; customTools: ToolDefinition[]; - includeGoalTools: boolean; - autoActivateGoalTools: boolean; + includeGoals: boolean; rlmDepth: number; rlmMaxDepth: number; rlmParentNodeId: string; diff --git a/packages/coding-agent/src/core/sdk.ts b/packages/coding-agent/src/core/sdk.ts index 76c560dbe..c6d9f3d0e 100644 --- a/packages/coding-agent/src/core/sdk.ts +++ b/packages/coding-agent/src/core/sdk.ts @@ -8,7 +8,6 @@ import { formatNoModelsAvailableMessage } from "./auth-guidance.js"; import { AuthStorage } from "./auth-storage.js"; import { DEFAULT_THINKING_LEVEL } from "./defaults.js"; import type { ExtensionRunner, LoadExtensionsResult, SessionStartEvent, ToolDefinition } from "./extensions/index.js"; -import { GOAL_TOOL_NAMES } from "./goals.js"; import { convertToLlm } from "./messages.js"; import { ModelRegistry } from "./model-registry.js"; import { findInitialModel } from "./model-resolver.js"; @@ -225,13 +224,9 @@ export async function createAgentSession(options: CreateAgentSessionOptions = {} } const allowedToolNames = options.allowedToolNames ?? options.tools ?? (options.noTools === "all" ? [] : undefined); - const includeGoalTools = options.includeGoalTools ?? (options.tools !== undefined || options.noTools !== "all"); - const autoActivateGoalTools = options.autoActivateGoalTools ?? (options.tools === undefined && !options.noTools); - const defaultActiveToolNames: string[] = - includeGoalTools && autoActivateGoalTools ? ["ipython", ...GOAL_TOOL_NAMES] : ["ipython"]; + const includeGoals = options.includeGoals ?? (options.tools !== undefined || options.noTools !== "all"); const initialActiveToolNames: string[] = - options.initialActiveToolNames ?? - (options.tools ? [...options.tools] : options.noTools ? [] : defaultActiveToolNames); + options.initialActiveToolNames ?? (options.tools ? [...options.tools] : options.noTools ? [] : ["ipython"]); let agent: Agent; @@ -353,8 +348,7 @@ export async function createAgentSession(options: CreateAgentSessionOptions = {} modelRegistry, initialActiveToolNames, allowedToolNames, - includeGoalTools, - autoActivateGoalTools, + includeGoals, extensionRunnerRef, rlmDepth: options.rlmDepth, rlmMaxDepth: options.rlmMaxDepth, diff --git a/packages/coding-agent/src/core/tools/ipython.ts b/packages/coding-agent/src/core/tools/ipython.ts index c58790183..fc5d25daf 100644 --- a/packages/coding-agent/src/core/tools/ipython.ts +++ b/packages/coding-agent/src/core/tools/ipython.ts @@ -3,8 +3,7 @@ import type { AgentTool } from "@earendil-works/pi-agent-core"; import { type Static, Type } from "typebox"; import type { ToolDefinition } from "../extensions/types.js"; import type { KernelBootstrapProgressHandler } from "../kernel/bootstrap.js"; -import { KernelManager } from "../kernel/index.js"; -import type { RlmRunHandler } from "../rlm-runtime.js"; +import { type HostRequestHandlers, KernelManager } from "../kernel/index.js"; import type { PythonSkillRuntimeInfo } from "../skills.js"; import { wrapToolDefinition } from "./tool-definition-wrapper.js"; @@ -137,7 +136,8 @@ export interface IpythonToolOptions { python?: string; env?: Record; sessionId?: string; - rlmRunHandler?: RlmRunHandler; + /** Typed host request handlers for the kernel↔host bridge (rlm.run, goal.*, …). */ + hostHandlers?: HostRequestHandlers; pythonSkills?: readonly PythonSkillRuntimeInfo[]; /** Filled after the first kernel start so the owning session can restart it after compaction. */ kernelManagerRef?: { current?: KernelManager }; @@ -245,7 +245,7 @@ export class IpythonKernelProvisioner { cwd: this.cwd, env: this.options?.env, sessionId: this.options?.sessionId, - rlmRunHandler: this.options?.rlmRunHandler, + hostHandlers: this.options?.hostHandlers, pythonSkills: this.options?.pythonSkills, }); this.emitStartupProgress("Starting IPython kernel..."); diff --git a/packages/coding-agent/src/main.ts b/packages/coding-agent/src/main.ts index fe4e37d9f..526cf7a57 100644 --- a/packages/coding-agent/src/main.ts +++ b/packages/coding-agent/src/main.ts @@ -1091,8 +1091,7 @@ export async function main(args: string[], options?: MainOptions) { customTools: runtimeSessionOptions?.customTools ?? sessionOptions.customTools, initialActiveToolNames: runtimeSessionOptions?.initialActiveToolNames, allowedToolNames: runtimeSessionOptions?.allowedToolNames, - includeGoalTools: runtimeSessionOptions?.includeGoalTools, - autoActivateGoalTools: runtimeSessionOptions?.autoActivateGoalTools, + includeGoals: runtimeSessionOptions?.includeGoals, rlmDepth: runtimeSessionOptions?.rlmDepth, rlmMaxDepth: runtimeSessionOptions?.rlmMaxDepth, rlmSessionDir: runtimeSessionOptions?.rlmSessionDir, diff --git a/packages/coding-agent/src/modes/daemon/daemon-mode.ts b/packages/coding-agent/src/modes/daemon/daemon-mode.ts index b989fda19..9f88b4ab0 100644 --- a/packages/coding-agent/src/modes/daemon/daemon-mode.ts +++ b/packages/coding-agent/src/modes/daemon/daemon-mode.ts @@ -402,8 +402,7 @@ class AgentDaemon { initialActiveToolNames: options.activeToolNames, allowedToolNames: options.allowedToolNames, customTools: options.customTools, - includeGoalTools: options.includeGoalTools, - autoActivateGoalTools: options.autoActivateGoalTools, + includeGoals: options.includeGoals, rlmDepth: options.rlmDepth, rlmMaxDepth: options.rlmMaxDepth, rlmSessionDir: options.sessionDir, diff --git a/packages/coding-agent/test/agent-session-recursion.test.ts b/packages/coding-agent/test/agent-session-recursion.test.ts index 2cf59b981..b4b6e27a6 100644 --- a/packages/coding-agent/test/agent-session-recursion.test.ts +++ b/packages/coding-agent/test/agent-session-recursion.test.ts @@ -17,6 +17,7 @@ import { AuthStorage } from "../src/core/auth-storage.js"; import { KernelManager } from "../src/core/kernel/index.js"; import { convertToLlm } from "../src/core/messages.js"; import { ModelRegistry } from "../src/core/model-registry.js"; +import { createRlmRunHostHandler } from "../src/core/rlm-runtime.js"; import { SessionManager } from "../src/core/session-manager.js"; import { SettingsManager } from "../src/core/settings-manager.js"; import { MISSING_RIPGREP_MESSAGE } from "../src/utils/tools-manager.js"; @@ -145,14 +146,14 @@ function rlmCommOpenData(commId: string, data: Record): TestCom metadata: {}, content: { comm_id: commId, - target_name: "rlm.run", + target_name: "host.request", data, }, }; } function rlmCommOpen(commId: string, prompt: string, kwargs: Record = {}): TestCommMessage { - return rlmCommOpenData(commId, { type: "run", prompt, kwargs }); + return rlmCommOpenData(commId, { type: "rlm.run", prompt, kwargs }); } function encodeTestMessage(message: TestCommMessage): Buffer[] { @@ -565,18 +566,20 @@ describe("AgentSession rlm recursion", () => { const replies: CapturedCommReply[] = []; const manager = new KernelManager({ python: process.execPath, - rlmRunHandler: async ({ prompt }) => { - active++; - started++; - maxActive = Math.max(maxActive, active); - await release; - active--; - return { - answer: `answer:${prompt}`, - usage: { prompt_tokens: 1, completion_tokens: 1 }, - turns: 1, - session_dir: null, - }; + hostHandlers: { + "rlm.run": createRlmRunHostHandler(async ({ prompt }) => { + active++; + started++; + maxActive = Math.max(maxActive, active); + await release; + active--; + return { + answer: `answer:${prompt}`, + usage: { prompt_tokens: 1, completion_tokens: 1 }, + turns: 1, + session_dir: null, + }; + }), }, }); @@ -620,14 +623,16 @@ describe("AgentSession rlm recursion", () => { let promptSeen = ""; const manager = new KernelManager({ python: process.execPath, - rlmRunHandler: async ({ prompt }) => { - promptSeen = prompt; - return { - answer: `answer:${prompt}`, - usage: { prompt_tokens: 1, completion_tokens: 1 }, - turns: 1, - session_dir: null, - }; + hostHandlers: { + "rlm.run": createRlmRunHostHandler(async ({ prompt }) => { + promptSeen = prompt; + return { + answer: `answer:${prompt}`, + usage: { prompt_tokens: 1, completion_tokens: 1 }, + turns: 1, + session_dir: null, + }; + }), }, }); @@ -662,14 +667,16 @@ describe("AgentSession rlm recursion", () => { const prompts: string[] = []; const manager = new KernelManager({ cwd: tempDir, - rlmRunHandler: async ({ prompt }) => { - prompts.push(prompt); - return { - answer: `answer:${prompt}`, - usage: { prompt_tokens: 1, completion_tokens: 1 }, - turns: 1, - session_dir: null, - }; + hostHandlers: { + "rlm.run": createRlmRunHostHandler(async ({ prompt }) => { + prompts.push(prompt); + return { + answer: `answer:${prompt}`, + usage: { prompt_tokens: 1, completion_tokens: 1 }, + turns: 1, + session_dir: null, + }; + }), }, }); @@ -739,12 +746,14 @@ print(_result.answer) const replies: CapturedCommReply[] = []; const manager = new KernelManager({ python: process.execPath, - rlmRunHandler: async () => ({ - answer: "unused", - usage: { prompt_tokens: 1, completion_tokens: 1 }, - turns: 1, - session_dir: null, - }), + hostHandlers: { + "rlm.run": createRlmRunHostHandler(async () => ({ + answer: "unused", + usage: { prompt_tokens: 1, completion_tokens: 1 }, + turns: 1, + session_dir: null, + })), + }, }); try { @@ -761,7 +770,7 @@ print(_result.answer) commId: "comm-bg", data: { status: "error", - error: "rlm.run comm payload must have a supported type", + error: 'host request type "background" is not available in this session', }, }); } finally { @@ -783,14 +792,16 @@ print(_result.answer) }); const manager = new KernelManager({ python: process.execPath, - rlmRunHandler: async () => { - started = true; - try { - await release; - throw new Error("child failed after dispose"); - } finally { - handlerSettled = true; - } + hostHandlers: { + "rlm.run": createRlmRunHostHandler(async () => { + started = true; + try { + await release; + throw new Error("child failed after dispose"); + } finally { + handlerSettled = true; + } + }), }, }); const stderrSpy = vi.spyOn(process.stderr, "write").mockImplementation(() => true); @@ -815,8 +826,8 @@ print(_result.answer) expect(handlerSettled).toBe(true); const kernelStderr = (manager as unknown as { kernelStderr: string }).kernelStderr; - expect(kernelStderr).toContain("[kernel] rlm.run failed for comm comm-dispose"); - expect(kernelStderr).toContain("[kernel] failed to send rlm.run error reply for comm comm-dispose"); + expect(kernelStderr).toContain("[kernel] host request failed for comm comm-dispose"); + expect(kernelStderr).toContain("[kernel] failed to send host request error reply for comm comm-dispose"); expect(stderrSpy).not.toHaveBeenCalled(); } finally { releaseChild(); diff --git a/packages/coding-agent/test/builtin-skills.test.ts b/packages/coding-agent/test/builtin-skills.test.ts index ec2a023f5..a8690222e 100644 --- a/packages/coding-agent/test/builtin-skills.test.ts +++ b/packages/coding-agent/test/builtin-skills.test.ts @@ -163,6 +163,15 @@ describe("builtin skills", () => { expect(skills.map((s) => s.name)).toContain("skill-creator"); }); + it("loads the bundled goal skill as a python skill", () => { + const { skills } = loadSkillsFromDir({ dir: getBundledSkillsDir(), source: "builtin" }); + + const goal = skills.find((s) => s.name === "goal"); + expect(goal).toBeDefined(); + expect(goal?.kind).toBe("python"); + expect(goal?.kind === "python" && goal.python.importName).toBe("goal"); + }); + it("ships the edit skill as a python skill importable as `edit`", () => { const { skills } = loadSkillsFromDir({ dir: getBundledSkillsDir(), source: "builtin" }); diff --git a/packages/coding-agent/test/kernel-bootstrap.test.ts b/packages/coding-agent/test/kernel-bootstrap.test.ts index 81e3ccab3..6bb719d93 100644 --- a/packages/coding-agent/test/kernel-bootstrap.test.ts +++ b/packages/coding-agent/test/kernel-bootstrap.test.ts @@ -13,6 +13,7 @@ import { let tempDir = ""; let originalEnv: NodeJS.ProcessEnv; + function pyprojectHash(pyprojectPath: string): string { return `sha256:${createHash("sha256").update(readFileSync(pyprojectPath)).digest("hex")}`; } diff --git a/packages/coding-agent/test/kernel-goal-skill.test.ts b/packages/coding-agent/test/kernel-goal-skill.test.ts new file mode 100644 index 000000000..b72c87af6 --- /dev/null +++ b/packages/coding-agent/test/kernel-goal-skill.test.ts @@ -0,0 +1,151 @@ +import { mkdirSync, rmSync } from "node:fs"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import { afterEach, beforeEach, describe, expect, it } from "vitest"; +import { getBundledSkillsDir } from "../src/config.js"; +import type { PythonSkillRuntimeInfo } from "../src/core/skills.js"; +import { IpythonKernelProvisioner } from "../src/core/tools/ipython.js"; + +function bundledGoalSkill(): PythonSkillRuntimeInfo { + const packagePath = join(getBundledSkillsDir(), "goal"); + return { + name: "goal", + importName: "goal", + packagePath, + pyprojectPath: join(packagePath, "pyproject.toml"), + }; +} + +describe("goal skill over the kernel host bridge", () => { + let tempDir: string; + let provisioner: IpythonKernelProvisioner | undefined; + + beforeEach(() => { + tempDir = join(tmpdir(), `pi-goal-skill-${Date.now()}-${Math.random().toString(36).slice(2)}`); + mkdirSync(tempDir, { recursive: true }); + }); + + afterEach(async () => { + await provisioner?.dispose(); + provisioner = undefined; + rmSync(tempDir, { recursive: true, force: true }); + }); + + it("round-trips goal.create and goal.complete through a live kernel", async () => { + const requests: Array<{ type: string; payload: Record }> = []; + provisioner = new IpythonKernelProvisioner(tempDir, { + pythonSkills: [bundledGoalSkill()], + hostHandlers: { + "goal.create": async (payload) => { + requests.push({ type: "goal.create", payload }); + return { + goal: { objective: payload.objective, status: "active", tokens_used: 0 }, + remaining_tokens: payload.token_budget ?? null, + completion_budget_report: null, + }; + }, + "goal.complete": async (payload) => { + requests.push({ type: "goal.complete", payload }); + return { + goal: { objective: "ship it", status: "complete", tokens_used: 7 }, + remaining_tokens: 3, + completion_budget_report: + "Goal achieved. Report final budget usage to the user: tokens used: 7 of 10.", + }; + }, + }, + }); + + const manager = await provisioner.ensure(); + const created = await manager.execute(` +import json +_created = await goal.create("ship it", token_budget=10) +print(json.dumps(_created, sort_keys=True)) +`); + expect(created.status).toBe("ok"); + expect(JSON.parse(created.stdout.trim())).toEqual({ + goal: { objective: "ship it", status: "active", tokens_used: 0 }, + remaining_tokens: 10, + completion_budget_report: null, + }); + + const completed = await manager.execute(` +_completed = await goal.complete() +print(_completed["goal"]["status"], _completed["completion_budget_report"]) +`); + expect(completed.status).toBe("ok"); + expect(completed.stdout.trim()).toBe( + "complete Goal achieved. Report final budget usage to the user: tokens used: 7 of 10.", + ); + + expect(requests.map((request) => request.type)).toEqual(["goal.create", "goal.complete"]); + expect(requests[0].payload).toMatchObject({ type: "goal.create", objective: "ship it", token_budget: 10 }); + }); + + it("surfaces host errors and missing handlers as Python exceptions", async () => { + provisioner = new IpythonKernelProvisioner(tempDir, { + pythonSkills: [bundledGoalSkill()], + hostHandlers: { + "goal.complete": async () => { + throw new Error("cannot complete goal because this thread has no goal"); + }, + }, + }); + + const manager = await provisioner.ensure(); + const completeError = await manager.execute(` +try: + await goal.complete() +except RuntimeError as error: + print(f"RuntimeError: {error}") +`); + expect(completeError.status).toBe("ok"); + expect(completeError.stdout.trim()).toBe("RuntimeError: cannot complete goal because this thread has no goal"); + + const unavailable = await manager.execute(` +try: + await goal.get() +except RuntimeError as error: + print(f"RuntimeError: {error}") +`); + expect(unavailable.status).toBe("ok"); + expect(unavailable.stdout.trim()).toBe( + 'RuntimeError: host request type "goal.get" is not available in this session', + ); + + // A "type" key smuggled into the payload must not reroute the request. + const reroute = await manager.execute(` +import rlm as _rlm +try: + await _rlm.host_request("goal.get", {"type": "goal.complete"}) +except RuntimeError as error: + print(f"RuntimeError: {error}") +`); + expect(reroute.status).toBe("ok"); + expect(reroute.stdout.trim()).toBe('RuntimeError: host request type "goal.get" is not available in this session'); + }); + + it("rejects replies with an unexpected status instead of hanging", async () => { + provisioner = new IpythonKernelProvisioner(tempDir, { + pythonSkills: [bundledGoalSkill()], + hostHandlers: { + "goal.get": async () => ({ goal: null, remaining_tokens: null, completion_budget_report: null }), + }, + }); + + const manager = await provisioner.ensure(); + type CommSender = { sendCommMessage: (commId: string, data: Record) => Promise }; + const kernel = manager as unknown as CommSender; + const originalSend = kernel.sendCommMessage.bind(manager); + kernel.sendCommMessage = (commId, _data) => originalSend(commId, { status: "partial" }); + + const result = await manager.execute(` +try: + await goal.get() +except RuntimeError as error: + print(f"RuntimeError: {error}") +`); + expect(result.status).toBe("ok"); + expect(result.stdout.trim()).toBe("RuntimeError: host request goal.get returned unexpected status: 'partial'"); + }); +}); diff --git a/packages/coding-agent/test/suite/agent-session-goal.test.ts b/packages/coding-agent/test/suite/agent-session-goal.test.ts index 775a8a825..cbf9c075f 100644 --- a/packages/coding-agent/test/suite/agent-session-goal.test.ts +++ b/packages/coding-agent/test/suite/agent-session-goal.test.ts @@ -2,8 +2,9 @@ import type { AgentContext, AgentTool } from "@earendil-works/pi-agent-core"; import { type AssistantMessage, fauxAssistantMessage, fauxToolCall, type Usage } from "@earendil-works/pi-ai"; import { Type } from "typebox"; import { afterEach, describe, expect, it, vi } from "vitest"; -import type { ExtensionContext, ExtensionFactory } from "../../src/core/extensions/types.js"; -import { GOAL_TOOL_NAMES } from "../../src/core/goals.js"; +import type { AgentSession } from "../../src/core/agent-session.js"; +import type { ExtensionFactory } from "../../src/core/extensions/types.js"; +import type { GoalHostResponse } from "../../src/core/goals.js"; import { createHarness, getAssistantTexts, getMessageText, type Harness } from "./harness.js"; function assistantWithUsage(message: string | AssistantMessage, usage: Partial): AssistantMessage { @@ -50,6 +51,43 @@ async function waitForCondition(predicate: () => boolean): Promise { throw new Error("condition was not met"); } +/** + * Stand-in for the real ipython tool. Goal calls reach the host over the + * kernel comm bridge while an ipython cell executes; this stub mirrors that + * timing by dispatching `goal.*` host requests from inside tool execution. + * + * Cell format: `goal.` optionally followed by a JSON payload, e.g. + * `goal.create {"objective": "write a note"}`. + */ +function createFauxIpythonTool(sessionRef: { current?: AgentSession }): AgentTool { + return { + name: "ipython", + label: "ipython", + description: "Execute Python code in the agent kernel.", + parameters: Type.Object({ code: Type.String() }), + execute: async (_toolCallId, params) => { + const session = sessionRef.current; + if (!session) { + throw new Error("test session is not initialized"); + } + const code = (params as { code: string }).code.trim(); + let text = ""; + if (code.startsWith("goal.")) { + const spaceIndex = code.indexOf(" "); + const type = spaceIndex < 0 ? code : code.slice(0, spaceIndex); + const payload = spaceIndex < 0 ? {} : JSON.parse(code.slice(spaceIndex + 1)); + text = JSON.stringify(session.handleGoalHostRequest(type, payload)); + } + return { + content: [{ type: "text", text }], + details: {}, + }; + }, + }; +} + +const COMPLETE_GOAL_CELL = { code: "goal.complete" }; + function createWaitingTool(): { tool: AgentTool; release: () => void; @@ -108,13 +146,20 @@ describe("AgentSession goals", () => { } }); - it("keeps continuing until the model calls update_goal complete", async () => { - const harness = await createHarness(); + async function createGoalHarness(extraTools: AgentTool[] = []): Promise { + const sessionRef: { current?: AgentSession } = {}; + const harness = await createHarness({ tools: [createFauxIpythonTool(sessionRef), ...extraTools] }); + sessionRef.current = harness.session; harnesses.push(harness); + return harness; + } + + it("keeps continuing until the model completes the goal through ipython", async () => { + const harness = await createGoalHarness(); harness.setResponses([ fauxAssistantMessage("I need another step."), fauxAssistantMessage("The work is complete."), - fauxAssistantMessage(fauxToolCall("update_goal", { status: "complete" }), { stopReason: "toolUse" }), + fauxAssistantMessage(fauxToolCall("ipython", COMPLETE_GOAL_CELL), { stopReason: "toolUse" }), fauxAssistantMessage("Goal complete."), ]); @@ -136,12 +181,11 @@ describe("AgentSession goals", () => { expect(harness.getPendingResponseCount()).toBe(0); }); - it("counts tokens from the update_goal completion turn", async () => { - const harness = await createHarness(); - harnesses.push(harness); + it("counts tokens from the goal completion turn", async () => { + const harness = await createGoalHarness(); harness.setResponses([ assistantWithUsage( - fauxAssistantMessage(fauxToolCall("update_goal", { status: "complete" }), { stopReason: "toolUse" }), + fauxAssistantMessage(fauxToolCall("ipython", COMPLETE_GOAL_CELL), { stopReason: "toolUse" }), { input: 4, output: 2, @@ -160,59 +204,151 @@ describe("AgentSession goals", () => { expect(harness.session.goalState.tokensUsed).toBeGreaterThan(0); }); - it("activates goal tools when a slash goal starts from an inactive tool set", async () => { - const harness = await createHarness(); - harnesses.push(harness); + it("does not count post-completion turns against the finished goal", async () => { + const harness = await createGoalHarness(); + harness.setResponses([ + assistantWithUsage( + fauxAssistantMessage(fauxToolCall("ipython", COMPLETE_GOAL_CELL), { stopReason: "toolUse" }), + { input: 4, output: 2, totalTokens: 6 }, + ), + assistantWithUsage( + "Goal complete; here is a long closing summary that must not be billed to the finished goal.", + { input: 20, output: 10, totalTokens: 30 }, + ), + ]); + + await harness.session.prompt("/goal finish the task"); + + expect(harness.session.goalState.status).toBe("complete"); + const completeUpdates = harness.eventsOfType("goal_update").filter((event) => event.goal.status === "complete"); + expect(completeUpdates.length).toBeGreaterThan(0); + const tokensAtCompletion = completeUpdates[0].goal.tokensUsed; + expect(tokensAtCompletion).toBeGreaterThan(0); + // The closing-summary turn runs after goal.complete() over the host bridge; + // it must not increase the finished goal's token usage. + expect(harness.session.goalState.tokensUsed).toBe(tokensAtCompletion); + }); + + it("returns the goal snapshot and completion report over the host bridge", async () => { + const harness = await createGoalHarness(); + + expect(harness.session.handleGoalHostRequest("goal.get")).toEqual({ + goal: null, + remaining_tokens: null, + completion_budget_report: null, + }); + + const created = harness.session.handleGoalHostRequest("goal.create", { + objective: "write a benchmark note", + token_budget: 50, + }); + expect(created.goal).toMatchObject({ + objective: "write a benchmark note", + status: "active", + token_budget: 50, + tokens_used: 0, + }); + expect(created.remaining_tokens).toBe(50); + + const completed = harness.session.handleGoalHostRequest("goal.complete"); + expect(completed.goal).toMatchObject({ status: "complete" }); + expect(completed.completion_budget_report).toContain("tokens used: 0 of 50"); + }); + + it("rejects malformed and unknown goal host requests", async () => { + const harness = await createGoalHarness(); + + expect(() => harness.session.handleGoalHostRequest("goal.create", {})).toThrow( + "goal.create objective must be a string", + ); + expect(() => harness.session.handleGoalHostRequest("goal.nonsense")).toThrow( + 'unknown goal request type "goal.nonsense"', + ); + expect(() => harness.session.handleGoalHostRequest("goal.complete")).toThrow( + "cannot complete goal because this thread has no goal", + ); + + harness.session.handleGoalHostRequest("goal.create", { objective: "first goal" }); + expect(() => harness.session.handleGoalHostRequest("goal.create", { objective: "second goal" })).toThrow( + "already has an active goal", + ); + }); + + it("lets the model create a fresh goal after the previous one completed", async () => { + const harness = await createGoalHarness(); + + const first = harness.session.handleGoalHostRequest("goal.create", { objective: "first goal" }); + harness.session.handleGoalHostRequest("goal.complete"); + + const second = harness.session.handleGoalHostRequest("goal.create", { objective: "second goal" }); + expect(second.goal).toMatchObject({ objective: "second goal", status: "active", tokens_used: 0 }); + expect(second.goal?.goal_id).not.toBe(first.goal?.goal_id); + expect(harness.session.goalState).toMatchObject({ + active: true, + status: "active", + objective: "second goal", + continuationsUsed: 0, + }); + }); + + it("rejects goal.create while a goal is paused", async () => { + const waiting = createWaitingTool(); + const harness = await createGoalHarness([waiting.tool]); + harness.setResponses([fauxAssistantMessage(fauxToolCall("wait", {}), { stopReason: "toolUse" })]); + + const waitForStart = waiting.waitForStart(harness); + const promptPromise = harness.session.prompt("/goal long task"); + await waitForStart; + await harness.session.prompt("/goal pause"); + waiting.release(); + await promptPromise; + + expect(harness.session.goalState.status).toBe("paused"); + expect(() => harness.session.handleGoalHostRequest("goal.create", { objective: "replacement" })).toThrow( + "a paused goal exists; ask the user to resume it with /goal resume or clear it with /goal clear", + ); + }); + + it("activates ipython when a slash goal starts from an inactive tool set", async () => { + const harness = await createGoalHarness(); harness.session.setActiveToolsByName([]); harness.setResponses([ - fauxAssistantMessage(fauxToolCall("update_goal", { status: "complete" }), { stopReason: "toolUse" }), + fauxAssistantMessage(fauxToolCall("ipython", COMPLETE_GOAL_CELL), { stopReason: "toolUse" }), fauxAssistantMessage("Goal complete."), ]); await harness.session.prompt("/goal finish the task"); - expect(harness.session.getActiveToolNames()).toEqual([...GOAL_TOOL_NAMES]); + expect(harness.session.getActiveToolNames()).toEqual(["ipython"]); expect(harness.session.goalState).toMatchObject({ active: false, status: "complete", }); }); - it("adds goal tools to the live continuation context when inactive at run start", async () => { - const harness = await createHarness(); - harnesses.push(harness); - const createGoalTool = harness.session.getToolDefinition("create_goal"); - if (!createGoalTool) { - throw new Error("expected create_goal tool"); - } - await createGoalTool.execute( - "create-goal", - { objective: "finish the active goal" }, - undefined, - undefined, - {} as ExtensionContext, - ); + it("adds ipython to the live continuation context when inactive at run start", async () => { + const harness = await createGoalHarness(); + harness.session.handleGoalHostRequest("goal.create", { objective: "finish the active goal" }); harness.session.setActiveToolsByName([]); harness.setResponses([ fauxAssistantMessage("Still working."), - fauxAssistantMessage(fauxToolCall("update_goal", { status: "complete" }), { stopReason: "toolUse" }), + fauxAssistantMessage(fauxToolCall("ipython", COMPLETE_GOAL_CELL), { stopReason: "toolUse" }), fauxAssistantMessage("Goal complete."), ]); await harness.session.prompt("continue"); expect(visibleAssistantTexts(harness)).toEqual(["Still working.", "Goal complete."]); - expect(harness.session.getActiveToolNames()).toEqual([...GOAL_TOOL_NAMES]); + expect(harness.session.getActiveToolNames()).toEqual(["ipython"]); expect(harness.session.goalState).toMatchObject({ active: false, status: "complete", }); }); - it("does not re-add deactivated goal tools on runtime rebuild without an active goal", async () => { - const harness = await createHarness(); - harnesses.push(harness); - expect(harness.session.getActiveToolNames()).toEqual(["ipython", ...GOAL_TOOL_NAMES]); + it("does not re-add deactivated tools on runtime rebuild without an active goal", async () => { + const harness = await createGoalHarness(); + expect(harness.session.getActiveToolNames()).toEqual(["ipython"]); harness.session.setActiveToolsByName([]); await harness.session.reload(); @@ -220,40 +356,18 @@ describe("AgentSession goals", () => { expect(harness.session.getActiveToolNames()).toEqual([]); }); - it("does not duplicate goal tools on active-goal runtime rebuild", async () => { - const harness = await createHarness(); - harnesses.push(harness); - const createGoalTool = harness.session.getToolDefinition("create_goal"); - if (!createGoalTool) { - throw new Error("expected create_goal tool"); - } - await createGoalTool.execute( - "create-goal", - { objective: "finish the active goal" }, - undefined, - undefined, - {} as ExtensionContext, - ); + it("keeps ipython active on active-goal runtime rebuild", async () => { + const harness = await createGoalHarness(); + harness.session.handleGoalHostRequest("goal.create", { objective: "finish the active goal" }); await harness.session.reload(); - expect(harness.session.getActiveToolNames()).toEqual(["ipython", ...GOAL_TOOL_NAMES]); + expect(harness.session.getActiveToolNames()).toEqual(["ipython"]); }); it("does not reject continuation when goal error update listeners throw", async () => { - const harness = await createHarness(); - harnesses.push(harness); - const createGoalTool = harness.session.getToolDefinition("create_goal"); - if (!createGoalTool) { - throw new Error("expected create_goal tool"); - } - await createGoalTool.execute( - "create-goal", - { objective: "finish the active goal" }, - undefined, - undefined, - {} as ExtensionContext, - ); + const harness = await createGoalHarness(); + harness.session.handleGoalHostRequest("goal.create", { objective: "finish the active goal" }); harness.session.subscribe((event) => { if (event.type === "goal_update") { throw new Error("listener failed"); @@ -268,12 +382,11 @@ describe("AgentSession goals", () => { }); }); - it("does not infer completion from an assistant claim without update_goal", async () => { - const harness = await createHarness(); - harnesses.push(harness); + it("does not infer completion from an assistant claim without goal.complete", async () => { + const harness = await createGoalHarness(); harness.setResponses([ fauxAssistantMessage("Done."), - fauxAssistantMessage(fauxToolCall("update_goal", { status: "complete" }), { stopReason: "toolUse" }), + fauxAssistantMessage(fauxToolCall("ipython", COMPLETE_GOAL_CELL), { stopReason: "toolUse" }), fauxAssistantMessage("Goal complete."), ]); @@ -288,15 +401,17 @@ describe("AgentSession goals", () => { }); }); - it("lets the model create a persistent goal with create_goal", async () => { - const harness = await createHarness(); - harnesses.push(harness); + it("lets the model create a persistent goal through ipython", async () => { + const harness = await createGoalHarness(); harness.setResponses([ - fauxAssistantMessage(fauxToolCall("create_goal", { objective: "write a benchmark note" }), { - stopReason: "toolUse", - }), + fauxAssistantMessage( + fauxToolCall("ipython", { code: 'goal.create {"objective": "write a benchmark note"}' }), + { + stopReason: "toolUse", + }, + ), fauxAssistantMessage("Started the note."), - fauxAssistantMessage(fauxToolCall("update_goal", { status: "complete" }), { stopReason: "toolUse" }), + fauxAssistantMessage(fauxToolCall("ipython", COMPLETE_GOAL_CELL), { stopReason: "toolUse" }), fauxAssistantMessage("Goal complete."), ]); @@ -312,8 +427,7 @@ describe("AgentSession goals", () => { }); it("reloads goal state after tree navigation", async () => { - const harness = await createHarness(); - harnesses.push(harness); + const harness = await createGoalHarness(); harness.setResponses([fauxAssistantMessage("before goal")]); await harness.session.prompt("normal prompt"); const beforeGoalEntry = harness.sessionManager @@ -324,7 +438,7 @@ describe("AgentSession goals", () => { } harness.setResponses([ - fauxAssistantMessage(fauxToolCall("update_goal", { status: "complete" }), { stopReason: "toolUse" }), + fauxAssistantMessage(fauxToolCall("ipython", COMPLETE_GOAL_CELL), { stopReason: "toolUse" }), fauxAssistantMessage("Goal complete."), ]); await harness.session.prompt("/goal finish the task"); @@ -340,12 +454,11 @@ describe("AgentSession goals", () => { it("keeps active goals sticky across normal user prompts", async () => { const waiting = createWaitingTool(); - const harness = await createHarness({ tools: [waiting.tool] }); - harnesses.push(harness); + const harness = await createGoalHarness([waiting.tool]); harness.setResponses([ fauxAssistantMessage(fauxToolCall("wait", {}), { stopReason: "toolUse" }), fauxAssistantMessage("answered the side question"), - fauxAssistantMessage(fauxToolCall("update_goal", { status: "complete" }), { stopReason: "toolUse" }), + fauxAssistantMessage(fauxToolCall("ipython", COMPLETE_GOAL_CELL), { stopReason: "toolUse" }), fauxAssistantMessage("Goal complete."), ]); @@ -372,8 +485,7 @@ describe("AgentSession goals", () => { { command: "/goal pause", status: "paused" }, ])("removes queued goal context after $command while streaming", async ({ command, status }) => { const waiting = createWaitingTool(); - const harness = await createHarness({ tools: [waiting.tool] }); - harnesses.push(harness); + const harness = await createGoalHarness([waiting.tool]); harness.setResponses([ fauxAssistantMessage(fauxToolCall("wait", {}), { stopReason: "toolUse" }), fauxAssistantMessage("stale goal response"), @@ -395,8 +507,7 @@ describe("AgentSession goals", () => { it("pauses an active goal with /goal pause", async () => { const waiting = createWaitingTool(); - const harness = await createHarness({ tools: [waiting.tool] }); - harnesses.push(harness); + const harness = await createGoalHarness([waiting.tool]); harness.setResponses([fauxAssistantMessage(fauxToolCall("wait", {}), { stopReason: "toolUse" })]); const waitForStart = waiting.waitForStart(harness); @@ -416,8 +527,7 @@ describe("AgentSession goals", () => { it("resumes a paused goal with /goal resume", async () => { const waiting = createWaitingTool(); - const harness = await createHarness({ tools: [waiting.tool] }); - harnesses.push(harness); + const harness = await createGoalHarness([waiting.tool]); harness.setResponses([fauxAssistantMessage(fauxToolCall("wait", {}), { stopReason: "toolUse" })]); const waitForStart = waiting.waitForStart(harness); @@ -428,7 +538,7 @@ describe("AgentSession goals", () => { await promptPromise; harness.setResponses([ - fauxAssistantMessage(fauxToolCall("update_goal", { status: "complete" }), { stopReason: "toolUse" }), + fauxAssistantMessage(fauxToolCall("ipython", COMPLETE_GOAL_CELL), { stopReason: "toolUse" }), fauxAssistantMessage("Goal complete."), ]); await harness.session.prompt("/goal resume"); @@ -441,10 +551,9 @@ describe("AgentSession goals", () => { }); it("does not resume a completed goal", async () => { - const harness = await createHarness(); - harnesses.push(harness); + const harness = await createGoalHarness(); harness.setResponses([ - fauxAssistantMessage(fauxToolCall("update_goal", { status: "complete" }), { stopReason: "toolUse" }), + fauxAssistantMessage(fauxToolCall("ipython", COMPLETE_GOAL_CELL), { stopReason: "toolUse" }), fauxAssistantMessage("Goal complete."), fauxAssistantMessage("should not run"), ]); @@ -477,13 +586,12 @@ describe("AgentSession goals", () => { expect(harness.getPendingResponseCount()).toBe(1); }); - it("reports active goal elapsed time on status reads and get_goal", async () => { + it("reports active goal elapsed time on status reads and goal.get", async () => { vi.useFakeTimers(); try { vi.setSystemTime(new Date("2026-01-01T00:00:00Z")); const waiting = createWaitingTool(); - const harness = await createHarness({ tools: [waiting.tool] }); - harnesses.push(harness); + const harness = await createGoalHarness([waiting.tool]); harness.setResponses([fauxAssistantMessage(fauxToolCall("wait", {}), { stopReason: "toolUse" })]); const waitForStart = waiting.waitForStart(harness); @@ -492,13 +600,8 @@ describe("AgentSession goals", () => { vi.setSystemTime(new Date("2026-01-01T00:00:05Z")); expect(harness.session.goalState.timeUsedSeconds).toBe(5); - const getGoalTool = harness.session.getToolDefinition("get_goal"); - if (!getGoalTool) { - throw new Error("expected get_goal tool"); - } - const result = await getGoalTool.execute("get-goal", {}, undefined, undefined, {} as ExtensionContext); - const details = result.details as { goal: { timeUsedSeconds: number } | null }; - expect(details.goal?.timeUsedSeconds).toBe(5); + const response: GoalHostResponse = harness.session.handleGoalHostRequest("goal.get"); + expect(response.goal?.time_used_seconds).toBe(5); await harness.session.prompt("/goal pause"); waiting.release(); @@ -533,6 +636,33 @@ describe("AgentSession goals", () => { expect(harness.session.messages).toEqual([]); }); + it("completes a goal whose completing turn crosses the budget without a stale budget-limit steer", async () => { + const harness = await createGoalHarness(); + harness.setResponses([ + assistantWithUsage( + fauxAssistantMessage(fauxToolCall("ipython", COMPLETE_GOAL_CELL), { stopReason: "toolUse" }), + { input: 6, output: 5, totalTokens: 11 }, + ), + fauxAssistantMessage("Goal complete."), + ]); + + await harness.session.prompt("/goal --budget 10 finish the task"); + + expect(visibleAssistantTexts(harness)).toEqual(["Goal complete."]); + const contextKinds = goalContextMessages(harness).map( + (message) => (message as { details?: { kind?: string } }).details?.kind, + ); + expect(contextKinds).not.toContain("budget_limit"); + expect(harness.session.goalState).toMatchObject({ + active: false, + status: "complete", + tokenBudget: 10, + lastReason: "Goal achieved", + }); + expect(harness.session.goalState.tokensUsed).toBeGreaterThanOrEqual(10); + expect(harness.getPendingResponseCount()).toBe(0); + }); + it("marks an active goal budget_limited when token budget is reached", async () => { const harness = await createHarness(); harnesses.push(harness); @@ -626,19 +756,8 @@ describe("AgentSession goals", () => { }); it("does not continue when a terminal error reaches the continuation hook", async () => { - const harness = await createHarness(); - harnesses.push(harness); - const createGoalTool = harness.session.getToolDefinition("create_goal"); - if (!createGoalTool) { - throw new Error("expected create_goal tool"); - } - await createGoalTool.execute( - "create-goal", - { objective: "finish the active goal" }, - undefined, - undefined, - {} as ExtensionContext, - ); + const harness = await createGoalHarness(); + harness.session.handleGoalHostRequest("goal.create", { objective: "finish the active goal" }); const errorMessage = fauxAssistantMessage("", { stopReason: "error", errorMessage: "invalid_api_key" }); const continuationMessages = await harness.session.agent.getContinuationMessages?.({ @@ -672,8 +791,7 @@ describe("AgentSession goals", () => { it("lets the user abort a goal turn, prompt in between, then resume the goal", async () => { const waiting = createWaitingTool(); - const harness = await createHarness({ tools: [waiting.tool] }); - harnesses.push(harness); + const harness = await createGoalHarness([waiting.tool]); harness.setResponses([fauxAssistantMessage(fauxToolCall("wait", {}), { stopReason: "toolUse" })]); const waitForStart = waiting.waitForStart(harness); @@ -689,7 +807,7 @@ describe("AgentSession goals", () => { harness.setResponses([ fauxAssistantMessage("answered the interjection"), - fauxAssistantMessage(fauxToolCall("update_goal", { status: "complete" }), { stopReason: "toolUse" }), + fauxAssistantMessage(fauxToolCall("ipython", COMPLETE_GOAL_CELL), { stopReason: "toolUse" }), fauxAssistantMessage("Goal complete."), ]); diff --git a/packages/coding-agent/test/suite/regressions/2835-tools-allowlist-filters-extension-tools.test.ts b/packages/coding-agent/test/suite/regressions/2835-tools-allowlist-filters-extension-tools.test.ts index 76592c165..dd4c46c83 100644 --- a/packages/coding-agent/test/suite/regressions/2835-tools-allowlist-filters-extension-tools.test.ts +++ b/packages/coding-agent/test/suite/regressions/2835-tools-allowlist-filters-extension-tools.test.ts @@ -91,12 +91,4 @@ describe("regression #2835: tool allowlists filter extension tools", () => { expect(session.systemPrompt).not.toContain("dynamic_tool"); session.dispose(); }); - - it("allows goal tools only when explicitly listed", async () => { - const session = await createSession(["get_goal"]); - - expect(session.getAllTools().map((tool) => tool.name)).toEqual(["get_goal"]); - expect(session.getActiveToolNames()).toEqual(["get_goal"]); - session.dispose(); - }); }); diff --git a/packages/coding-agent/test/suite/regressions/3592-no-builtin-tools-keeps-extension-tools.test.ts b/packages/coding-agent/test/suite/regressions/3592-no-builtin-tools-keeps-extension-tools.test.ts index 7844382a3..b1e4f9b5c 100644 --- a/packages/coding-agent/test/suite/regressions/3592-no-builtin-tools-keeps-extension-tools.test.ts +++ b/packages/coding-agent/test/suite/regressions/3592-no-builtin-tools-keeps-extension-tools.test.ts @@ -78,7 +78,7 @@ describe("regression #3592: no-builtin-tools keeps extension tools enabled", () .getAllTools() .map((tool) => tool.name) .sort(), - ).toEqual(["bash", "create_goal", "dynamic_tool", "edit", "get_goal", "ipython", "update_goal"]); + ).toEqual(["bash", "dynamic_tool", "edit", "ipython"]); expect(session.getActiveToolNames()).toEqual(["dynamic_tool"]); expect(session.systemPrompt).not.toContain("- dynamic_tool: Run dynamic test behavior"); expect(session.systemPrompt).not.toContain("- ipython:"); diff --git a/packages/coding-agent/test/suite/regressions/3885-subagent-runtime-host.test.ts b/packages/coding-agent/test/suite/regressions/3885-subagent-runtime-host.test.ts index 85dd33731..f73435cce 100644 --- a/packages/coding-agent/test/suite/regressions/3885-subagent-runtime-host.test.ts +++ b/packages/coding-agent/test/suite/regressions/3885-subagent-runtime-host.test.ts @@ -94,8 +94,7 @@ describe("ENG-3885 subagent runtime host", () => { initialActiveToolNames: sessionOptions?.initialActiveToolNames, allowedToolNames: sessionOptions?.allowedToolNames, customTools: sessionOptions?.customTools, - includeGoalTools: sessionOptions?.includeGoalTools, - autoActivateGoalTools: sessionOptions?.autoActivateGoalTools, + includeGoals: sessionOptions?.includeGoals, rlmDepth: sessionOptions?.rlmDepth, rlmMaxDepth: sessionOptions?.rlmMaxDepth, rlmSessionDir: sessionOptions?.rlmSessionDir, diff --git a/prime-agent-runtime/src/rlm/__init__.py b/prime-agent-runtime/src/rlm/__init__.py index bbcc2d397..43236a4d3 100644 --- a/prime-agent-runtime/src/rlm/__init__.py +++ b/prime-agent-runtime/src/rlm/__init__.py @@ -22,6 +22,8 @@ try: except Exception: # pragma: no cover - only available in kernels get_ipython = None # type: ignore[assignment] +HOST_COMM_TARGET = "host.request" + @dataclass class TokenUsage: @@ -94,14 +96,25 @@ def _result_from_payload(payload: dict[str, Any]) -> RLMResult: ) -async def _host_request(data: dict[str, Any]) -> dict[str, Any]: +async def host_request(request_type: str, payload: dict[str, Any] | None = None) -> dict[str, Any]: + """Send a typed request to the Prime Agent host and await its reply. + + This is the kernel side of the generic host bridge: Python skills call + ``await host_request("", {...})`` and the TypeScript host dispatches + on the type. Raises RuntimeError when the host reports an error or when no + handler for the type is registered in this session. + """ + if not isinstance(request_type, str) or not request_type: + raise TypeError("request_type must be a non-empty str") + if payload is not None and not isinstance(payload, dict): + raise TypeError(f"payload must be a dict or None, got {type(payload).__name__}") if Comm is None: raise RuntimeError("Jupyter comm support is unavailable in this kernel") _install_control_comm_handlers() loop = asyncio.get_running_loop() future: asyncio.Future[dict[str, Any]] = loop.create_future() - comm = Comm(target_name="rlm.run", primary=False) + comm = Comm(target_name=HOST_COMM_TARGET, primary=False) def _on_msg(msg: dict[str, Any]) -> None: content = msg.get("content", {}) @@ -113,22 +126,32 @@ async def _host_request(data: dict[str, Any]) -> dict[str, Any]: if status == "ok": def _resolve_result() -> None: if not future.done(): - future.set_result(reply) + future.set_result({k: v for k, v in reply.items() if k != "status"}) comm.close() loop.call_soon_threadsafe(_resolve_result) return if status == "error": - message = reply.get("error") or "rlm.run failed" + message = reply.get("error") or f"host request {request_type} failed" def _resolve_error() -> None: if not future.done(): future.set_exception(RuntimeError(str(message))) comm.close() loop.call_soon_threadsafe(_resolve_error) + return + + unexpected = f"host request {request_type} returned unexpected status: {status!r}" + def _resolve_unexpected() -> None: + if not future.done(): + future.set_exception(RuntimeError(unexpected)) + comm.close() + + loop.call_soon_threadsafe(_resolve_unexpected) comm.on_msg(_on_msg) - comm.open(data=data) + # request_type goes last so a payload "type" key cannot reroute the request. + comm.open(data={**(payload or {}), "type": request_type}) return await future @@ -137,7 +160,7 @@ async def run(prompt: str, **kwargs: Any) -> RLMResult: if not isinstance(prompt, str): raise TypeError(f"prompt must be str, got {type(prompt).__name__}") _ensure_recursion_allowed() - payload = await _host_request({"type": "run", "prompt": prompt, "kwargs": kwargs}) + payload = await host_request("rlm.run", {"prompt": prompt, "kwargs": kwargs}) return _result_from_payload(payload) @@ -180,6 +203,7 @@ __all__ = [ "TokenUsage", "get_harness_state", "harness", + "host_request", "rlm", "run", ]