From a2d090a11305ce8b50200df5947730388d5d86a1 Mon Sep 17 00:00:00 2001 From: Seth Date: Fri, 4 Sep 2026 12:40:46 -0400 Subject: [PATCH 1/6] feat(coding-agent): wake agent after async bash completion --- .../.changes/async-bash-completion-message.md | 1 + packages/coding-agent/docs/rlm.md | 9 +++ .../coding-agent/src/core/agent-session.ts | 15 +++++ packages/coding-agent/src/core/messages.ts | 29 +++++++++ packages/coding-agent/src/core/prompts/rlm.ts | 2 +- packages/coding-agent/src/core/rlm-runtime.ts | 29 +++++++++ .../components/injected-prompt-message.ts | 12 +++- .../test/agent-session-recursion.test.ts | 29 ++++++++- .../test/async-bash-completion.test.ts | 47 ++++++++++++++ .../coding-agent/test/system-prompt.test.ts | 1 + prime-agent-runtime/src/rlm/bash.py | 62 +++++++++++++++++++ prime-agent-runtime/src/rlm/repl.py | 29 +++++++++ prime-agent-runtime/test/test_repl.py | 40 ++++++++++++ 13 files changed, 302 insertions(+), 3 deletions(-) create mode 100644 packages/coding-agent/.changes/async-bash-completion-message.md create mode 100644 packages/coding-agent/test/async-bash-completion.test.ts diff --git a/packages/coding-agent/.changes/async-bash-completion-message.md b/packages/coding-agent/.changes/async-bash-completion-message.md new file mode 100644 index 0000000000..c4d1286ade --- /dev/null +++ b/packages/coding-agent/.changes/async-bash-completion-message.md @@ -0,0 +1 @@ +- Added completion follow-ups for unawaited kernel `bash()` handles so the agent inspects finished async commands and continues the task. diff --git a/packages/coding-agent/docs/rlm.md b/packages/coding-agent/docs/rlm.md index 7d558073b4..2f6e3c2e23 100644 --- a/packages/coding-agent/docs/rlm.md +++ b/packages/coding-agent/docs/rlm.md @@ -48,6 +48,15 @@ result = await bash("npm run check") print(result.output) ``` +For a long command, keep the live handle and let the turn end instead of blocking: + +```python +checks = bash("npm test") +checks.pid +``` + +When an unawaited handle finishes, Prime Agent sends the agent a follow-up with its PID and exit code. The follow-up asks the agent to inspect the saved handle with `poll()`, `output()`, or `tail()` and continue the task. `await bash(...)` stays synchronous from the agent's perspective and does not send a second completion follow-up. + Each `bash()` call is its own process, while Python state, `os.chdir(...)`, and `os.environ[...]` changes persist in the kernel and apply to later `bash()` calls. Prime Agent extensions may intentionally add custom tools, but the built-in RLM design does not require a separate model tool for every capability. ### 2. Subagents are native RLM calls diff --git a/packages/coding-agent/src/core/agent-session.ts b/packages/coding-agent/src/core/agent-session.ts index be375cc048..e6022ebeb2 100644 --- a/packages/coding-agent/src/core/agent-session.ts +++ b/packages/coding-agent/src/core/agent-session.ts @@ -165,11 +165,13 @@ import { type RestoreResult, snapshotPathIn } from "./kernel/state-snapshot.js"; import type { AcpMcpServerConfig } from "./mcp/acp-mcp-types.js"; import type { McpManager } from "./mcp/mcp-manager.js"; import { + ASYNC_BASH_COMPLETION_CUSTOM_TYPE, type BashExecutionMessage, type CompactionOutcome, type CompactionOutcomeReason, type CustomMessage, convertToLlm, + createAsyncBashCompletionMessage, createCompactionOutcomeMessage, createHeartbeatPromptMessage, createRefinementOutcomeMessage, @@ -214,6 +216,7 @@ import { resolveConfigValue } from "./resolve-config-value.js"; import type { ResourceExtensionPaths, ResourceLoader } from "./resource-loader.js"; import { type CreateRlmSubagentRuntimeOptions, + createAsyncBashCompletionHostHandler, createDefaultRlmSubagentSessionName, createRlmDeleteSubagentHostHandler, createRlmFindModelsHostHandler, @@ -811,6 +814,8 @@ function injectedMessagePreviewLabel(message: CustomMessage): string | undefined switch (message.customType) { case HEARTBEAT_PROMPT_CUSTOM_TYPE: return HEARTBEAT_PROMPT_PREVIEW_LABEL; + case ASYNC_BASH_COMPLETION_CUSTOM_TYPE: + return "Async shell completed"; case GOAL_CONTEXT_CUSTOM_TYPE: return GOAL_CONTEXT_PREVIEW_LABEL; default: @@ -9346,6 +9351,16 @@ export class AgentSession { "rlm.run": createRlmRunHostHandler(async ({ prompt, kwargs, cellSourceCode }) => ({ ...(await this.runRlmChild(prompt, kwargs, cellSourceCode)), })), + "bash.completed": createAsyncBashCompletionHostHandler(async (details) => { + const message = createAsyncBashCompletionMessage(details); + await this._promptInjectedMessage(message.content, message, { + streamingBehavior: "followUp", + queueIfBusy: true, + resumeIfIdle: true, + returnAfterAccepted: true, + suppressAutonomousContinuation: true, + }); + }), "rlm.find_models": createRlmFindModelsHostHandler((query, limit) => this.findRlmModels(query, limit)), "rlm.list_subagents": createRlmListSubagentsHostHandler(() => this.listRlmSubagents()), "rlm.delete_subagent": createRlmDeleteSubagentHostHandler((target) => this.deleteRlmSubagent(target)), diff --git a/packages/coding-agent/src/core/messages.ts b/packages/coding-agent/src/core/messages.ts index 81a77117f4..c81416f279 100644 --- a/packages/coding-agent/src/core/messages.ts +++ b/packages/coding-agent/src/core/messages.ts @@ -35,6 +35,7 @@ export const COMPACTION_OUTCOME_CUSTOM_TYPE = "compaction_outcome"; export const REFINEMENT_OUTCOME_CUSTOM_TYPE = "refinement_outcome"; export const RLM_CHILD_FAILURE_CUSTOM_TYPE = "rlm_child_failure"; export const RLM_CHILD_TERMINAL_NOTICE_CUSTOM_TYPE = "rlm_child_terminal_notice"; +export const ASYNC_BASH_COMPLETION_CUSTOM_TYPE = "async_bash_completion"; export interface SessionSlashCommandDetails { command: SessionSlashCommand; @@ -109,6 +110,34 @@ export type RlmChildTerminalNoticeDetails = lastAssistantTextPreview?: string; }; +export interface AsyncBashCompletionDetails { + pid: number; + command: string; + exitCode: number; +} + +export interface AsyncBashCompletionMessage extends CustomMessage { + customType: typeof ASYNC_BASH_COMPLETION_CUSTOM_TYPE; + content: string; +} + +export function createAsyncBashCompletionMessage( + details: AsyncBashCompletionDetails, + timestamp = Date.now(), +): AsyncBashCompletionMessage { + return { + role: "custom", + customType: ASYNC_BASH_COMPLETION_CUSTOM_TYPE, + content: `Async bash command completed (pid ${details.pid}, exit code ${details.exitCode}). +Command: ${JSON.stringify(details.command)} + +Inspect the saved BashHandle with .poll(), .output(), or .tail(), then continue the task.`, + display: true, + details, + timestamp, + }; +} + export function createRlmChildFailureMessage( details: RlmChildFailureDetails, timestamp = Date.now(), diff --git a/packages/coding-agent/src/core/prompts/rlm.ts b/packages/coding-agent/src/core/prompts/rlm.ts index ef610268e6..a9c13946da 100644 --- a/packages/coding-agent/src/core/prompts/rlm.ts +++ b/packages/coding-agent/src/core/prompts/rlm.ts @@ -12,7 +12,7 @@ export interface RlmPromptOptions { } const LONG_RUNNING_WORK_PROMPT = [ - "For slow or independently completing work, use a nonblocking control loop: start the work, record its handle or output location, then end your turn. Read the result on a later turn or when a reply arrives.", + "For slow or independently completing work, use a nonblocking control loop: start the work, record its handle or output location, then end your turn. An unawaited `bash()` handle sends a completion follow-up; when it arrives, inspect the saved handle and continue.", "When delegation is available and useful, assign independent substantive tasks to separate workers. Start independent workers without waiting for each one sequentially, and let them run in parallel.", "Do not keep the turn open by polling with `time.sleep()` or shell `sleep`, and do not replace polling with a long blocking `await`. Await only the short operation needed to start work or inspect a result that is already available; otherwise end the turn.", ].join("\n"); diff --git a/packages/coding-agent/src/core/rlm-runtime.ts b/packages/coding-agent/src/core/rlm-runtime.ts index 470224cc53..c7fd2500ce 100644 --- a/packages/coding-agent/src/core/rlm-runtime.ts +++ b/packages/coding-agent/src/core/rlm-runtime.ts @@ -51,6 +51,14 @@ export interface RlmFindModelsResult { } export type RlmRunHandler = (request: RlmRunRequest) => Promise>; + +export interface AsyncBashCompletionRequest { + pid: number; + command: string; + exitCode: number; +} + +export type AsyncBashCompletionHandler = (request: AsyncBashCompletionRequest) => void | Promise; export type RlmListSubagentsHandler = () => RlmListSubagentsResult | Promise; export type RlmDeleteSubagentHandler = (target: string) => Promise; export type RlmFindModelsHandler = (query: string, limit: number) => RlmFindModelsResult | Promise; @@ -178,6 +186,27 @@ export function createRlmRunHostHandler(handler: RlmRunHandler): HostRequestHand }; } +/** Adapt detached kernel bash completions into a validated host notification. */ +export function createAsyncBashCompletionHostHandler(handler: AsyncBashCompletionHandler): HostRequestHandler { + return async (payload) => { + if (!Number.isInteger(payload.pid) || (payload.pid as number) <= 0) { + throw new Error("bash.completed pid must be a positive integer"); + } + if (typeof payload.command !== "string" || !payload.command) { + throw new Error("bash.completed command must be a non-empty string"); + } + if (!Number.isInteger(payload.exitCode)) { + throw new Error("bash.completed exitCode must be an integer"); + } + await handler({ + pid: payload.pid as number, + command: payload.command, + exitCode: payload.exitCode as number, + }); + return {}; + }; +} + /** Search a bounded authenticated model catalog without adding it to the system prompt. */ export function createRlmFindModelsHostHandler(handler: RlmFindModelsHandler): HostRequestHandler { return async (payload) => { diff --git a/packages/coding-agent/src/modes/interactive/components/injected-prompt-message.ts b/packages/coding-agent/src/modes/interactive/components/injected-prompt-message.ts index 6893496870..93f0758023 100644 --- a/packages/coding-agent/src/modes/interactive/components/injected-prompt-message.ts +++ b/packages/coding-agent/src/modes/interactive/components/injected-prompt-message.ts @@ -10,6 +10,8 @@ import { } from "@earendil-works/pi-tui"; import { GOAL_CONTEXT_CUSTOM_TYPE, type GoalContextDetails } from "../../../core/goals.js"; import { + ASYNC_BASH_COMPLETION_CUSTOM_TYPE, + type AsyncBashCompletionDetails, type CustomMessage, HEARTBEAT_PROMPT_CUSTOM_TYPE, type HeartbeatPromptDetails, @@ -24,6 +26,7 @@ import { getMarkdownTheme, theme } from "../theme/theme.js"; import { expandCollapseHint } from "./keybinding-hints.js"; type InjectedPromptDetails = + | AsyncBashCompletionDetails | GoalContextDetails | HeartbeatPromptDetails | IpythonStateRestoredDetails @@ -34,7 +37,8 @@ type InjectedPromptMessage = CustomMessage; export function isInjectedPromptMessage(message: AgentMessage): message is InjectedPromptMessage { return ( message.role === "custom" && - (message.customType === HEARTBEAT_PROMPT_CUSTOM_TYPE || + (message.customType === ASYNC_BASH_COMPLETION_CUSTOM_TYPE || + message.customType === HEARTBEAT_PROMPT_CUSTOM_TYPE || message.customType === GOAL_CONTEXT_CUSTOM_TYPE || message.customType === IPYTHON_STATE_RESTORED_CUSTOM_TYPE || message.customType === RLM_CHILD_FAILURE_CUSTOM_TYPE || @@ -125,6 +129,12 @@ export class InjectedPromptMessageComponent extends Container { if (this.message.customType === HEARTBEAT_PROMPT_CUSTOM_TYPE) { return this.heartbeatHeaderText(); } + if (this.message.customType === ASYNC_BASH_COMPLETION_CUSTOM_TYPE) { + const details = this.message.details as AsyncBashCompletionDetails | undefined; + const status = details ? ` ยท exit ${details.exitCode}` : ""; + const hint = this.expanded ? "" : ` ${expandCollapseHint("app.tools.expand", false)}`; + return theme.fg("muted", "Async shell completed") + theme.fg("dim", status + hint); + } if (this.message.customType === IPYTHON_STATE_RESTORED_CUSTOM_TYPE) { const details = this.message.details as IpythonStateRestoredDetails | undefined; const label = details?.restored === false ? "Started fresh Python kernel" : "Restored Python kernel state"; diff --git a/packages/coding-agent/test/agent-session-recursion.test.ts b/packages/coding-agent/test/agent-session-recursion.test.ts index 2dad4fcd9d..ba75944913 100644 --- a/packages/coding-agent/test/agent-session-recursion.test.ts +++ b/packages/coding-agent/test/agent-session-recursion.test.ts @@ -22,7 +22,7 @@ import { AgentSession, type RlmChildAgentSnapshot } from "../src/core/agent-sess import { AuthStorage } from "../src/core/auth-storage.js"; import type { LoadExtensionsResult } from "../src/core/extensions/index.js"; import { type HostRequestHandlers, ReplKernelManager } from "../src/core/kernel/index.js"; -import { convertToLlm } from "../src/core/messages.js"; +import { ASYNC_BASH_COMPLETION_CUSTOM_TYPE, convertToLlm } from "../src/core/messages.js"; import { ModelRegistry } from "../src/core/model-registry.js"; import { createDefaultRlmSubagentSessionName, @@ -820,6 +820,33 @@ describe("AgentSession rlm recursion", () => { expect(doneUpdate?.toolUseCount).toBeUndefined(); }); + it("wakes the agent with a follow-up when a detached bash handle completes", async () => { + const prompts: string[] = []; + const root = createSession({ + streamFn: (_model, context) => { + prompts.push(userText(context)); + return streamAnswer("checked shell result"); + }, + }); + const handlers = (root as unknown as InspectableRlmSession)._createKernelHostHandlers(); + const completed = handlers["bash.completed"]; + if (!completed) throw new Error("Missing bash.completed host handler"); + + await expect(completed({ pid: 42, command: "npm test", exitCode: 1 })).resolves.toEqual({}); + await root.waitForIdle(); + + expect(prompts).toEqual([ + expect.stringContaining("Inspect the saved BashHandle with .poll(), .output(), or .tail()"), + ]); + expect(root.messages).toContainEqual( + expect.objectContaining({ + role: "custom", + customType: ASYNC_BASH_COMPLETION_CUSTOM_TYPE, + details: { pid: 42, command: "npm test", exitCode: 1 }, + }), + ); + }); + it("marks an in-cell roled send to the parent as replied", async () => { const sendAgentMessage = vi.fn(async () => ({ id: "agentmsg-reply", diff --git a/packages/coding-agent/test/async-bash-completion.test.ts b/packages/coding-agent/test/async-bash-completion.test.ts new file mode 100644 index 0000000000..b90aab677b --- /dev/null +++ b/packages/coding-agent/test/async-bash-completion.test.ts @@ -0,0 +1,47 @@ +import { describe, expect, it, vi } from "vitest"; +import { + ASYNC_BASH_COMPLETION_CUSTOM_TYPE, + convertToLlm, + createAsyncBashCompletionMessage, +} from "../src/core/messages.js"; +import { createAsyncBashCompletionHostHandler } from "../src/core/rlm-runtime.js"; + +describe("async bash completion", () => { + it("creates a model-visible instruction to inspect the saved handle", () => { + const message = createAsyncBashCompletionMessage({ + pid: 42, + command: "npm test", + exitCode: 1, + }); + + expect(message.customType).toBe(ASYNC_BASH_COMPLETION_CUSTOM_TYPE); + expect(message.content).toContain("pid 42, exit code 1"); + expect(message.content).toContain("npm test"); + expect(message.content).toContain(".poll(), .output(), or .tail()"); + expect(convertToLlm([message])).toEqual([ + { + role: "user", + content: [{ type: "text", text: message.content }], + timestamp: message.timestamp, + }, + ]); + }); + + it("validates and forwards kernel completion payloads", async () => { + const completion = vi.fn(); + const handler = createAsyncBashCompletionHostHandler(completion); + const payload = { pid: 42, command: "npm test", exitCode: 0 }; + + await expect(handler(payload)).resolves.toEqual({}); + expect(completion).toHaveBeenCalledWith(payload); + }); + + it.each([ + [{ pid: 0, command: "ok", exitCode: 0 }, "positive integer"], + [{ pid: 1, command: "", exitCode: 0 }, "non-empty string"], + [{ pid: 1, command: "ok", exitCode: 0.5 }, "exitCode"], + ])("rejects an invalid payload %#", async (payload, error) => { + const handler = createAsyncBashCompletionHostHandler(() => undefined); + await expect(handler(payload)).rejects.toThrow(error); + }); +}); diff --git a/packages/coding-agent/test/system-prompt.test.ts b/packages/coding-agent/test/system-prompt.test.ts index ca39619eb1..a0b79b6f0d 100644 --- a/packages/coding-agent/test/system-prompt.test.ts +++ b/packages/coding-agent/test/system-prompt.test.ts @@ -168,6 +168,7 @@ describe("buildRlmPrompt", () => { }); expect(prompt).toContain("Use `bash()` to invoke programs, not to write shell programs"); + expect(prompt).toContain("An unawaited `bash()` handle sends a completion follow-up"); }); test("documents preferring Python for reading and searching files when ipython is active", () => { diff --git a/prime-agent-runtime/src/rlm/bash.py b/prime-agent-runtime/src/rlm/bash.py index 5795b6ae72..f3499b38c4 100644 --- a/prime-agent-runtime/src/rlm/bash.py +++ b/prime-agent-runtime/src/rlm/bash.py @@ -43,6 +43,7 @@ # wait for a confirmed group exit before CancelledError propagates. _CANCEL_TERM_GRACE = 0.5 _CANCEL_KILL_WAIT = 2.0 +_COMPLETION_NOTICE_COMMAND_CAP = 1000 _live_handles: set["BashHandle"] = set() _live_lock = threading.Lock() @@ -50,6 +51,24 @@ _hook_lock = threading.Lock() +def _current_cell_finished_event() -> asyncio.Event | None: + """Get the creating REPL cell's completion barrier without coupling standalone use to repl.""" + try: + from . import repl + + if repl.is_active(): + return repl.current_cell_finished_event() + except (ImportError, RuntimeError): + pass + return None + + +def _consume_notice_task(task: asyncio.Task[None]) -> None: + """Retrieve detached notifier failures so they never become loop warnings.""" + if not task.cancelled(): + task.exception() + + @dataclass(frozen=True) class BashResult: exit_code: int @@ -117,6 +136,8 @@ class BashHandle: def __init__(self, command: str) -> None: self.command = command + self._creating_cell_finished = _current_cell_finished_event() + self._await_started = False self._buffer = _BoundedBuffer() self._done = threading.Event() self._eof = threading.Event() @@ -235,6 +256,7 @@ def __init__(self, command: str) -> None: threading.Thread(target=self._pump, daemon=True).start() threading.Thread(target=self._report, daemon=True).start() threading.Thread(target=self._watch, daemon=True).start() + self._schedule_background_completion_notice() @property def pid(self) -> int: @@ -515,6 +537,45 @@ def _add_done_callback(self, callback: Callable[[], None]) -> None: return callback() + def _schedule_background_completion_notice(self) -> None: + cell_finished = self._creating_cell_finished + if cell_finished is None: + return + try: + loop = asyncio.get_running_loop() + except RuntimeError: + return + task = loop.create_task(self._notify_background_completion(cell_finished)) + task.add_done_callback(_consume_notice_task) + + async def _notify_background_completion(self, cell_finished: asyncio.Event) -> None: + result = await self._wait() + # The cell may do other work before awaiting this handle. Do not classify + # it as detached until that whole cell has crossed its completion barrier. + await cell_finished.wait() + if self._await_started: + return + try: + from . import repl + + if not repl.is_active(): + return + command = self.command + if len(command) > _COMPLETION_NOTICE_COMMAND_CAP: + command = command[:_COMPLETION_NOTICE_COMMAND_CAP] + "\n... [command truncated]" + await repl.host_request( + { + "type": "bash.completed", + "pid": self._pid, + "command": command, + "exitCode": result.exit_code, + } + ) + except (OSError, RuntimeError): + # Standalone runtimes have no host handler, and teardown can close + # the bridge while a process is finishing. Shell results stay usable. + return + async def _wait(self) -> BashResult: # Asyncio-native wakeup: no executor thread is parked for the command's # duration, so many concurrent awaits cannot exhaust the default pool. @@ -639,6 +700,7 @@ def __await__(self) -> Generator[Any, None, BashResult]: # A handle awaited before any other API use is a one-shot command tied # to the await (kill-on-cancel); touching the handle API first marks it # as a deliberate background handle whose awaits only wait. + self._await_started = True if self._released: return self._wait().__await__() self._released = True diff --git a/prime-agent-runtime/src/rlm/repl.py b/prime-agent-runtime/src/rlm/repl.py index ea04ca1766..51317e8f94 100644 --- a/prime-agent-runtime/src/rlm/repl.py +++ b/prime-agent-runtime/src/rlm/repl.py @@ -52,6 +52,10 @@ _active: dict[str, Any] = {"task": None, "rid": None, "interrupted": False} _cell_counter = 0 _pending_host: dict[str, "asyncio.Future[dict[str, Any]]"] = {} +# Bash handles use this per-cell barrier to distinguish a detached command from +# one that the creating cell eventually awaits. Entries live only until the +# request finishes; late detached tasks receive an already-set event. +_cell_finished_events: dict[str, asyncio.Event] = {} # Set on the loop thread once stdin hits EOF or a shutdown request arrives; no # host reply can arrive after that, so waiting (and future) host_request calls fail. _host_closed = False @@ -98,6 +102,28 @@ def is_active() -> bool: return _protocol_fd >= 0 +def current_cell_finished_event() -> asyncio.Event | None: + """Return a barrier that opens when the calling cell has fully finished. + + Detached asyncio tasks inherit the creating cell id. If such a task starts + after that cell finished, return an already-open barrier instead of retaining + completed request ids indefinitely. + """ + cell_id = _current_cell.get() + if cell_id is None or _loop is None: + return None + with _interrupt_lock: + if cell_id not in _inflight: + event = asyncio.Event() + event.set() + return event + event = _cell_finished_events.get(cell_id) + if event is None: + event = asyncio.Event() + _cell_finished_events[cell_id] = event + return event + + async def host_request(data: dict[str, Any]) -> dict[str, Any]: """Send one typed request to the host and await its raw reply dict.""" if _loop is None: @@ -406,6 +432,9 @@ def _consume_handoff_interrupt() -> bool: def _finish_locked(rid: str) -> None: """Drop a finished request; a parked untargeted interrupt survives while others are inflight.""" global _finishing_rid, _handoff_interrupted, _sigint_target + cell_finished = _cell_finished_events.pop(rid, None) + if cell_finished is not None: + cell_finished.set() if _finishing_rid == rid: # An unconsumed handoff interrupt dies with its request (state requests # have no cancellable post-run work); it must never hit the next request. diff --git a/prime-agent-runtime/test/test_repl.py b/prime-agent-runtime/test/test_repl.py index 26b7aca89d..212e2b79e0 100644 --- a/prime-agent-runtime/test/test_repl.py +++ b/prime-agent-runtime/test/test_repl.py @@ -713,6 +713,46 @@ def test_bash_integration(self): time.sleep(0.05) self.fail(f"bash child {pid} survived runtime shutdown") + def test_async_bash_completion_notifies_only_after_an_unawaited_creating_cell(self): + direct = self.repl.execute( + "bash-direct", "from rlm import bash\n(await bash('printf direct')).exit_code" + ) + self.assertIsNone(one(direct, "host_request")) + + eventual = self.repl.execute( + "bash-eventual", + "import asyncio\nhandle = bash('printf eventual')\nawait asyncio.sleep(0.1)\n(await handle).exit_code", + ) + self.assertIsNone(one(eventual, "host_request")) + + started = self.repl.execute( + "bash-detached", "handle = bash('sleep 0.05; printf detached')\nhandle.pid" + ) + pid = int(one(started, "result")["text"]) + request = one(started, "host_request") + while request is None: + event = self.repl.read_event() + if event.get("event") == "host_request": + request = event + self.assertEqual( + request["data"], + { + "type": "bash.completed", + "pid": pid, + "command": "sleep 0.05; printf detached", + "exitCode": 0, + }, + ) + self.repl.send( + { + "type": "host_reply", + "id": request["id"], + "data": {"status": "ok", "result": {}}, + } + ) + inspected = self.repl.execute("bash-inspect", "handle.poll().output") + self.assertIn("detached", one(inspected, "result")["text"]) + def test_protocol_framing_under_noise(self): setup = "\n".join( [ From 305f88fa995789c4e9326d8afc556585d17a5e82 Mon Sep 17 00:00:00 2001 From: Seth Date: Sat, 5 Sep 2026 19:47:03 -0400 Subject: [PATCH 2/6] fix(runtime): harden async bash completion attribution --- .../.changes/async-bash-completion-message.md | 2 +- packages/coding-agent/docs/rlm.md | 2 +- packages/coding-agent/src/core/prompts/rlm.ts | 2 +- .../coding-agent/test/system-prompt.test.ts | 2 +- prime-agent-runtime/src/rlm/bash.py | 21 ++++-- prime-agent-runtime/src/rlm/repl.py | 63 ++++++++++------- prime-agent-runtime/test/test_repl.py | 68 +++++++++++++++++++ 7 files changed, 126 insertions(+), 34 deletions(-) diff --git a/packages/coding-agent/.changes/async-bash-completion-message.md b/packages/coding-agent/.changes/async-bash-completion-message.md index c4d1286ade..37671554ab 100644 --- a/packages/coding-agent/.changes/async-bash-completion-message.md +++ b/packages/coding-agent/.changes/async-bash-completion-message.md @@ -1 +1 @@ -- Added completion follow-ups for unawaited kernel `bash()` handles so the agent inspects finished async commands and continues the task. +- Added completion follow-ups for kernel `bash()` handles left running beyond their creating cell so the agent inspects finished async commands and continues the task. diff --git a/packages/coding-agent/docs/rlm.md b/packages/coding-agent/docs/rlm.md index 2f6e3c2e23..1dbebe6d83 100644 --- a/packages/coding-agent/docs/rlm.md +++ b/packages/coding-agent/docs/rlm.md @@ -55,7 +55,7 @@ checks = bash("npm test") checks.pid ``` -When an unawaited handle finishes, Prime Agent sends the agent a follow-up with its PID and exit code. The follow-up asks the agent to inspect the saved handle with `poll()`, `output()`, or `tail()` and continue the task. `await bash(...)` stays synchronous from the agent's perspective and does not send a second completion follow-up. +When a handle left running beyond its creating cell finishes, Prime Agent sends the agent a follow-up with its PID and exit code. The follow-up asks the agent to inspect the saved handle with `poll()`, `output()`, or `tail()` and continue the task. `await bash(...)` stays synchronous from the agent's perspective and does not send a second completion follow-up. Each `bash()` call is its own process, while Python state, `os.chdir(...)`, and `os.environ[...]` changes persist in the kernel and apply to later `bash()` calls. Prime Agent extensions may intentionally add custom tools, but the built-in RLM design does not require a separate model tool for every capability. diff --git a/packages/coding-agent/src/core/prompts/rlm.ts b/packages/coding-agent/src/core/prompts/rlm.ts index a9c13946da..d325fb0221 100644 --- a/packages/coding-agent/src/core/prompts/rlm.ts +++ b/packages/coding-agent/src/core/prompts/rlm.ts @@ -12,7 +12,7 @@ export interface RlmPromptOptions { } const LONG_RUNNING_WORK_PROMPT = [ - "For slow or independently completing work, use a nonblocking control loop: start the work, record its handle or output location, then end your turn. An unawaited `bash()` handle sends a completion follow-up; when it arrives, inspect the saved handle and continue.", + "For slow or independently completing work, use a nonblocking control loop: start the work, record its handle or output location, then end your turn. A `bash()` handle left running beyond its creating cell sends a completion follow-up; when it arrives, inspect the saved handle and continue.", "When delegation is available and useful, assign independent substantive tasks to separate workers. Start independent workers without waiting for each one sequentially, and let them run in parallel.", "Do not keep the turn open by polling with `time.sleep()` or shell `sleep`, and do not replace polling with a long blocking `await`. Await only the short operation needed to start work or inspect a result that is already available; otherwise end the turn.", ].join("\n"); diff --git a/packages/coding-agent/test/system-prompt.test.ts b/packages/coding-agent/test/system-prompt.test.ts index a0b79b6f0d..41b038c672 100644 --- a/packages/coding-agent/test/system-prompt.test.ts +++ b/packages/coding-agent/test/system-prompt.test.ts @@ -168,7 +168,7 @@ describe("buildRlmPrompt", () => { }); expect(prompt).toContain("Use `bash()` to invoke programs, not to write shell programs"); - expect(prompt).toContain("An unawaited `bash()` handle sends a completion follow-up"); + expect(prompt).toContain("A `bash()` handle left running beyond its creating cell sends a completion follow-up"); }); test("documents preferring Python for reading and searching files when ipython is active", () => { diff --git a/prime-agent-runtime/src/rlm/bash.py b/prime-agent-runtime/src/rlm/bash.py index f3499b38c4..7e1583545e 100644 --- a/prime-agent-runtime/src/rlm/bash.py +++ b/prime-agent-runtime/src/rlm/bash.py @@ -51,13 +51,13 @@ _hook_lock = threading.Lock() -def _current_cell_finished_event() -> asyncio.Event | None: - """Get the creating REPL cell's completion barrier without coupling standalone use to repl.""" +def _current_cell_completion_context() -> tuple[asyncio.Event, asyncio.Task[Any] | None] | None: + """Get the creating REPL cell's lifecycle without coupling standalone use to repl.""" try: from . import repl if repl.is_active(): - return repl.current_cell_finished_event() + return repl.current_cell_completion_context() except (ImportError, RuntimeError): pass return None @@ -136,8 +136,10 @@ class BashHandle: def __init__(self, command: str) -> None: self.command = command - self._creating_cell_finished = _current_cell_finished_event() - self._await_started = False + completion_context = _current_cell_completion_context() + self._creating_cell_finished = completion_context[0] if completion_context else None + self._creating_cell_task = completion_context[1] if completion_context else None + self._awaited_by_creating_cell = False self._buffer = _BoundedBuffer() self._done = threading.Event() self._eof = threading.Event() @@ -553,7 +555,7 @@ async def _notify_background_completion(self, cell_finished: asyncio.Event) -> N # The cell may do other work before awaiting this handle. Do not classify # it as detached until that whole cell has crossed its completion barrier. await cell_finished.wait() - if self._await_started: + if self._awaited_by_creating_cell: return try: from . import repl @@ -700,7 +702,12 @@ def __await__(self) -> Generator[Any, None, BashResult]: # A handle awaited before any other API use is a one-shot command tied # to the await (kill-on-cancel); touching the handle API first marks it # as a deliberate background handle whose awaits only wait. - self._await_started = True + try: + current_task = asyncio.current_task() + except RuntimeError: + current_task = None + if current_task is not None and current_task is self._creating_cell_task: + self._awaited_by_creating_cell = True if self._released: return self._wait().__await__() self._released = True diff --git a/prime-agent-runtime/src/rlm/repl.py b/prime-agent-runtime/src/rlm/repl.py index 51317e8f94..063f957750 100644 --- a/prime-agent-runtime/src/rlm/repl.py +++ b/prime-agent-runtime/src/rlm/repl.py @@ -49,13 +49,17 @@ # detached task spawned by a cell keeps writing under that cell's id after # the cell finishes. Threads start with a fresh context and emit id null. _current_cell: contextvars.ContextVar[str | None] = contextvars.ContextVar("_current_cell", default=None) +_current_cell_generation: contextvars.ContextVar[int | None] = contextvars.ContextVar( + "_current_cell_generation", default=None +) _active: dict[str, Any] = {"task": None, "rid": None, "interrupted": False} _cell_counter = 0 _pending_host: dict[str, "asyncio.Future[dict[str, Any]]"] = {} -# Bash handles use this per-cell barrier to distinguish a detached command from -# one that the creating cell eventually awaits. Entries live only until the -# request finishes; late detached tasks receive an already-set event. -_cell_finished_events: dict[str, asyncio.Event] = {} +# Request ids may be reused. Bash completion barriers therefore use the +# runtime's monotonic cell generation rather than the caller-supplied id. +_active_cell_generations: set[int] = set() +_cell_owner_tasks: dict[int, asyncio.Task[Any]] = {} +_cell_finished_events: dict[int, asyncio.Event] = {} # Set on the loop thread once stdin hits EOF or a shutdown request arrives; no # host reply can arrive after that, so waiting (and future) host_request calls fail. _host_closed = False @@ -102,26 +106,25 @@ def is_active() -> bool: return _protocol_fd >= 0 -def current_cell_finished_event() -> asyncio.Event | None: - """Return a barrier that opens when the calling cell has fully finished. +def current_cell_completion_context() -> tuple[asyncio.Event, asyncio.Task[Any] | None] | None: + """Return the calling cell's completion barrier and owning execution task. - Detached asyncio tasks inherit the creating cell id. If such a task starts - after that cell finished, return an already-open barrier instead of retaining - completed request ids indefinitely. + Detached asyncio tasks inherit the internal cell generation. If such a task + runs after that cell finished, return an already-open barrier and no owner. """ - cell_id = _current_cell.get() - if cell_id is None or _loop is None: + cell_generation = _current_cell_generation.get() + if cell_generation is None or _loop is None: return None with _interrupt_lock: - if cell_id not in _inflight: + if cell_generation not in _active_cell_generations: event = asyncio.Event() event.set() - return event - event = _cell_finished_events.get(cell_id) + return event, None + event = _cell_finished_events.get(cell_generation) if event is None: event = asyncio.Event() - _cell_finished_events[cell_id] = event - return event + _cell_finished_events[cell_generation] = event + return event, _cell_owner_tasks.get(cell_generation) async def host_request(data: dict[str, Any]) -> dict[str, Any]: @@ -432,9 +435,6 @@ def _consume_handoff_interrupt() -> bool: def _finish_locked(rid: str) -> None: """Drop a finished request; a parked untargeted interrupt survives while others are inflight.""" global _finishing_rid, _handoff_interrupted, _sigint_target - cell_finished = _cell_finished_events.pop(rid, None) - if cell_finished is not None: - cell_finished.set() if _finishing_rid == rid: # An unconsumed handoff interrupt dies with its request (state requests # have no cancellable post-run work); it must never hit the next request. @@ -455,6 +455,15 @@ def _finish_request(rid: str) -> None: _finish_locked(rid) +def _finish_cell_generation(cell_generation: int) -> None: + with _interrupt_lock: + _active_cell_generations.discard(cell_generation) + _cell_owner_tasks.pop(cell_generation, None) + cell_finished = _cell_finished_events.pop(cell_generation, None) + if cell_finished is not None: + cell_finished.set() + + _RUNTIME_FILE = __file__ @@ -565,14 +574,20 @@ async def _handle_execute(req: dict[str, Any], ns: dict[str, Any]) -> None: global _cell_counter cell_id = req["id"] _cell_counter += 1 - filename = f"" + cell_generation = _cell_counter + filename = f"" + with _interrupt_lock: + _active_cell_generations.add(cell_generation) # The cell task (created below) copies this context, so writes made from - # the cell and from asyncio tasks it spawns carry this cell's id. - token = _current_cell.set(cell_id) + # the cell and from asyncio tasks it spawns carry this cell's identity. + cell_token = _current_cell.set(cell_id) + generation_token = _current_cell_generation.set(cell_generation) try: codes, has_trailing = _compile_cell(req["code"], filename) assert _loop is not None task = _loop.create_task(_run_codes(codes, ns)) + with _interrupt_lock: + _cell_owner_tasks[cell_generation] = task status, value, error = await _run_guarded(task, cell_id) result_text: str | None = None try: @@ -597,7 +612,9 @@ async def _handle_execute(req: dict[str, Any], ns: dict[str, Any]) -> None: _send(error) _send({"event": "done", "id": cell_id, "status": status}) finally: - _current_cell.reset(token) + _finish_cell_generation(cell_generation) + _current_cell_generation.reset(generation_token) + _current_cell.reset(cell_token) def _drain_output() -> None: diff --git a/prime-agent-runtime/test/test_repl.py b/prime-agent-runtime/test/test_repl.py index 212e2b79e0..141d72f46d 100644 --- a/prime-agent-runtime/test/test_repl.py +++ b/prime-agent-runtime/test/test_repl.py @@ -753,6 +753,74 @@ def test_async_bash_completion_notifies_only_after_an_unawaited_creating_cell(se inspected = self.repl.execute("bash-inspect", "handle.poll().output") self.assertIn("detached", one(inspected, "result")["text"]) + def test_background_task_await_does_not_suppress_bash_completion(self): + code = "\n".join( + [ + "from rlm import bash", + "import asyncio", + "task_handle = bash('sleep 0.05; printf background-waiter')", + "async def consume():", + " globals()['task_result'] = await task_handle", + "waiter = asyncio.create_task(consume())", + "task_handle.pid", + ] + ) + started = self.repl.execute("bash-task-waiter", code) + pid = int(one(started, "result")["text"]) + request = one(started, "host_request") + while request is None: + event = self.repl.read_event() + if event.get("event") == "host_request": + request = event + self.assertEqual(request["data"]["type"], "bash.completed") + self.assertEqual(request["data"]["pid"], pid) + self.repl.send( + { + "type": "host_reply", + "id": request["id"], + "data": {"status": "ok", "result": {}}, + } + ) + + def test_reused_request_id_does_not_capture_old_cell_bash_completion(self): + setup = "\n".join( + [ + "from rlm import bash", + "import asyncio", + "reuse_gate = asyncio.Event()", + "async def launch_after_cell():", + " await reuse_gate.wait()", + " globals()['reused_handle'] = bash('printf reused-id')", + "asyncio.create_task(launch_after_cell())", + ] + ) + self.repl.execute("reused-cell-id", setup) + + self.repl.send( + { + "type": "execute", + "id": "reused-cell-id", + "code": "reuse_gate.set()\nawait asyncio.sleep(0.2)", + } + ) + request = None + while True: + event = self.repl.read_event() + if event.get("event") == "host_request": + request = event + break + self.assertFalse(event.get("event") == "done" and event.get("id") == "reused-cell-id") + self.assertEqual(request["data"]["type"], "bash.completed") + self.assertEqual(request["data"]["command"], "printf reused-id") + self.repl.send( + { + "type": "host_reply", + "id": request["id"], + "data": {"status": "ok", "result": {}}, + } + ) + self.assertEqual(one(self.repl.until_done("reused-cell-id"), "done")["status"], "ok") + def test_protocol_framing_under_noise(self): setup = "\n".join( [ From ffb28fc95badec99cbe5f75cb927e1f332633dd3 Mon Sep 17 00:00:00 2001 From: Seth Date: Sat, 5 Sep 2026 19:57:26 -0400 Subject: [PATCH 3/6] fix(runtime): suppress wrapped bash await notices --- prime-agent-runtime/src/rlm/bash.py | 89 ++++++++++++++++++++++++++- prime-agent-runtime/test/test_repl.py | 27 ++++++++ 2 files changed, 115 insertions(+), 1 deletion(-) diff --git a/prime-agent-runtime/src/rlm/bash.py b/prime-agent-runtime/src/rlm/bash.py index 7e1583545e..536446782a 100644 --- a/prime-agent-runtime/src/rlm/bash.py +++ b/prime-agent-runtime/src/rlm/bash.py @@ -4,6 +4,7 @@ import asyncio import atexit +import functools import json import os import secrets @@ -69,6 +70,92 @@ def _consume_notice_task(task: asyncio.Task[None]) -> None: task.exception() +def _referenced_futures( + value: Any, *, depth: int = 0, seen: set[int] | None = None +) -> set[asyncio.Future[Any]]: + """Find futures captured by asyncio's small internal completion callbacks.""" + if isinstance(value, asyncio.Future): + return {value} + if depth >= 4: + return set() + if seen is None: + seen = set() + identity = id(value) + if identity in seen: + return set() + seen.add(identity) + + nested: list[Any] = [] + if isinstance(value, functools.partial): + nested.extend((value.func, value.args, value.keywords)) + elif isinstance(value, dict): + nested.extend(value.keys()) + nested.extend(value.values()) + elif isinstance(value, (tuple, list, set, frozenset)): + nested.extend(value) + else: + closure = getattr(value, "__closure__", None) + if closure: + for cell in closure: + try: + nested.append(cell.cell_contents) + except ValueError: + pass + bound_self = getattr(value, "__self__", None) + if bound_self is not None: + nested.append(bound_self) + + futures: set[asyncio.Future[Any]] = set() + for item in nested: + futures.update(_referenced_futures(item, depth=depth + 1, seen=seen)) + return futures + + +def _future_contains(root: asyncio.Future[Any], target: asyncio.Future[Any], seen: set[int]) -> bool: + """Follow aggregate children, such as asyncio.gather's private future list.""" + if root is target: + return True + if id(root) in seen: + return False + seen.add(id(root)) + children = getattr(root, "_children", None) or () + return any( + isinstance(child, asyncio.Future) and _future_contains(child, target, seen) for child in children + ) + + +def _completion_reaches(start: asyncio.Future[Any], target: asyncio.Future[Any]) -> bool: + """Follow callback-captured futures from an inner await to its wrapper future.""" + pending = [start] + seen: set[int] = set() + while pending: + future = pending.pop() + if future is target: + return True + if id(future) in seen: + continue + seen.add(id(future)) + callbacks = getattr(future, "_callbacks", None) or () + for entry in callbacks: + callback = entry[0] if isinstance(entry, tuple) else entry + pending.extend(_referenced_futures(callback)) + return False + + +def _creating_cell_waits_for( + owner: asyncio.Task[Any] | None, awaiter: asyncio.Task[Any] | None +) -> bool: + """Return whether the cell owner directly or transitively waits for awaiter.""" + if owner is None or awaiter is None: + return False + if owner is awaiter: + return True + waiter = getattr(owner, "_fut_waiter", None) + if not isinstance(waiter, asyncio.Future): + return False + return _future_contains(waiter, awaiter, set()) or _completion_reaches(awaiter, waiter) + + @dataclass(frozen=True) class BashResult: exit_code: int @@ -706,7 +793,7 @@ def __await__(self) -> Generator[Any, None, BashResult]: current_task = asyncio.current_task() except RuntimeError: current_task = None - if current_task is not None and current_task is self._creating_cell_task: + if _creating_cell_waits_for(self._creating_cell_task, current_task): self._awaited_by_creating_cell = True if self._released: return self._wait().__await__() diff --git a/prime-agent-runtime/test/test_repl.py b/prime-agent-runtime/test/test_repl.py index 141d72f46d..8e56e10748 100644 --- a/prime-agent-runtime/test/test_repl.py +++ b/prime-agent-runtime/test/test_repl.py @@ -753,6 +753,32 @@ def test_async_bash_completion_notifies_only_after_an_unawaited_creating_cell(se inspected = self.repl.execute("bash-inspect", "handle.poll().output") self.assertIn("detached", one(inspected, "result")["text"]) + def test_wrapper_awaits_in_creating_cell_suppress_bash_completion(self): + expressions = { + "gather": "(await asyncio.gather(bash('printf gather')))[0].output", + "wait-for": "(await asyncio.wait_for(bash('printf wait-for'), 1)).output", + "shield": "(await asyncio.shield(bash('printf shield'))).output", + "nested": "(await asyncio.shield(asyncio.gather(bash('printf nested'))))[0].output", + } + for label, expression in expressions.items(): + with self.subTest(label=label): + completed = self.repl.execute( + f"bash-wrapper-{label}", + f"from rlm import bash\nimport asyncio\n{expression}", + ) + self.assertIn(label, one(completed, "result")["text"]) + probe = self.repl.execute(f"bash-wrapper-{label}-probe", "await asyncio.sleep(0.05)") + request = one(probe, "host_request") + if request is not None: + self.repl.send( + { + "type": "host_reply", + "id": request["id"], + "data": {"status": "ok", "result": {}}, + } + ) + self.assertIsNone(request) + def test_background_task_await_does_not_suppress_bash_completion(self): code = "\n".join( [ @@ -762,6 +788,7 @@ def test_background_task_await_does_not_suppress_bash_completion(self): "async def consume():", " globals()['task_result'] = await task_handle", "waiter = asyncio.create_task(consume())", + "await asyncio.sleep(0.02)", "task_handle.pid", ] ) From cbb97a2925c4d66e708f75225993632a7f2ff235 Mon Sep 17 00:00:00 2001 From: Seth Date: Sat, 5 Sep 2026 20:42:12 -0400 Subject: [PATCH 4/6] refactor(runtime): simplify async bash completion --- .../.changes/async-bash-completion-message.md | 2 +- packages/coding-agent/src/core/messages.ts | 2 +- packages/coding-agent/src/core/rlm-runtime.ts | 17 ++-- prime-agent-runtime/src/rlm/bash.py | 94 +++++++------------ prime-agent-runtime/src/rlm/repl.py | 70 +++++--------- prime-agent-runtime/test/test_repl.py | 57 ++++------- 6 files changed, 86 insertions(+), 156 deletions(-) diff --git a/packages/coding-agent/.changes/async-bash-completion-message.md b/packages/coding-agent/.changes/async-bash-completion-message.md index 37671554ab..6e9a4b4e06 100644 --- a/packages/coding-agent/.changes/async-bash-completion-message.md +++ b/packages/coding-agent/.changes/async-bash-completion-message.md @@ -1 +1 @@ -- Added completion follow-ups for kernel `bash()` handles left running beyond their creating cell so the agent inspects finished async commands and continues the task. +- Added completion follow-ups when background kernel `bash()` commands finish so agents can inspect results and continue. diff --git a/packages/coding-agent/src/core/messages.ts b/packages/coding-agent/src/core/messages.ts index c81416f279..b974a25862 100644 --- a/packages/coding-agent/src/core/messages.ts +++ b/packages/coding-agent/src/core/messages.ts @@ -116,7 +116,7 @@ export interface AsyncBashCompletionDetails { exitCode: number; } -export interface AsyncBashCompletionMessage extends CustomMessage { +interface AsyncBashCompletionMessage extends CustomMessage { customType: typeof ASYNC_BASH_COMPLETION_CUSTOM_TYPE; content: string; } diff --git a/packages/coding-agent/src/core/rlm-runtime.ts b/packages/coding-agent/src/core/rlm-runtime.ts index c7fd2500ce..33c8ac84b0 100644 --- a/packages/coding-agent/src/core/rlm-runtime.ts +++ b/packages/coding-agent/src/core/rlm-runtime.ts @@ -52,13 +52,13 @@ export interface RlmFindModelsResult { export type RlmRunHandler = (request: RlmRunRequest) => Promise>; -export interface AsyncBashCompletionRequest { +interface AsyncBashCompletionRequest { pid: number; command: string; exitCode: number; } -export type AsyncBashCompletionHandler = (request: AsyncBashCompletionRequest) => void | Promise; +type AsyncBashCompletionHandler = (request: AsyncBashCompletionRequest) => void | Promise; export type RlmListSubagentsHandler = () => RlmListSubagentsResult | Promise; export type RlmDeleteSubagentHandler = (target: string) => Promise; export type RlmFindModelsHandler = (query: string, limit: number) => RlmFindModelsResult | Promise; @@ -189,20 +189,17 @@ export function createRlmRunHostHandler(handler: RlmRunHandler): HostRequestHand /** Adapt detached kernel bash completions into a validated host notification. */ export function createAsyncBashCompletionHostHandler(handler: AsyncBashCompletionHandler): HostRequestHandler { return async (payload) => { - if (!Number.isInteger(payload.pid) || (payload.pid as number) <= 0) { + const { pid, command, exitCode } = payload; + if (typeof pid !== "number" || !Number.isInteger(pid) || pid <= 0) { throw new Error("bash.completed pid must be a positive integer"); } - if (typeof payload.command !== "string" || !payload.command) { + if (typeof command !== "string" || !command) { throw new Error("bash.completed command must be a non-empty string"); } - if (!Number.isInteger(payload.exitCode)) { + if (typeof exitCode !== "number" || !Number.isInteger(exitCode)) { throw new Error("bash.completed exitCode must be an integer"); } - await handler({ - pid: payload.pid as number, - command: payload.command, - exitCode: payload.exitCode as number, - }); + await handler({ pid, command, exitCode }); return {}; }; } diff --git a/prime-agent-runtime/src/rlm/bash.py b/prime-agent-runtime/src/rlm/bash.py index 536446782a..87558b64e7 100644 --- a/prime-agent-runtime/src/rlm/bash.py +++ b/prime-agent-runtime/src/rlm/bash.py @@ -70,75 +70,51 @@ def _consume_notice_task(task: asyncio.Task[None]) -> None: task.exception() -def _referenced_futures( - value: Any, *, depth: int = 0, seen: set[int] | None = None -) -> set[asyncio.Future[Any]]: - """Find futures captured by asyncio's small internal completion callbacks.""" - if isinstance(value, asyncio.Future): - return {value} - if depth >= 4: - return set() - if seen is None: - seen = set() - identity = id(value) - if identity in seen: - return set() - seen.add(identity) - - nested: list[Any] = [] - if isinstance(value, functools.partial): - nested.extend((value.func, value.args, value.keywords)) - elif isinstance(value, dict): - nested.extend(value.keys()) - nested.extend(value.values()) - elif isinstance(value, (tuple, list, set, frozenset)): - nested.extend(value) - else: - closure = getattr(value, "__closure__", None) - if closure: +def _completion_reaches(start: asyncio.Future[Any], target: asyncio.Future[Any]) -> bool: + """Follow futures captured by asyncio wrapper completion callbacks.""" + pending = [start] + seen_futures: set[int] = set() + seen_values: set[int] = set() + + def collect(value: Any, depth: int = 0) -> None: + if isinstance(value, asyncio.Future): + pending.append(value) + return + identity = id(value) + if depth >= 4 or identity in seen_values: + return + seen_values.add(identity) + + nested: list[Any] = [] + if isinstance(value, functools.partial): + nested.extend((value.func, value.args, value.keywords)) + elif isinstance(value, dict): + nested.extend(value.keys()) + nested.extend(value.values()) + elif isinstance(value, (tuple, list, set, frozenset)): + nested.extend(value) + else: + closure = getattr(value, "__closure__", None) or () for cell in closure: try: nested.append(cell.cell_contents) except ValueError: pass - bound_self = getattr(value, "__self__", None) - if bound_self is not None: - nested.append(bound_self) - - futures: set[asyncio.Future[Any]] = set() - for item in nested: - futures.update(_referenced_futures(item, depth=depth + 1, seen=seen)) - return futures - + bound_self = getattr(value, "__self__", None) + if bound_self is not None: + nested.append(bound_self) + for item in nested: + collect(item, depth + 1) -def _future_contains(root: asyncio.Future[Any], target: asyncio.Future[Any], seen: set[int]) -> bool: - """Follow aggregate children, such as asyncio.gather's private future list.""" - if root is target: - return True - if id(root) in seen: - return False - seen.add(id(root)) - children = getattr(root, "_children", None) or () - return any( - isinstance(child, asyncio.Future) and _future_contains(child, target, seen) for child in children - ) - - -def _completion_reaches(start: asyncio.Future[Any], target: asyncio.Future[Any]) -> bool: - """Follow callback-captured futures from an inner await to its wrapper future.""" - pending = [start] - seen: set[int] = set() while pending: future = pending.pop() if future is target: return True - if id(future) in seen: + if id(future) in seen_futures: continue - seen.add(id(future)) - callbacks = getattr(future, "_callbacks", None) or () - for entry in callbacks: - callback = entry[0] if isinstance(entry, tuple) else entry - pending.extend(_referenced_futures(callback)) + seen_futures.add(id(future)) + for entry in getattr(future, "_callbacks", None) or (): + collect(entry[0] if isinstance(entry, tuple) else entry) return False @@ -153,7 +129,7 @@ def _creating_cell_waits_for( waiter = getattr(owner, "_fut_waiter", None) if not isinstance(waiter, asyncio.Future): return False - return _future_contains(waiter, awaiter, set()) or _completion_reaches(awaiter, waiter) + return _completion_reaches(awaiter, waiter) @dataclass(frozen=True) diff --git a/prime-agent-runtime/src/rlm/repl.py b/prime-agent-runtime/src/rlm/repl.py index 063f957750..d3324be774 100644 --- a/prime-agent-runtime/src/rlm/repl.py +++ b/prime-agent-runtime/src/rlm/repl.py @@ -45,21 +45,23 @@ _write_lock = threading.Lock() _loop: asyncio.AbstractEventLoop | None = None _serve_task: asyncio.Task[Any] | None = None -# Attribution rides task context: asyncio tasks copy it at creation, so a -# detached task spawned by a cell keeps writing under that cell's id after -# the cell finishes. Threads start with a fresh context and emit id null. + + +class _CellExecution: + def __init__(self) -> None: + self.finished = asyncio.Event() + self.owner: asyncio.Task[Any] | None = None + + +# Asyncio tasks copy cell context at creation, so detached tasks retain their +# output attribution and completion barrier. Threads start with a fresh context. _current_cell: contextvars.ContextVar[str | None] = contextvars.ContextVar("_current_cell", default=None) -_current_cell_generation: contextvars.ContextVar[int | None] = contextvars.ContextVar( - "_current_cell_generation", default=None +_current_cell_execution: contextvars.ContextVar[_CellExecution | None] = contextvars.ContextVar( + "_current_cell_execution", default=None ) _active: dict[str, Any] = {"task": None, "rid": None, "interrupted": False} _cell_counter = 0 _pending_host: dict[str, "asyncio.Future[dict[str, Any]]"] = {} -# Request ids may be reused. Bash completion barriers therefore use the -# runtime's monotonic cell generation rather than the caller-supplied id. -_active_cell_generations: set[int] = set() -_cell_owner_tasks: dict[int, asyncio.Task[Any]] = {} -_cell_finished_events: dict[int, asyncio.Event] = {} # Set on the loop thread once stdin hits EOF or a shutdown request arrives; no # host reply can arrive after that, so waiting (and future) host_request calls fail. _host_closed = False @@ -107,24 +109,11 @@ def is_active() -> bool: def current_cell_completion_context() -> tuple[asyncio.Event, asyncio.Task[Any] | None] | None: - """Return the calling cell's completion barrier and owning execution task. - - Detached asyncio tasks inherit the internal cell generation. If such a task - runs after that cell finished, return an already-open barrier and no owner. - """ - cell_generation = _current_cell_generation.get() - if cell_generation is None or _loop is None: + """Return the calling cell's completion barrier and owning execution task.""" + execution = _current_cell_execution.get() + if execution is None: return None - with _interrupt_lock: - if cell_generation not in _active_cell_generations: - event = asyncio.Event() - event.set() - return event, None - event = _cell_finished_events.get(cell_generation) - if event is None: - event = asyncio.Event() - _cell_finished_events[cell_generation] = event - return event, _cell_owner_tasks.get(cell_generation) + return execution.finished, execution.owner async def host_request(data: dict[str, Any]) -> dict[str, Any]: @@ -455,15 +444,6 @@ def _finish_request(rid: str) -> None: _finish_locked(rid) -def _finish_cell_generation(cell_generation: int) -> None: - with _interrupt_lock: - _active_cell_generations.discard(cell_generation) - _cell_owner_tasks.pop(cell_generation, None) - cell_finished = _cell_finished_events.pop(cell_generation, None) - if cell_finished is not None: - cell_finished.set() - - _RUNTIME_FILE = __file__ @@ -574,20 +554,15 @@ async def _handle_execute(req: dict[str, Any], ns: dict[str, Any]) -> None: global _cell_counter cell_id = req["id"] _cell_counter += 1 - cell_generation = _cell_counter - filename = f"" - with _interrupt_lock: - _active_cell_generations.add(cell_generation) - # The cell task (created below) copies this context, so writes made from - # the cell and from asyncio tasks it spawns carry this cell's identity. + filename = f"" + execution = _CellExecution() cell_token = _current_cell.set(cell_id) - generation_token = _current_cell_generation.set(cell_generation) + execution_token = _current_cell_execution.set(execution) try: codes, has_trailing = _compile_cell(req["code"], filename) assert _loop is not None task = _loop.create_task(_run_codes(codes, ns)) - with _interrupt_lock: - _cell_owner_tasks[cell_generation] = task + execution.owner = task status, value, error = await _run_guarded(task, cell_id) result_text: str | None = None try: @@ -612,8 +587,9 @@ async def _handle_execute(req: dict[str, Any], ns: dict[str, Any]) -> None: _send(error) _send({"event": "done", "id": cell_id, "status": status}) finally: - _finish_cell_generation(cell_generation) - _current_cell_generation.reset(generation_token) + execution.owner = None + execution.finished.set() + _current_cell_execution.reset(execution_token) _current_cell.reset(cell_token) diff --git a/prime-agent-runtime/test/test_repl.py b/prime-agent-runtime/test/test_repl.py index 8e56e10748..b07ce070a0 100644 --- a/prime-agent-runtime/test/test_repl.py +++ b/prime-agent-runtime/test/test_repl.py @@ -109,6 +109,19 @@ def one(events: list[dict], kind: str) -> dict | None: return matches[0] if matches else None +def wait_for_host_request(repl: ReplProcess, events: list[dict]) -> dict: + request = one(events, "host_request") + while request is None: + event = repl.read_event() + if event.get("event") == "host_request": + request = event + return request + + +def reply_ok(repl: ReplProcess, request: dict) -> None: + repl.send({"type": "host_reply", "id": request["id"], "data": {"status": "ok", "result": {}}}) + + class ReplTest(unittest.TestCase): def setUp(self) -> None: self.repl = ReplProcess() @@ -729,11 +742,7 @@ def test_async_bash_completion_notifies_only_after_an_unawaited_creating_cell(se "bash-detached", "handle = bash('sleep 0.05; printf detached')\nhandle.pid" ) pid = int(one(started, "result")["text"]) - request = one(started, "host_request") - while request is None: - event = self.repl.read_event() - if event.get("event") == "host_request": - request = event + request = wait_for_host_request(self.repl, started) self.assertEqual( request["data"], { @@ -743,13 +752,7 @@ def test_async_bash_completion_notifies_only_after_an_unawaited_creating_cell(se "exitCode": 0, }, ) - self.repl.send( - { - "type": "host_reply", - "id": request["id"], - "data": {"status": "ok", "result": {}}, - } - ) + reply_ok(self.repl, request) inspected = self.repl.execute("bash-inspect", "handle.poll().output") self.assertIn("detached", one(inspected, "result")["text"]) @@ -770,13 +773,7 @@ def test_wrapper_awaits_in_creating_cell_suppress_bash_completion(self): probe = self.repl.execute(f"bash-wrapper-{label}-probe", "await asyncio.sleep(0.05)") request = one(probe, "host_request") if request is not None: - self.repl.send( - { - "type": "host_reply", - "id": request["id"], - "data": {"status": "ok", "result": {}}, - } - ) + reply_ok(self.repl, request) self.assertIsNone(request) def test_background_task_await_does_not_suppress_bash_completion(self): @@ -794,20 +791,10 @@ def test_background_task_await_does_not_suppress_bash_completion(self): ) started = self.repl.execute("bash-task-waiter", code) pid = int(one(started, "result")["text"]) - request = one(started, "host_request") - while request is None: - event = self.repl.read_event() - if event.get("event") == "host_request": - request = event + request = wait_for_host_request(self.repl, started) self.assertEqual(request["data"]["type"], "bash.completed") self.assertEqual(request["data"]["pid"], pid) - self.repl.send( - { - "type": "host_reply", - "id": request["id"], - "data": {"status": "ok", "result": {}}, - } - ) + reply_ok(self.repl, request) def test_reused_request_id_does_not_capture_old_cell_bash_completion(self): setup = "\n".join( @@ -839,13 +826,7 @@ def test_reused_request_id_does_not_capture_old_cell_bash_completion(self): self.assertFalse(event.get("event") == "done" and event.get("id") == "reused-cell-id") self.assertEqual(request["data"]["type"], "bash.completed") self.assertEqual(request["data"]["command"], "printf reused-id") - self.repl.send( - { - "type": "host_reply", - "id": request["id"], - "data": {"status": "ok", "result": {}}, - } - ) + reply_ok(self.repl, request) self.assertEqual(one(self.repl.until_done("reused-cell-id"), "done")["status"], "ok") def test_protocol_framing_under_noise(self): From 1eab8bc5177100146fd8beca466e04a7191dd101 Mon Sep 17 00:00:00 2001 From: Seth Date: Sat, 5 Sep 2026 21:06:00 -0400 Subject: [PATCH 5/6] fix(runtime): follow explicit asyncio ownership --- prime-agent-runtime/src/rlm/bash.py | 30 ++++++++++++++++++++------- prime-agent-runtime/test/test_repl.py | 23 ++++++++++++++++---- 2 files changed, 42 insertions(+), 11 deletions(-) diff --git a/prime-agent-runtime/src/rlm/bash.py b/prime-agent-runtime/src/rlm/bash.py index 87558b64e7..0c4536c202 100644 --- a/prime-agent-runtime/src/rlm/bash.py +++ b/prime-agent-runtime/src/rlm/bash.py @@ -45,6 +45,11 @@ _CANCEL_TERM_GRACE = 0.5 _CANCEL_KILL_WAIT = 2.0 _COMPLETION_NOTICE_COMMAND_CAP = 1000 +_ASYNCIO_WRAPPER_CALLBACKS = { + ("asyncio.tasks", "gather.._done_callback"), + ("asyncio.tasks", "shield.._inner_done_callback"), + ("asyncio.tasks", "_release_waiter"), +} _live_handles: set["BashHandle"] = set() _live_lock = threading.Lock() @@ -70,8 +75,10 @@ def _consume_notice_task(task: asyncio.Task[None]) -> None: task.exception() -def _completion_reaches(start: asyncio.Future[Any], target: asyncio.Future[Any]) -> bool: - """Follow futures captured by asyncio wrapper completion callbacks.""" +def _completion_reaches( + start: asyncio.Future[Any], targets: tuple[asyncio.Future[Any], ...] +) -> bool: + """Follow asyncio's wrapper and TaskGroup ownership callbacks.""" pending = [start] seen_futures: set[int] = set() seen_values: set[int] = set() @@ -108,13 +115,21 @@ def collect(value: Any, depth: int = 0) -> None: while pending: future = pending.pop() - if future is target: + if any(future is target for target in targets): return True if id(future) in seen_futures: continue seen_futures.add(id(future)) for entry in getattr(future, "_callbacks", None) or (): - collect(entry[0] if isinstance(entry, tuple) else entry) + callback = entry[0] if isinstance(entry, tuple) else entry + base = callback.func if isinstance(callback, functools.partial) else callback + identity = (getattr(base, "__module__", None), getattr(base, "__qualname__", None)) + if identity in _ASYNCIO_WRAPPER_CALLBACKS: + collect(callback) + elif identity == ("asyncio.taskgroups", "TaskGroup._on_task_done"): + parent = getattr(getattr(callback, "__self__", None), "_parent_task", None) + if isinstance(parent, asyncio.Future): + pending.append(parent) return False @@ -127,9 +142,10 @@ def _creating_cell_waits_for( if owner is awaiter: return True waiter = getattr(owner, "_fut_waiter", None) - if not isinstance(waiter, asyncio.Future): - return False - return _completion_reaches(awaiter, waiter) + targets: tuple[asyncio.Future[Any], ...] = (owner,) + if isinstance(waiter, asyncio.Future): + targets += (waiter,) + return _completion_reaches(awaiter, targets) @dataclass(frozen=True) diff --git a/prime-agent-runtime/test/test_repl.py b/prime-agent-runtime/test/test_repl.py index b07ce070a0..2ab4901787 100644 --- a/prime-agent-runtime/test/test_repl.py +++ b/prime-agent-runtime/test/test_repl.py @@ -757,17 +757,27 @@ def test_async_bash_completion_notifies_only_after_an_unawaited_creating_cell(se self.assertIn("detached", one(inspected, "result")["text"]) def test_wrapper_awaits_in_creating_cell_suppress_bash_completion(self): - expressions = { + snippets = { "gather": "(await asyncio.gather(bash('printf gather')))[0].output", "wait-for": "(await asyncio.wait_for(bash('printf wait-for'), 1)).output", "shield": "(await asyncio.shield(bash('printf shield'))).output", "nested": "(await asyncio.shield(asyncio.gather(bash('printf nested'))))[0].output", + "task-group": "\n".join( + [ + "handle = bash('printf task-group')", + "async def consume():", + " return await handle", + "async with asyncio.TaskGroup() as group:", + " task = group.create_task(consume())", + "task.result().output", + ] + ), } - for label, expression in expressions.items(): + for label, snippet in snippets.items(): with self.subTest(label=label): completed = self.repl.execute( f"bash-wrapper-{label}", - f"from rlm import bash\nimport asyncio\n{expression}", + f"from rlm import bash\nimport asyncio\n{snippet}", ) self.assertIn(label, one(completed, "result")["text"]) probe = self.repl.execute(f"bash-wrapper-{label}-probe", "await asyncio.sleep(0.05)") @@ -785,7 +795,12 @@ def test_background_task_await_does_not_suppress_bash_completion(self): "async def consume():", " globals()['task_result'] = await task_handle", "waiter = asyncio.create_task(consume())", - "await asyncio.sleep(0.02)", + "bookkeeping = asyncio.get_running_loop().create_future()", + "def callback_for(marker):", + " return lambda _: marker.cancelled()", + "waiter.add_done_callback(callback_for(bookkeeping))", + "asyncio.get_running_loop().call_later(0.02, bookkeeping.set_result, None)", + "await bookkeeping", "task_handle.pid", ] ) From 1bb4d43a0a64366217f8fec38d574e3917fa550d Mon Sep 17 00:00:00 2001 From: Seth Date: Sat, 5 Sep 2026 21:15:15 -0400 Subject: [PATCH 6/6] fix(runtime): follow nested asyncio wait ownership --- prime-agent-runtime/src/rlm/bash.py | 5 +++++ prime-agent-runtime/test/test_repl.py | 25 +++++++++++++++++++------ 2 files changed, 24 insertions(+), 6 deletions(-) diff --git a/prime-agent-runtime/src/rlm/bash.py b/prime-agent-runtime/src/rlm/bash.py index 0c4536c202..36cdd77594 100644 --- a/prime-agent-runtime/src/rlm/bash.py +++ b/prime-agent-runtime/src/rlm/bash.py @@ -48,6 +48,7 @@ _ASYNCIO_WRAPPER_CALLBACKS = { ("asyncio.tasks", "gather.._done_callback"), ("asyncio.tasks", "shield.._inner_done_callback"), + ("asyncio.tasks", "_wait.._on_completion"), ("asyncio.tasks", "_release_waiter"), } @@ -126,6 +127,10 @@ def collect(value: Any, depth: int = 0) -> None: identity = (getattr(base, "__module__", None), getattr(base, "__qualname__", None)) if identity in _ASYNCIO_WRAPPER_CALLBACKS: collect(callback) + elif identity == (None, "Task.task_wakeup"): + task = getattr(callback, "__self__", None) + if isinstance(task, asyncio.Task): + pending.append(task) elif identity == ("asyncio.taskgroups", "TaskGroup._on_task_done"): parent = getattr(getattr(callback, "__self__", None), "_parent_task", None) if isinstance(parent, asyncio.Future): diff --git a/prime-agent-runtime/test/test_repl.py b/prime-agent-runtime/test/test_repl.py index 2ab4901787..7aa284a6e3 100644 --- a/prime-agent-runtime/test/test_repl.py +++ b/prime-agent-runtime/test/test_repl.py @@ -757,21 +757,34 @@ def test_async_bash_completion_notifies_only_after_an_unawaited_creating_cell(se self.assertIn("detached", one(inspected, "result")["text"]) def test_wrapper_awaits_in_creating_cell_suppress_bash_completion(self): + def task_group(label: str, await_expression: str) -> str: + return "\n".join( + [ + f"handle = bash('printf {label}')", + "async def consume():", + f" return {await_expression}", + "async with asyncio.TaskGroup() as group:", + " task = group.create_task(consume())", + "task.result().output", + ] + ) + snippets = { "gather": "(await asyncio.gather(bash('printf gather')))[0].output", "wait-for": "(await asyncio.wait_for(bash('printf wait-for'), 1)).output", "shield": "(await asyncio.shield(bash('printf shield'))).output", "nested": "(await asyncio.shield(asyncio.gather(bash('printf nested'))))[0].output", - "task-group": "\n".join( + "wait": "\n".join( [ - "handle = bash('printf task-group')", - "async def consume():", - " return await handle", - "async with asyncio.TaskGroup() as group:", - " task = group.create_task(consume())", + "handle = bash('printf wait')", + "task = asyncio.ensure_future(handle)", + "await asyncio.wait({task})", "task.result().output", ] ), + "task-group": task_group("task-group", "await handle"), + "task-group-gather": task_group("task-group-gather", "(await asyncio.gather(handle))[0]"), + "task-group-shield": task_group("task-group-shield", "await asyncio.shield(handle)"), } for label, snippet in snippets.items(): with self.subTest(label=label):