diff --git a/README.md b/README.md index d6094ec..37bc1cd 100644 --- a/README.md +++ b/README.md @@ -550,3 +550,63 @@ stops with a failure event when the inbox event limit is reached. There are no model-authored callbacks, event selectors, or output predicates. Activity subscriptions become `completed` after their target permanently terminates, flushing pending activity and releasing their active slot. Watches of idle persistent agents remain active. Watching an already-finished target returns a completed subscription. + +### Separate tool execution + + + + Colocated, the agent loop, model credentials and the agent's tools share one machine. Split, the harness keeps the loop, credentials, supervisor and ledger, and a worker in the workspace runs the tools over one authenticated connection. + + +The ACP runtime metadata accepts `execution_command` (a trusted argv launching +`python -m rlm.execution`) and `execution_cwd` (the task workspace). For example, +an orchestrator can start the worker with `docker exec -i +/path/to/python -m rlm.execution`. Both processes must have the same nano-rlm +version installed. The orchestrator owns container provisioning and teardown. + +The engine, model credentials, supervisor, and authoritative conversation logs +remain in the harness process. Native tools, persistent IPython kernels, and builtin skills run +in the worker environment. Recursive agents inherit the execution destination, +keep separate kernels, and share the task workspace. With the stdio transport, +communication and recursive callbacks use the worker's stdio, so a task container +can have networking disabled. +The command itself defines the isolation boundary; launching a local Python +worker alone provides process separation, not filesystem isolation. + +ACP initialization advertises `ai.prime.rlm/execution-v1` with the supported +transports (`stdio`, `tcp`) and tools (`ipython`, `bash`, `edit`, `fetch`). Tool +calls cross the boundary as a name and arguments; filesystem operations happen +inside the worker. `ExecutionTransport` supplies `connect()` (reader and writer) +and `close()` independently of tool dispatch. + +For a trusted private TCP network, launch `python -m rlm.execution --listen +0.0.0.0 --port 8765` with a random `RLM_EXECUTION_TOKEN` of at least 32 characters. +With `--port 0` the kernel picks a free port, which the worker prints on stderr as +`listening on HOST:PORT` once it accepts connections. +Set `execution_tcp` to `{"host": "worker-address", "port": 8765, "token": "..."}` +instead of `execution_command`. Each connection owns a separate execution session. +The token authenticates connections; this transport does not encrypt traffic, so +use a private network or an encrypted tunnel. The worker removes its token from +the environment before starting tools. The orchestrator owns listener teardown. + +Supervisor shell jobs (`rlm.shell`), services (`rlm.services`) and path watches +(`rlm.watch.path`) run in the worker's environment, beside the kernels; the +supervisor keeps their records and inbox events, so the agent gets the same +runtime prompt and API as in local execution. The current agent's history is a +local snapshot refreshed before each cell; `await child.history()` reads a child's +ledger from the supervisor in pages. Frames are limited to 16 MiB: a command or +cell keeps the first and last 512 KiB of its output, and a long ledger syncs in +chunks. Kernel recovery preserves the existing no-replay behavior. Worker +disconnects fail the session. + + + + One agent step in split mode: the harness sends the tool call with the new ledger records, the cell runs in the workspace and calls back for rlm.* operations, and the result returns to the harness. A tool call costs one round trip, each rlm.* call one more. + + +Each tool call costs one round trip between harness and worker, each `rlm.*` +call one more, and `rlm.shell.run` two. + +A session that closes stops what is still attached to its kernels, as a closing +terminal does; daemons and bash background jobs keep running, as they do in local +execution, so the orchestrator stops what remains when the rollout ends. diff --git a/docs/images/split-execution-dark.svg b/docs/images/split-execution-dark.svg new file mode 100644 index 0000000..64e42bc --- /dev/null +++ b/docs/images/split-execution-dark.svg @@ -0,0 +1 @@ +Colocatedthe agent's commands run beside the harnessModel APItask VMharnessagent loop · model credentialssupervisor · conversation ledgernano-rlm installagent's toolsIPython kernel · bash · edit · fetchshell jobs · services · watchestask filesone machine, one filesystem, one process treeA command can read the API key, kill orstarve the loop, or edit the ledger.Splitthe harness keeps the loop; the workspace runs the toolsModel APIharness VMagent loopthe model's turnsmodel credentialsnever leave this VMsupervisoragents · inbox · job recordsconversation ledgerauthoritative copyits own image and installworkspace VM · task imagepython -m rlm.executionIPython kernelsone per agent, persistentbash · edit · fetchas the agent userjobs · services · watchesrun here, recorded theretask files · no credentialsgrading runs heretool callresultrlm.* callreplyjob eventsledger recordsone TCP connectiontoken-authenticatedprivate network(or a worker's stdio) diff --git a/docs/images/split-execution-step-dark.svg b/docs/images/split-execution-step-dark.svg new file mode 100644 index 0000000..010b542 --- /dev/null +++ b/docs/images/split-execution-step-dark.svg @@ -0,0 +1 @@ +Model APIharness VMengine · supervisorworkspace VMworker · IPython kernelconversationtool call ipython(code)tool call + new ledger recordscell runsawait rlm.agent.spawn(...)child handleresult: first and last 512 KiB + whole sizenext turn shows its truncation1 round tripper tool call+1 round tripper rlm.* call diff --git a/docs/images/split-execution-step.svg b/docs/images/split-execution-step.svg new file mode 100644 index 0000000..719ca21 --- /dev/null +++ b/docs/images/split-execution-step.svg @@ -0,0 +1 @@ +Model APIharness VMengine · supervisorworkspace VMworker · IPython kernelconversationtool call ipython(code)tool call + new ledger recordscell runsawait rlm.agent.spawn(...)child handleresult: first and last 512 KiB + whole sizenext turn shows its truncation1 round tripper tool call+1 round tripper rlm.* call diff --git a/docs/images/split-execution.svg b/docs/images/split-execution.svg new file mode 100644 index 0000000..b815d8a --- /dev/null +++ b/docs/images/split-execution.svg @@ -0,0 +1 @@ +Colocatedthe agent's commands run beside the harnessModel APItask VMharnessagent loop · model credentialssupervisor · conversation ledgernano-rlm installagent's toolsIPython kernel · bash · edit · fetchshell jobs · services · watchestask filesone machine, one filesystem, one process treeA command can read the API key, kill orstarve the loop, or edit the ledger.Splitthe harness keeps the loop; the workspace runs the toolsModel APIharness VMagent loopthe model's turnsmodel credentialsnever leave this VMsupervisoragents · inbox · job recordsconversation ledgerauthoritative copyits own image and installworkspace VM · task imagepython -m rlm.executionIPython kernelsone per agent, persistentbash · edit · fetchas the agent userjobs · services · watchesrun here, recorded theretask files · no credentialsgrading runs heretool callresultrlm.* callreplyjob eventsledger recordsone TCP connectiontoken-authenticatedprivate network(or a worker's stdio) diff --git a/src/rlm/acp.py b/src/rlm/acp.py index 8aacfde..7f173e0 100644 --- a/src/rlm/acp.py +++ b/src/rlm/acp.py @@ -45,6 +45,7 @@ from rlm.engine import RLMEngine from rlm.config import ( ExecutionPolicy, + TcpExecutionConfig, InvocationContext, ProviderConfig, RuntimeConfig, @@ -75,6 +76,9 @@ class _RuntimeMetadata(_ContractModel): leaf_append_to_system_prompt: str | None = None skills: list[Annotated[str, Field(min_length=1)]] builtin_tools: list[Annotated[str, Field(min_length=1)]] | None = None + execution_tcp: TcpExecutionConfig | None = None + execution_cwd: str | None = None + execution_command: list[Annotated[str, Field(min_length=1)]] | None = None kernel_env: dict[str, str] search_api_key: str | None network_block: list[str] = Field(default_factory=list) @@ -189,6 +193,33 @@ def _runtime_config(meta_kwargs: Any) -> tuple[RuntimeConfig, str]: payload = _RuntimeMetadata.model_validate( meta_kwargs.get(RUNTIME_METADATA_KEY), ) + result = ( + RuntimeConfig( + model=payload.model, + provider=payload.provider, + invocation=InvocationContext(), + policy=payload.policy, + system_prompt_path=payload.system_prompt_path, + append_to_system_prompt=payload.append_to_system_prompt, + subagent_append_to_system_prompt=payload.subagent_append_to_system_prompt, + leaf_append_to_system_prompt=payload.leaf_append_to_system_prompt, + skills=tuple(payload.skills), + builtin_tools=( + tuple(payload.builtin_tools) + if payload.builtin_tools is not None + else None + ), + kernel_env=tuple(payload.kernel_env.items()), + execution_tcp=payload.execution_tcp, + execution_cwd=payload.execution_cwd, + execution_command=tuple(payload.execution_command) + if payload.execution_command is not None + else None, + search_api_key=payload.search_api_key, + network_block=tuple(payload.network_block), + ), + payload.session_id, + ) except ValidationError as error: raise RequestError.invalid_params( { @@ -199,28 +230,7 @@ def _runtime_config(meta_kwargs: Any) -> tuple[RuntimeConfig, str]: } ) from error - return ( - RuntimeConfig( - model=payload.model, - provider=payload.provider, - invocation=InvocationContext(), - policy=payload.policy, - system_prompt_path=payload.system_prompt_path, - append_to_system_prompt=payload.append_to_system_prompt, - subagent_append_to_system_prompt=payload.subagent_append_to_system_prompt, - leaf_append_to_system_prompt=payload.leaf_append_to_system_prompt, - skills=tuple(payload.skills), - builtin_tools=( - tuple(payload.builtin_tools) - if payload.builtin_tools is not None - else None - ), - kernel_env=tuple(payload.kernel_env.items()), - search_api_key=payload.search_api_key, - network_block=tuple(payload.network_block), - ), - payload.session_id, - ) + return result def _mcp_servers( @@ -293,7 +303,14 @@ async def initialize( ), ), agent_info=Implementation(name="rlm", title="RLM", version=version("rlm")), - field_meta={CONTRACT_METADATA_KEY: True, "steering": {"supported": True}}, + field_meta={ + CONTRACT_METADATA_KEY: True, + "steering": {"supported": True}, + "ai.prime.rlm/execution-v1": { + "transports": ["stdio", "tcp"], + "tools": ["ipython", "bash", "edit", "fetch"], + }, + }, ) async def new_session( diff --git a/src/rlm/agent.py b/src/rlm/agent.py index 507eb7c..36bf9eb 100644 --- a/src/rlm/agent.py +++ b/src/rlm/agent.py @@ -48,7 +48,18 @@ async def info(self) -> AgentInfo: async def history(self) -> History: """Read a fresh snapshot of this agent's local conversation history.""" - return await history(session_dir=self.session_dir) + if not broker.is_active(): + return await history(session_dir=self.session_dir) + # The supervisor returns a long ledger in pages that each fit one frame. + events: list[dict] = [] + offset: int | None = 0 + while offset is not None: + page = await broker.agent_request( + "agent.history", agent_id=self.id, offset=offset + ) + events += page["records"] + offset = page["next"] + return History.from_events(self.session_dir, events) async def result(self, *, yield_after: float = 300) -> AgentResult: """Wait up to yield_after seconds for the child to finish, then return its latest state. diff --git a/src/rlm/broker.py b/src/rlm/broker.py index 354702e..f6d5862 100644 --- a/src/rlm/broker.py +++ b/src/rlm/broker.py @@ -85,6 +85,15 @@ class BrokerAgentHandleRequest(TypedDict): agent_id: Annotated[str, Field(min_length=1)] +class BrokerAgentHistoryRequest(TypedDict): + __pydantic_config__ = ConfigDict(extra="forbid", strict=True) + op: Literal["agent.history"] + capability: Annotated[str, Field(min_length=1)] + scope_id: Annotated[str, Field(min_length=1)] + agent_id: Annotated[str, Field(min_length=1)] + offset: Annotated[int, Field(ge=0)] + + class BrokerAgentResultRequest(TypedDict): __pydantic_config__ = ConfigDict(extra="forbid", strict=True) op: Literal["agent.result"] @@ -324,6 +333,7 @@ class BrokerSkillRequest(TypedDict): | BrokerAgentGetRequest | BrokerAgentListRequest | BrokerAgentHandleRequest + | BrokerAgentHistoryRequest | BrokerAgentResultRequest | BrokerAgentWaitRequest | BrokerSkillRequest @@ -403,6 +413,11 @@ def result_from_payload(value: dict[str, Any]) -> AgentResult: raise RuntimeError("invalid response from RLM supervisor") from None +def is_active() -> bool: + """Whether the current process is inside a brokered execution scope.""" + return _endpoint is not None and _scope_id is not None + + def configure(endpoint: BrokerEndpoint | None) -> None: global _endpoint _endpoint = endpoint diff --git a/src/rlm/compaction.py b/src/rlm/compaction.py index 9330c79..98d5ae6 100644 --- a/src/rlm/compaction.py +++ b/src/rlm/compaction.py @@ -5,6 +5,8 @@ from openai import APIError, APIStatusError, AsyncOpenAI +from rlm.tools.output import KEPT_PER_END, OutputSize + CHECKPOINT_PROMPT = """Create a concise continuation summary for the current task. Preserve what is needed to resume accurately: - The user's objective, exact requirements, constraints, and unresolved decisions. @@ -136,24 +138,33 @@ async def discover_threshold(client: AsyncOpenAI, model: str) -> int | None: return default_threshold(window) if window is not None else None -def truncate_tool_output(text: str, max_bytes: int = TOOL_OUTPUT_MAX_BYTES) -> str: - """Keep the head and tail of an oversized tool result and say what was cut.""" +def truncate_tool_output( + text: str, + max_bytes: int = TOOL_OUTPUT_MAX_BYTES, + omitted: OutputSize | None = None, +) -> str: + """Keep the head and tail of an oversized tool result and say what was cut. + `omitted` is how much more the whole result measures than `text`, which kept only + its start and end; the note describes the whole result.""" data = text.encode("utf-8") - if len(data) <= max_bytes: + size = OutputSize.of(text) + (omitted or OutputSize()) + if size.nbytes <= max_bytes: return text keep = max_bytes // 2 + if omitted is not None: + keep = min(keep, KEPT_PER_END) head = data[:keep].decode("utf-8", errors="ignore") tail = data[-keep:].decode("utf-8", errors="ignore") return ( - f"Warning: truncated output (original token count: {estimated_tokens(text)})\n" - f"Total output lines: {text.count(chr(10)) + 1}\n\n" - f"{head}\n[... {len(data) - 2 * keep} bytes truncated ...]\n{tail}" + f"Warning: truncated output (original token count: {estimated_tokens(size.nchars)})\n" + f"Total output lines: {size.newlines + 1}\n\n" + f"{head}\n[... {size.nbytes - 2 * keep} bytes truncated ...]\n{tail}" ) -def estimated_tokens(chars: str) -> int: - """Rough token count at four characters per token.""" - return (len(chars) + 3) // 4 +def estimated_tokens(chars: str | int) -> int: + """Rough token count at four characters per token, of text or a character count.""" + return ((chars if isinstance(chars, int) else len(chars)) + 3) // 4 def compactable(messages: list[dict]) -> bool: diff --git a/src/rlm/config.py b/src/rlm/config.py index 0a819b5..7aa932b 100644 --- a/src/rlm/config.py +++ b/src/rlm/config.py @@ -7,7 +7,7 @@ from __future__ import annotations -from pydantic import BaseModel, ConfigDict, Field, field_validator +from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator from rlm.network import NetworkBlocklist from rlm.semantic import ACP_EXTENSION_HEADER_NAMES @@ -96,6 +96,14 @@ class ExecutionPolicy(_ConfigModel): allow_git: bool = False +class TcpExecutionConfig(_ConfigModel): + """Authenticated worker endpoint on a trusted private network.""" + + host: str = Field(min_length=1) + port: int = Field(gt=0, le=65535) + token: str = Field(min_length=32, repr=False) + + class RuntimeConfig(_ConfigModel): """Configuration resolved once at an RLM process boundary.""" @@ -112,10 +120,29 @@ class RuntimeConfig(_ConfigModel): """Builtin tool set for every engine in the tree; None = the registry default (`ipython` alone). Validated against the registry when the engine starts.""" kernel_env: tuple[tuple[str, str], ...] = Field(default=(), repr=False) + execution_cwd: str | None = None + execution_command: tuple[str, ...] | None = None + """Trusted argv launching an rlm.execution worker in the task environment.""" + execution_tcp: TcpExecutionConfig | None = None search_api_key: str | None = Field(default=None, repr=False) network_block: NetworkBlocklist = () """Search-result destination rules shared by every agent in the session tree.""" + @model_validator(mode="after") + def _validate_execution(self): + if self.execution_command is not None: + if not self.execution_command or any( + not arg for arg in self.execution_command + ): + raise ValueError("execution_command must contain nonempty arguments") + if self.execution_tcp is not None: + raise ValueError("choose execution_command or execution_tcp") + return self + + @property + def split_execution(self) -> bool: + return self.execution_command is not None or self.execution_tcp is not None + @property def resolved_append_to_system_prompt(self) -> str | None: """The append for this engine's role in the session tree. diff --git a/src/rlm/engine.py b/src/rlm/engine.py index 8e9ebee..9139d13 100644 --- a/src/rlm/engine.py +++ b/src/rlm/engine.py @@ -201,7 +201,7 @@ def __init__( self.runtime_config = runtime_config config = self.runtime_config self.model = config.model - self.cwd = cwd or os.getcwd() + self.cwd = config.execution_cwd or cwd or os.getcwd() self.exec_timeout = config.policy.exec_timeout self.max_total_turns = config.policy.max_total_turns self.max_tool_output_bytes = config.policy.max_tool_output_bytes @@ -520,6 +520,10 @@ async def _start(self, prompt: str) -> None: self._active_tools = get_active_builtin_tools( self.exec_timeout, self.builtin_tools ) + if self.runtime_config.split_execution: + from rlm.execution import validate_tools + + validate_tools(tool.name for tool in self._active_tools) self._active_tool_schemas = [tool.schema() for tool in self._active_tools] if self.compaction and self.summarize_at_tokens is None: @@ -579,7 +583,17 @@ async def _start(self, prompt: str) -> None: self._owns_supervisor = False raise - self._repl = IPythonREPL( + repl_type = IPythonREPL + repl_options = {} + if self.runtime_config.split_execution: + from rlm.execution import RemoteREPL + + repl_type = RemoteREPL + repl_options["command"] = self.runtime_config.execution_command + repl_options["tcp"] = self.runtime_config.execution_tcp + repl_options["tools"] = [tool.name for tool in self._active_tools] + self._repl = repl_type( + **repl_options, cwd=self.cwd, session=self.session, kernel_env=self.kernel_env, @@ -602,6 +616,14 @@ async def _start(self, prompt: str) -> None: cancelled = True if cancelled: raise asyncio.CancelledError + if ( + self.runtime_config.split_execution + and self.depth == 0 + and self._supervisor is not None + ): + # The root's worker lives for the session, so the session's jobs and + # services run through it, wherever an agent started them. + self._supervisor.attach_host(self._repl.host_call) system_prompt = self._load_system_prompt(self._active_tools) @@ -924,9 +946,14 @@ async def _run_loop(self) -> RLMResult: if scope_id is not None: repl.set_broker_scope(scope_id) run = repl.run_in_thread if repl is not None else asyncio.to_thread - tool_task = asyncio.create_task( - run(tool.execute, tool_args, self._tool_context(messages)) - ) + if self.runtime_config.split_execution: + tool_task = asyncio.create_task( + run(repl.execute_tool, tool.name, tool_args) + ) + else: + tool_task = asyncio.create_task( + run(tool.execute, tool_args, self._tool_context(messages)) + ) try: tool_result = await asyncio.shield(tool_task) except asyncio.CancelledError: @@ -974,7 +1001,9 @@ async def _run_loop(self) -> RLMResult: result = tool_result.content content = truncate_tool_output( - result, self.max_tool_output_bytes or TOOL_OUTPUT_MAX_BYTES + result, + self.max_tool_output_bytes or TOOL_OUTPUT_MAX_BYTES, + tool_result.omitted, ) self.session.log_tool_result( turn, @@ -1378,7 +1407,7 @@ async def _compact_branch( + "\n\n" + summary_text + "\n\nFull conversation history is available in " - + str(self.session.dir / "messages.jsonl") + + str(Path(self._kernel_session_dir()) / "messages.jsonl") + ". Use `from rlm import history; h = await history()` to inspect " "`h.windows[w].messages[i]`, `h.messages[i]`, or `h.user_messages()`. " "Search or read relevant records with Python when the summary lacks context. The log includes failed attempts: prompt_rollback.prompt_id " @@ -1476,12 +1505,17 @@ def execution_snapshot(self) -> dict: } return snapshot + def _kernel_session_dir(self) -> str: + """The session directory as the kernel sees it: in split execution, the + worker's, since the harness's own is on another machine.""" + return getattr(self._repl, "session_dir", None) or str(self.session.dir) + def _load_system_prompt(self, active_tools: list[BuiltinTool]) -> str: return build_system_prompt( self.cwd, discover_skills(self.session.dir), depth=self.depth, - session_dir=str(self.session.dir), + session_dir=self._kernel_session_dir(), allow_recursion=self.depth < self.max_depth, delegation_prompt=self.runtime_config.policy.delegation_prompt, allow_git=self.allow_git, @@ -1501,6 +1535,7 @@ def _load_system_prompt(self, active_tools: list[BuiltinTool]) -> str: ) def _tool_context(self, messages: list[dict]) -> ToolContext: + repl = self._repl if isinstance(self._repl, IPythonREPL) else None return ToolContext( messages=messages, metrics=self._metrics, @@ -1511,6 +1546,8 @@ def _tool_context(self, messages: list[dict]) -> ToolContext: repl=self._repl, state=self._tool_state, cwd=self.cwd, + # A shell command stops with a cancelled turn. + interrupt_event=repl._interrupt_requested if repl else None, ) diff --git a/src/rlm/execution.py b/src/rlm/execution.py new file mode 100644 index 0000000..6d2ad1c --- /dev/null +++ b/src/rlm/execution.py @@ -0,0 +1,523 @@ +"""Tool-call execution sessions over orchestrator-supplied transports. + +The command is supplied by the orchestrator (for example docker exec -i). Only +kernel configuration and generated skills cross into the worker; model credentials, +conversation ownership, and recursive agent loops stay in the calling process. +""" + +from __future__ import annotations + +import concurrent.futures +import json +import socket +import struct +import argparse +import base64 +import hmac +import logging +import os +from dataclasses import asdict +import sys +import tempfile +import threading +import uuid +from pathlib import Path +from types import SimpleNamespace + +from rlm.broker import BrokerEndpoint +from rlm.tools.ipython import IPythonREPL, build_kernel_env +from rlm.tools.base import ToolContext, ToolOutcome +from rlm.tools.output import OutputSize +from rlm.tools.registry import get_builtin_tool +from rlm.types import IpythonExecuted, RLMMetrics, TokenUsage +from rlm.workspace_host import WorkspaceHost +from rlm.execution_transport import ExecutionTransport, StdioTransport, TcpTransport + +logger = logging.getLogger(__name__) + +WORKSPACE_TOOLS = ("ipython", "bash", "edit", "fetch") + + +def validate_tools(names): + unknown = set(names) - set(WORKSPACE_TOOLS) + if unknown: + raise ValueError(f"split execution does not support tools: {sorted(unknown)}") + + +MAX_FRAME = 16 * 1024 * 1024 +HISTORY_CHUNK = 4 * 1024 * 1024 +"""Raw bytes of the session ledger per frame; base64 keeps a chunk well under +MAX_FRAME, and a long session's ledger syncs as several.""" + + +def _read(stream, size): + chunks = bytearray() + while len(chunks) < size: + chunk = stream.read(size - len(chunks)) + if not chunk: + raise EOFError("execution worker disconnected") + chunks.extend(chunk) + return bytes(chunks) + + +def _receive(stream): + size = struct.unpack("!I", _read(stream, 4))[0] + if size > MAX_FRAME: + raise ValueError("execution frame exceeds 16 MiB") + value = json.loads(_read(stream, size)) + if not isinstance(value, dict): + raise ValueError("execution frame must be an object") + return value + + +def _send(stream, value): + data = json.dumps(value).encode() + if len(data) > MAX_FRAME: + raise ValueError("execution frame exceeds 16 MiB") + stream.write(struct.pack("!I", len(data)) + data) + stream.flush() + + +class _Peer: + """Bidirectional RPC: callbacks must keep flowing during a running cell.""" + + def __init__(self, reader, writer, handler): + self.reader, self.writer, self.handler = reader, writer, handler + self.lock = threading.Lock() + self.pending = {} + self.closed = threading.Event() + self.error = None + self.slots = threading.BoundedSemaphore(32) + self.thread = threading.Thread(target=self._listen, daemon=True) + self.thread.start() + + def send(self, message): + with self.lock: + if self.closed.is_set(): + raise RuntimeError("execution connection closed") from self.error + _send(self.writer, message) + + def call(self, method, params=None, timeout=None): + key = uuid.uuid4().hex + future = concurrent.futures.Future() + with self.lock: + if self.closed.is_set(): + raise RuntimeError("execution connection closed") from self.error + self.pending[key] = future + try: + self.send({"id": key, "method": method, "params": params}) + return future.result(timeout=timeout) + finally: + with self.lock: + self.pending.pop(key, None) + + def _answer(self, message): + try: + result = self.handler(message["method"], message.get("params")) + reply = {"id": message["id"], "result": result} + except Exception as exc: + reply = {"id": message["id"], "error": str(exc)} + try: + try: + self.send(reply) + except ValueError as exc: + self.send({"id": message["id"], "error": str(exc)}) + except (OSError, RuntimeError): + pass # the caller disconnected while its operation was settling + finally: + self.slots.release() + + def _listen(self): + try: + while True: + message = _receive(self.reader) + if "method" in message: + if not self.slots.acquire(blocking=False): + self.send( + { + "id": message["id"], + "error": "execution request limit reached", + } + ) + continue + threading.Thread( + target=self._answer, args=(message,), daemon=True + ).start() + else: + with self.lock: + future = self.pending.get(message["id"]) + if future is not None: + if "error" in message: + future.set_exception(RuntimeError(message["error"])) + else: + future.set_result(message["result"]) + except Exception as exc: + with self.lock: + self.error = exc + self.closed.set() + for future in self.pending.values(): + if not future.done(): + future.set_exception(exc) + + +class RemoteREPL(IPythonREPL): + """The IPython execution-session interface backed by a separate worker.""" + + def __init__( + self, + *, + command=None, + tcp=None, + transport: ExecutionTransport | None = None, + tools=("ipython",), + **kwargs, + ): + super().__init__(**kwargs) + if sum(option is not None for option in (command, tcp, transport)) != 1: + raise ValueError("choose one execution transport") + validate_tools(tools) + self.tools = list(tools) + self.transport = transport or ( + TcpTransport(tcp) if tcp else StdioTransport(command) + ) + self._peer = None + self.session_dir = None + self._history_sent = 0 + # Interrupts and teardown run off the caller's thread, which can be the + # supervisor's event loop, in order and without waiting on a slow worker. + self._control = concurrent.futures.ThreadPoolExecutor( + 1, thread_name_prefix="rlm-control" + ) + self._control_done: concurrent.futures.Future | None = None + self._control_closed = False + + def _control_call(self, method, timeout): + if self._control_closed: + return + peer = self._peer + + def call(): + if peer is not None and not peer.closed.is_set(): + try: + peer.call(method, timeout=timeout) + except Exception: + logger.warning("execution worker %s failed", method, exc_info=True) + + self._control_done = self._control.submit(call) + + def _broker(self, method, request): + if method != "broker" or self.broker_endpoint is None: + raise PermissionError("unknown execution callback") + if request.get("capability") != self.broker_endpoint.capability: + raise PermissionError("invalid execution capability") + with socket.socket(socket.AF_UNIX) as connection: + connection.connect(self.broker_endpoint.socket_path) + with connection.makefile("rwb", buffering=0) as stream: + _send(stream, request) + return _receive(stream) + + def start(self): + reader, writer = self.transport.connect() + self._peer = _Peer(reader, writer, self._broker) + self._history_sent = 0 + modules = { + path.name: path.read_text() for path in self.session.dir.glob("*.py") + } + self.session_dir = self._peer.call( + "start", + { + "tools": self.tools, + "cwd": self.cwd, + "kernel_env": self.kernel_env, + "depth": self.depth, + "max_depth": self.max_depth, + "exec_timeout": self.exec_timeout, + "allow_git": self.allow_git, + "capability": self.broker_endpoint.capability + if self.broker_endpoint + else None, + "modules": modules, + }, + timeout=120, + ) + + def execute_tool(self, name, args): + if self._control_done is not None: + # A call must not overtake a pending interrupt or its release. + self._control_done.result() + # The ledger only grows: the worker gets what it hasn't seen, in chunks + # that fit a frame however long the session runs. + with open(self.session.dir / "messages.jsonl", "rb") as ledger: + ledger.seek(self._history_sent) + new = ledger.read() + *head, last = [ + new[start : start + HISTORY_CHUNK] + for start in range(0, len(new), HISTORY_CHUNK) + ] or [b""] + for chunk in head: + self._peer.call("history", self._history_chunk(chunk)) + self._history_sent += len(chunk) + result = self._peer.call( + "tool", + { + "name": name, + "args": args, + "scope_id": self._scope_id, + "history": self._history_chunk(last), + }, + ) + self._history_sent += len(last) + self._recovery_notices.extend(result["notices"]) + with open(self.session.dir / "programmatic_tool_calls.jsonl", "ab") as log: + log.write(base64.b64decode(result["tool_log"])) + return ToolOutcome( + content=result["output"], + metric_events=[ + IpythonExecuted(**event) for event in result["metric_events"] + ], + omitted=result["omitted"] and OutputSize(**result["omitted"]), + ) + + def _history_chunk(self, data: bytes) -> dict: + return {"offset": self._history_sent, "data": base64.b64encode(data).decode()} + + def execute(self, code, timeout=None): + outcome = self.execute_tool("ipython", {"code": code, "timeout": timeout}) + self.omitted = outcome.omitted + return outcome.content + + def host_call(self, method, params): + """Call the worker's WorkspaceHost: the supervisor's jobs, services and path + watches in the workspace.""" + return self._peer.call(method, params) + + def interrupt(self): + self._control_call("interrupt", 5) + + def finish_interrupt(self): + self._control_call("finish_interrupt", 5) + + def shutdown(self): + def close(): + try: + if self._peer is not None and not self._peer.closed.is_set(): + self._peer.call("close", timeout=30) + except Exception: + logger.warning("execution worker close failed", exc_info=True) + finally: + self.transport.close() + self._executor.shutdown(wait=False) + + if self._control_closed: + return + self._control_closed = True + self._control.submit(close) + self._control.shutdown(wait=False) + + +class _Worker: + def __init__(self): + self.repl = None + self.tools = [] + self.context = None + self.peer = None + self.directory = tempfile.TemporaryDirectory(prefix="rlm-execution-") + self.session_dir = Path(self.directory.name) + self.listener = None + self.closed = threading.Event() + self._tool_log_sent = 0 + self.host = None + + def _accept(self): + while not self.closed.is_set(): + try: + connection, _ = self.listener.accept() + except OSError: + return + threading.Thread( + target=self._proxy, args=(connection,), daemon=True + ).start() + + def _proxy(self, connection): + with connection, connection.makefile("rwb", buffering=0) as stream: + try: + request = _receive(stream) + result = self.peer.call("broker", request) + except Exception as exc: + result = {"error": str(exc)} + try: + _send(stream, result) + except OSError: + pass # interrupted kernels close their pending callbacks + + def handle(self, method, params): + if method == "start": + if self.repl is not None: + raise ValueError("execution session already started") + params = dict(params) + self.tools = params.pop("tools") + validate_tools(self.tools) + modules = params.pop("modules") + for name, source in modules.items(): + if Path(name).name != name or not name.endswith(".py"): + raise ValueError("invalid generated skill name") + (self.session_dir / name).write_text(source) + capability = params.pop("capability") + endpoint = None + if capability: + path = str(self.session_dir / "broker.sock") + self.listener = socket.socket(socket.AF_UNIX) + self.listener.bind(path) + self.listener.listen(32) + threading.Thread(target=self._accept, daemon=True).start() + endpoint = BrokerEndpoint(path, capability) + self.repl = IPythonREPL( + session=SimpleNamespace(dir=self.session_dir), + broker_endpoint=endpoint, + **params, + ) + self.context = ToolContext( + messages=[], + metrics=RLMMetrics(), + total_usage=TokenUsage(), + last_prompt_tokens=0, + exec_timeout=self.repl.exec_timeout or 300, + allow_git=self.repl.allow_git, + repl=self.repl, + cwd=self.repl.cwd, + env=build_kernel_env(self.repl.kernel_env), + interrupt_event=self.repl._interrupt_requested, + ) + # All Jupyter operations use the same dedicated thread. + self.repl._executor.submit(self.repl.start).result() + return str(self.session_dir) + if self.repl is None: + raise ValueError("execution session has not started") + if method == "history": + self._write_history(params) + return {} + if method.split(".", 1)[0] in ("host", "jobs", "services", "watch"): + if self.host is None: + self.host = WorkspaceHost(self.session_dir) + return self.host.handle(method, params) + if method == "tool": + + def execute(): + self._write_history(params["history"]) + self.repl.set_broker_scope(params["scope_id"]) + tool = get_builtin_tool(params["name"], self.tools) + if tool is None: + raise ValueError("tool is not enabled for this execution session") + outcome = tool.execute(params["args"], self.context) + return { + "output": outcome.content, + "omitted": outcome.omitted and asdict(outcome.omitted), + "metric_events": [asdict(event) for event in outcome.metric_events], + "notices": self.repl.take_recovery_notices(), + "tool_log": base64.b64encode(self._new_tool_log()).decode(), + } + + return self.repl._executor.submit(execute).result() + if method == "interrupt": + self.repl.interrupt() + elif method == "finish_interrupt": + self.repl.finish_interrupt() + elif method == "close": + self.repl.shutdown() + self.repl = None + else: + raise ValueError(f"unknown execution method: {method}") + + def _write_history(self, chunk): + """Write a chunk of the harness's ledger at its offset. Rewriting the same + bytes is harmless, so a retried chunk can't duplicate history.""" + path = self.session_dir / "messages.jsonl" + path.touch() + with path.open("r+b") as ledger: + if chunk["offset"] > ledger.seek(0, os.SEEK_END): + raise ValueError("execution history is out of sync") + ledger.seek(chunk["offset"]) + ledger.write(base64.b64decode(chunk["data"])) + ledger.truncate() + + def _new_tool_log(self) -> bytes: + """The programmatic tool calls logged since the last tool call returned.""" + path = self.session_dir / "programmatic_tool_calls.jsonl" + if not path.exists(): + return b"" + with path.open("rb") as log: + log.seek(self._tool_log_sent) + new = log.read() + self._tool_log_sent += len(new) + return new + + def close(self): + self.closed.set() + if self.listener is not None: + self.listener.close() + if self.host is not None: + self.host.close() + if self.repl is not None: + self.repl.interrupt() + self.repl.shutdown() + self.directory.cleanup() + + +def _serve(reader, writer): + worker = _Worker() + worker.peer = _Peer(reader, writer, worker.handle) + try: + worker.peer.closed.wait() + finally: + worker.close() + + +def _serve_tcp(connection, token, slots): + try: + with connection, connection.makefile("rwb", buffering=0) as stream: + connection.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1) + connection.settimeout(10) + supplied = _receive(stream).get("token") + if not isinstance(supplied, str) or not hmac.compare_digest( + supplied, token + ): + _send(stream, {"authenticated": False}) + return + _send(stream, {"authenticated": True}) + connection.settimeout(None) + _serve(stream, stream) + except (OSError, EOFError, ValueError): + pass # unauthenticated or disconnected clients have no remaining session + finally: + slots.release() + + +def main(): + parser = argparse.ArgumentParser(description="Tool execution worker") + parser.add_argument("--listen", metavar="HOST") + parser.add_argument("--port", type=int, default=8765) + args = parser.parse_args() + output = sys.stdout.buffer + sys.stdout = sys.stderr # dependency diagnostics must not corrupt the wire + if args.listen is None: + _serve(sys.stdin.buffer, output) + return + token = os.environ.pop("RLM_EXECUTION_TOKEN", "") + if len(token) < 32: + parser.error("TCP workers require RLM_EXECUTION_TOKEN (at least 32 characters)") + slots = threading.BoundedSemaphore(32) + with socket.create_server((args.listen, args.port)) as listener: + # With --port 0 the kernel picks the port; the orchestrator reads it here. + print(f"listening on {args.listen}:{listener.getsockname()[1]}", flush=True) + while True: + connection, _ = listener.accept() + if not slots.acquire(blocking=False): + connection.close() + continue + threading.Thread( + target=_serve_tcp, args=(connection, token, slots), daemon=True + ).start() + + +if __name__ == "__main__": + main() diff --git a/src/rlm/execution_transport.py b/src/rlm/execution_transport.py new file mode 100644 index 0000000..f7db163 --- /dev/null +++ b/src/rlm/execution_transport.py @@ -0,0 +1,83 @@ +"""Byte-stream transports for tool execution, independent of provisioning.""" + +from __future__ import annotations + +import socket +import subprocess +from typing import BinaryIO, Protocol + +from rlm.config import TcpExecutionConfig + + +class ExecutionTransport(Protocol): + def connect(self) -> tuple[BinaryIO, BinaryIO]: ... + + def close(self) -> None: ... + + +class StdioTransport: + def __init__(self, command): + self.command = command + self.process = None + + def connect(self): + self.process = subprocess.Popen( + self.command, + stdin=subprocess.PIPE, + stdout=subprocess.PIPE, + stderr=None, + start_new_session=True, + ) + return self.process.stdout, self.process.stdin + + def close(self): + process, self.process = self.process, None + if process is not None: + process.stdin.close() + try: + process.wait(timeout=10) + except subprocess.TimeoutExpired: + process.kill() + process.wait() + process.stdout.close() + + +class TcpTransport: + """One execution session per connection on a trusted private network.""" + + def __init__(self, endpoint: TcpExecutionConfig): + self.endpoint = endpoint + self.connection = None + self.stream = None + + def connect(self): + from rlm.execution import _receive, _send + + self.connection = socket.create_connection( + (self.endpoint.host, self.endpoint.port), timeout=10 + ) + # Each frame is one write, and calls overlap (the event long-poll, kernel + # callbacks): Nagle would hold a frame until the previous one is acknowledged. + self.connection.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1) + self.stream = self.connection.makefile("rwb", buffering=0) + try: + _send(self.stream, {"token": self.endpoint.token}) + if _receive(self.stream) != {"authenticated": True}: + raise PermissionError("execution worker authentication failed") + self.connection.settimeout(None) + except BaseException: + self.close() + raise + return self.stream, self.stream + + def close(self): + if self.connection is not None: + try: + self.connection.shutdown(socket.SHUT_RDWR) + except OSError: + pass # the worker may already have disconnected + self.connection.close() + self.connection = None + if self.stream is not None: + self.stream.close() + self.stream = None diff --git a/src/rlm/history.py b/src/rlm/history.py index a118729..65c8a08 100644 --- a/src/rlm/history.py +++ b/src/rlm/history.py @@ -9,6 +9,11 @@ from pathlib import Path +HISTORY_PAGE_BYTES = 8 * 1024 * 1024 +"""JSON bytes of ledger records one page of another agent's history holds, beyond its +first record: half a frame, so a page always crosses the broker and execution wire.""" + + def read_records(path: Path) -> Iterator[dict]: with path.open("rb") as stream: for line in stream: @@ -18,6 +23,27 @@ def read_records(path: Path) -> Iterator[dict]: yield json.loads(line) +def read_page( + path: Path, offset: int, limit: int = HISTORY_PAGE_BYTES +) -> tuple[list[dict], int | None]: + """The complete records from byte `offset` that fit `limit` bytes of JSON, at least + one, and the offset of the record after them, or None at the ledger's end.""" + records: list[dict] = [] + size = 0 + with path.open("rb") as stream: + stream.seek(offset) + for line in stream: + if not line.endswith(b"\n"): + break + record = json.loads(line) + size += len(json.dumps(record)) + if records and size > limit: + return records, offset + records.append(record) + offset += len(line) + return records, None + + @dataclass class ContextWindow: index: int @@ -35,7 +61,18 @@ class History: def __init__(self, session_dir: str | Path): self.session_dir = Path(session_dir).resolve() - self.events = list(read_records(self.session_dir / "messages.jsonl")) + self._load(list(read_records(self.session_dir / "messages.jsonl"))) + + @classmethod + def from_events(cls, session_dir: str | Path, events: list[dict]) -> History: + """Construct a history snapshot received from the owning supervisor.""" + history = cls.__new__(cls) + history.session_dir = Path(session_dir) + history._load(events) + return history + + def _load(self, events: list[dict]) -> None: + self.events = events self.messages: list[dict] = [] self.windows: list[ContextWindow] = [] for event in self.events: diff --git a/src/rlm/shell_jobs.py b/src/rlm/shell_jobs.py index e1b5f42..f254443 100644 --- a/src/rlm/shell_jobs.py +++ b/src/rlm/shell_jobs.py @@ -123,7 +123,7 @@ def read(self, job: JobRecord, cursor: int, max_bytes: int) -> dict: # A cursor past the retained output is not an error: it returns an empty # chunk positioned at the end, so "read until done" loops terminate. cursor = min(cursor, job.info.output_bytes) - with open(job.info.output_path, "rb") as stream: + with open(self.output_file(job), "rb") as stream: stream.seek(cursor) data = stream.read(max_bytes) return { @@ -134,6 +134,10 @@ def read(self, job: JobRecord, cursor: int, max_bytes: int) -> dict: "truncated": job.info.output_truncated, } + def output_file(self, job: JobRecord) -> str: + """Where this supervisor reads the job's output.""" + return job.info.output_path + def run_text(self, job: JobRecord) -> dict: """Text for a finished run(): the whole output when it fits RUN_TEXT_BYTES, otherwise its first and last halves around an explicit gap marker (test diff --git a/src/rlm/subscriptions.py b/src/rlm/subscriptions.py index e412921..bac713a 100644 --- a/src/rlm/subscriptions.py +++ b/src/rlm/subscriptions.py @@ -6,7 +6,7 @@ import json import time import uuid -from collections.abc import Callable +from collections.abc import AsyncIterator, Awaitable, Callable from dataclasses import asdict, dataclass, field, replace from pathlib import Path @@ -31,6 +31,36 @@ class Subscription: thresholds: dict = field(default_factory=dict) +PathChanges = AsyncIterator[tuple[set[str], bool]] +"""Changes under a watched path: the changed paths and whether the path still exists.""" + + +async def watch_path(target: Path, recursive: bool) -> tuple[Path, PathChanges]: + """Watch `target` on this machine; returns the resolved path and its changes.""" + target = target.resolve() + if not target.exists(): + raise FileNotFoundError(target) + return target, _changes(target, recursive) + + +async def _changes(target: Path, recursive: bool) -> PathChanges: + is_file = target.is_file() + async for changes in awatch( + target.parent if is_file else target, + recursive=recursive if not is_file else False, + watch_filter=None, + debounce=1, + step=1, + rust_timeout=100, + yield_on_timeout=True, + ignore_permission_denied=False, + ): + yield ( + {name for _, name in changes if not is_file or name == str(target)}, + target.exists(), + ) + + class Subscriptions: def __init__( self, @@ -40,6 +70,10 @@ def __init__( self.items: dict[str, Subscription] = {} self.publish = publish self.record = record + self.watch: Callable[[Path, bool], Awaitable[tuple[Path, PathChanges]]] = ( + watch_path + ) + """Where path watches run: this machine, or the execution workspace.""" def register( self, @@ -115,10 +149,13 @@ def progress( self._fail(sub, str(exc)) async def path(self, owner_id: str, target: Path, recursive: bool) -> Subscription: - if not target.exists(): - raise FileNotFoundError(target) - sub = self.register(owner_id, "path", str(target), recursive=recursive) - sub.task = asyncio.create_task(self._watch_path(sub, target)) + target, changes = await self.watch(target, recursive) + try: + sub = self.register(owner_id, "path", str(target), recursive=recursive) + except BaseException: + await changes.aclose() + raise + sub.task = asyncio.create_task(self._watch_path(sub, changes)) try: await sub.ready.wait() except BaseException: @@ -193,26 +230,12 @@ def _fail(self, sub: Subscription, error: str) -> None: self.record(sub) self.publish(sub, "watch.failed", {"error": error}) - async def _watch_path(self, sub: Subscription, target: Path) -> None: - is_file = target.is_file() - watched = target.parent if is_file else target + async def _watch_path(self, sub: Subscription, changes: PathChanges) -> None: try: - async for changes in awatch( - watched, - recursive=sub.info.recursive if not is_file else False, - watch_filter=None, - debounce=1, - step=1, - rust_timeout=100, - yield_on_timeout=True, - ignore_permission_denied=False, - ): + async for paths, exists in changes: sub.ready.set() if sub.info.status != "active": return - paths = { - name for _, name in changes if not is_file or name == str(target) - } if paths: pending = sub.pending or {"paths": set(), "truncated": False} for name in sorted(paths): @@ -225,7 +248,7 @@ async def _watch_path(self, sub: Subscription, target: Path) -> None: pending["truncated"] = True sub.pending = pending self._schedule(sub) - if not target.exists(): + if not exists: self._flush(sub) if sub.info.status == "active": self._fail( @@ -239,6 +262,7 @@ async def _watch_path(self, sub: Subscription, target: Path) -> None: self._fail(sub, str(exc)) finally: sub.ready.set() + await changes.aclose() async def cancel(self, sub: Subscription) -> dict: if sub.info.status == "active": diff --git a/src/rlm/supervisor.py b/src/rlm/supervisor.py index 945859e..c8df178 100644 --- a/src/rlm/supervisor.py +++ b/src/rlm/supervisor.py @@ -47,6 +47,7 @@ from rlm.service_manager import ServiceManager from rlm.subscriptions import Subscription, Subscriptions from rlm.tools.ipython import build_kernel_env +from rlm.workspace_host import RemoteHost, RemoteShellJobs from rlm.tools.git_block import find_blocked_command, refusal from rlm.session import Session from rlm.skills.search import run_with_api_key as run_search @@ -228,6 +229,7 @@ def __init__( self._services = ServiceManager( Path(os.environ.get("RLM_SERVICES_DIR", "/tmp/rlm-services")) ) + self._host: RemoteHost | None = None self._root_config = runtime_config if "search" in runtime_config.skills: capability = secrets.token_urlsafe(24) @@ -835,9 +837,8 @@ async def _watch_operation(self, parent: _Invocation, request: dict) -> Any: completed=job.finished is not None, ) elif op == "watch.path": - target = (Path(parent.cwd) / request["path"]).resolve() sub = await self._subscriptions.path( - parent.id, target, request["recursive"] + parent.id, Path(parent.cwd) / request["path"], request["recursive"] ) elif op == "watch.list": return [ @@ -948,16 +949,10 @@ async def _shell_operation(self, parent: _Invocation, request: dict) -> Any: if yield_after is None else min(float(yield_after), RUN_BLOCK_MAX_SECONDS) ) - info = self._shell_jobs.start( + job_options = dict( owner_id=parent.id, command=command, - cwd=str(cwd.resolve()), directory=parent.session.dir, - env={ - **build_kernel_env(dict(parent.runtime_config.kernel_env)), - **parent.shell_env, - **(request.get("env") or {}), - }, source_request_id=self._scopes[request["scope_id"]].request_id, # yield_after= bounds how long run() blocks; timeout= is the kill deadline. timeout=request.get("timeout"), @@ -965,6 +960,17 @@ async def _shell_operation(self, parent: _Invocation, request: dict) -> Any: # handed back while still running posts shell.completed when it ends. notify=block == 0, ) + if isinstance(self._shell_jobs, RemoteShellJobs): + # The workspace resolves the directory and environment it runs in. + info = await self._shell_jobs.start_remote( + cwd=str(cwd), env=self._env_layers(parent, request), **job_options + ) + else: + info = self._shell_jobs.start( + cwd=str(cwd.resolve()), + env=self._command_env(parent, request), + **job_options, + ) job = self._shell_jobs.get(parent.id, info["id"]) if block == 0: return self._running_payload( @@ -1051,6 +1057,22 @@ async def _shell_operation(self, parent: _Invocation, request: dict) -> Any: return await self._shell_jobs.cancel(job) return job.snapshot() + @staticmethod + def _env_layers(parent: _Invocation, request: dict) -> dict: + return { + "kernel_env": dict(parent.runtime_config.kernel_env), + "shell_env": dict(parent.shell_env), + "env": dict(request.get("env") or {}), + } + + def _command_env(self, parent: _Invocation, request: dict) -> dict[str, str]: + layers = self._env_layers(parent, request) + return { + **build_kernel_env(layers["kernel_env"]), + **layers["shell_env"], + **layers["env"], + } + async def _services_operation(self, parent: _Invocation, request: dict) -> Any: op = request["op"] if op == "services.start": @@ -1063,17 +1085,38 @@ async def _services_operation(self, parent: _Invocation, request: dict) -> Any: if blocked: raise PermissionError(refusal(blocked)) cwd = Path(parent.cwd) / (request["cwd"] or ".") + if self._host is not None: + return await self._host.call( + op, + { + "name": request["name"], + "command": command, + "cwd": str(cwd), + "env": self._env_layers(parent, request), + "owner_id": parent.id, + }, + ) return self._services.start( name=request["name"], command=command, cwd=str(cwd.resolve()), - env={ - **build_kernel_env(dict(parent.runtime_config.kernel_env)), - **parent.shell_env, - **(request.get("env") or {}), - }, + env=self._command_env(parent, request), owner_id=parent.id, ) + if self._host is not None and op in ( + "services.list", + "services.get", + "services.tail", + "services.stop", + ): + return await self._host.call( + op, + { + key: request[key] + for key in ("name", "lines", "grace") + if key in request + }, + ) if op == "services.list": return self._services.list() if op == "services.get": @@ -1194,6 +1237,13 @@ async def _agent_operation(self, request: dict) -> Any: ancestor = self._parents[ancestor] return agents child = self._child(parent, request.get("name_or_id", request.get("agent_id"))) + if op == "agent.history": + from rlm.history import read_page + + records, offset = read_page( + child.session.dir / "messages.jsonl", request["offset"] + ) + return {"records": records, "next": offset} if op in {"agent.send", "agent.steer"}: if child.status in {"completed", "failed", "cancelled"}: raise RuntimeError("agent is no longer active") @@ -1500,6 +1550,16 @@ async def _handle_connection( pass self._connection_writers.discard(writer) + def attach_host(self, call: Callable[[str, dict], Any]) -> None: + """Run shell jobs, services and path watches in the execution workspace, through + `call` on the root agent's worker, as split execution needs.""" + self._host = RemoteHost(call) + self._shell_jobs = RemoteShellJobs( + self._host, self._publish_job, self._publish_job_output + ) + self._host.follow(self._shell_jobs) + self._subscriptions.watch = self._host.watch + async def aclose(self) -> None: if self._closed and self._close_task is None: return @@ -1537,6 +1597,8 @@ async def _aclose_impl(self) -> None: return_exceptions=True, ) ) + if self._host is not None: + await self._host.close() tasks = list(self._tasks) for task in tasks: task.cancel() diff --git a/src/rlm/tools/base.py b/src/rlm/tools/base.py index 295d6ff..6298a33 100644 --- a/src/rlm/tools/base.py +++ b/src/rlm/tools/base.py @@ -4,7 +4,9 @@ from dataclasses import dataclass, field from typing import Any, Protocol +from threading import Event +from rlm.tools.output import OutputSize from rlm.types import BuiltinMetricEvent, RLMMetrics, TokenUsage @@ -14,6 +16,9 @@ class ToolOutcome: content: str metric_events: list[BuiltinMetricEvent] = field(default_factory=list) + omitted: OutputSize | None = None + """How much more the whole output measures than `content`, when `content` kept + only its start and end.""" @dataclass @@ -29,6 +34,8 @@ class ToolContext: repl: Any | None = None state: dict[str, Any] = field(default_factory=dict) cwd: str = "" + env: dict[str, str] | None = None + interrupt_event: Event | None = None class BuiltinTool(Protocol): diff --git a/src/rlm/tools/bash.py b/src/rlm/tools/bash.py index 945209d..90537a9 100644 --- a/src/rlm/tools/bash.py +++ b/src/rlm/tools/bash.py @@ -7,12 +7,28 @@ from __future__ import annotations +import codecs +import io import os +import selectors +import signal +import sys +import time +import contextlib import subprocess +from threading import Event from typing import Any from rlm.tools.base import ToolContext, ToolOutcome from rlm.tools.git_block import find_blocked_command, refusal +from rlm.tools.output import OUTPUT_KEEP_BYTES, OutputCapture, OutputSize + +_OWN_GROUP = ( + {"process_group": 0} if sys.version_info >= (3, 11) else {"preexec_fn": os.setpgrp} +) +"""A command leads a process group in its caller's session: a timeout or a cancel stops +the command and what it started, a kernel restart or shutdown stops a command the kernel +ran, and a daemon, in a session of its own, outlives both.""" BASH_SCHEMA = { "type": "function", @@ -61,28 +77,96 @@ def run_bash( timeout: int, cwd: str | None = None, allow_git: bool | None = None, + env: dict[str, str] | None = None, + interrupt_event: Event | None = None, ) -> str: """Guarded ``bash -c`` execution shared by the bash tool and the bash skill.""" + return bash_outcome(command, timeout, cwd, allow_git, env, interrupt_event).content + + +def bash_outcome( + command: str, + timeout: int, + cwd: str | None = None, + allow_git: bool | None = None, + env: dict[str, str] | None = None, + interrupt_event: Event | None = None, +) -> ToolOutcome: + """`run_bash`, with how much of a long output its content left out.""" if not isinstance(command, str) or not command.strip(): - return "Error: empty command" + return ToolOutcome(content="Error: empty command") blocked = find_blocked_command(command, allow_git=allow_git) if blocked: - return refusal(blocked) - try: - proc = subprocess.run( + return ToolOutcome(content=refusal(blocked)) + deadline = time.monotonic() + timeout + streams = ("stdout", "stderr") + # stdout and stderr each keep half of one result's bound. + captures = {name: OutputCapture(OUTPUT_KEEP_BYTES // 2) for name in streams} + decoders = { + name: io.IncrementalNewlineDecoder( + codecs.getincrementaldecoder("utf-8")(errors="replace"), translate=True + ) + for name in streams + } + reason = "" + with ( + subprocess.Popen( ["bash", "-c", command], - capture_output=True, - text=True, - errors="replace", - timeout=timeout, + # The caller's stdin may be its terminal, where a background group stops on + # a read, or its own protocol stream; a command reads end-of-file instead. + stdin=subprocess.DEVNULL, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, cwd=cwd or None, - ) - except subprocess.TimeoutExpired: - return f"Error: command timed out after {timeout}s" - out = proc.stdout + (("\n" + proc.stderr) if proc.stderr else "") - if proc.returncode != 0: + env=env, + **_OWN_GROUP, + ) as proc, + selectors.DefaultSelector() as selector, + ): + selector.register(proc.stdout, selectors.EVENT_READ, "stdout") + selector.register(proc.stderr, selectors.EVENT_READ, "stderr") + try: + while selector.get_map(): + if interrupt_event is not None and interrupt_event.is_set(): + reason = "[command interrupted]" + break + remaining = deadline - time.monotonic() + if remaining <= 0: + reason = f"Error: command timed out after {timeout}s" + break + for key, _ in selector.select(min(remaining, 0.1)): + data = os.read(key.fileobj.fileno(), 65536) + if not data: + selector.unregister(key.fileobj) + continue + captures[key.data].add(decoders[key.data].decode(data)) + if not reason: + try: + proc.wait(timeout=max(0.01, deadline - time.monotonic())) + except subprocess.TimeoutExpired: + reason = f"Error: command timed out after {timeout}s" + finally: + # Once bash itself has exited, what still holds its output is a job it put in + # the background, and that job outlives the call. + if proc.poll() is None: + with contextlib.suppress(ProcessLookupError): + os.killpg(proc.pid, signal.SIGKILL) + proc.wait() + for name in streams: + captures[name].add(decoders[name].decode(b"", final=True)) + (stdout, stdout_omitted), (stderr, stderr_omitted) = ( + captures[name].result() for name in streams + ) + out = stdout + (("\n" + stderr) if stderr else "") + if reason: + out += "\n" + reason + elif proc.returncode != 0: out += f"\n[exit code {proc.returncode}]" - return out.strip() or "(no output)" + omitted = [size for size in (stdout_omitted, stderr_omitted) if size is not None] + return ToolOutcome( + content=out.strip() or "(no output)", + omitted=sum(omitted, OutputSize()) if omitted else None, + ) class BashTool: @@ -94,13 +178,13 @@ def schema(self) -> dict[str, Any]: return BASH_SCHEMA def execute(self, args: dict[str, Any], context: ToolContext) -> ToolOutcome: - return ToolOutcome( - content=run_bash( - args.get("command", ""), - context.exec_timeout, - cwd=context.cwd, - allow_git=context.allow_git, - ) + return bash_outcome( + args.get("command", ""), + context.exec_timeout, + cwd=context.cwd, + allow_git=context.allow_git, + env=context.env, + interrupt_event=context.interrupt_event, ) diff --git a/src/rlm/tools/ipython.py b/src/rlm/tools/ipython.py index 8c15216..374598a 100644 --- a/src/rlm/tools/ipython.py +++ b/src/rlm/tools/ipython.py @@ -3,12 +3,14 @@ from __future__ import annotations import asyncio +import contextlib import copy import functools import os from queue import Empty import re import shutil +import signal import subprocess import sys import tempfile @@ -21,6 +23,7 @@ from rlm.tools.base import ToolContext, ToolOutcome from rlm.tools.git_block import find_blocked_in_ipython, refusal +from rlm.tools.output import OutputCapture from rlm.tools.skills import discover_skills from rlm.types import IpythonExecuted @@ -237,8 +240,10 @@ def execute(self, args: dict[str, Any], context: ToolContext) -> ToolOutcome: metric_events=metric_events, ) + content = context.repl.execute(code, timeout=timeout) return ToolOutcome( - content=context.repl.execute(code, timeout=timeout), + content=content, + omitted=getattr(context.repl, "omitted", None), metric_events=metric_events, ) @@ -256,7 +261,29 @@ def __init__(self, output: str = ""): MAX_RECOVERY_ATTEMPTS = 3 +def kill_session(session: int) -> None: + """SIGKILL the process group of every process in `session`, read from /proc; a no-op + where there is none. Reparented processes stay in their session, so this also + reaches the commands of a dead kernel, while a daemon, in a session of its own, is + left running. The caller's own group is spared.""" + try: + entries = [entry for entry in os.listdir("/proc") if entry.isdigit()] + except FileNotFoundError: + return + groups = set() + for entry in entries: + with contextlib.suppress(OSError): + if os.getsid(int(entry)) == session: + groups.add(os.getpgid(int(entry))) + groups -= {0, 1, os.getpgrp()} + for group in groups: + with contextlib.suppress(ProcessLookupError, PermissionError): + os.killpg(group, signal.SIGKILL) + + class IPythonREPL: + _kernel_session: int | None = None + """Persistent IPython kernel communicating via Jupyter protocol.""" def __init__( @@ -287,6 +314,9 @@ def __init__( self._recovery_attempts = 0 self._recovery_failed = False self._cell_submitted = False + self.omitted = None + """How much more the last cell's whole output measured than what `execute` + returned, when that kept only its start and end.""" self._recovery_notices: list[str] = [] self._executor = ThreadPoolExecutor( max_workers=1, thread_name_prefix="rlm-repl" @@ -337,6 +367,7 @@ def start(self): stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, ) + self._kernel_session = self._km.provisioner.pid self._kc = self._km.client() self._kc.start_channels() self._kc.wait_for_ready(timeout=30) @@ -522,11 +553,20 @@ def _wait_for_idle(self, msg_id: str, timeout: float) -> bool: ): return True + def _kill_kernel_commands(self): + """Kill what the kernel started that is still attached to it, the way a closing + terminal does: everything in the kernel's session. The kernel leads a session + of its own, and a shell command leads a process group in it.""" + if self._kernel_session: + kill_session(self._kernel_session) + def restart_kernel(self): """Restart the kernel and restore the initial REPL state.""" if self._kc: self._kc.stop_channels() + self._kill_kernel_commands() self._km.restart_kernel(now=True) + self._kernel_session = self._km.provisioner.pid self._kc = self._km.client() self._kc.start_channels() self._kc.wait_for_ready(timeout=30) @@ -555,6 +595,7 @@ def execute(self, code: str, timeout: int | None = None) -> str: """Execute once; a lost kernel is recovered without replaying the cell.""" with self._lock: self._cell_submitted = False + self.omitted = None try: if self._interrupt_requested.is_set(): return "" @@ -592,11 +633,12 @@ def _execute_locked(self, code: str, timeout: int | None) -> str: self._cell_submitted = True deadline = None if timeout is None else time.monotonic() + timeout - outputs: list[str] = [] + capture = OutputCapture() try: while True: if not self._km.is_alive(): - raise _KernelDied("".join(outputs)) + text, self.omitted = capture.result() + raise _KernelDied(text) if self._interrupt_requested.is_set(): self._interrupt_and_recover(msg_id) break @@ -606,7 +648,7 @@ def _execute_locked(self, code: str, timeout: int | None) -> str: remaining = deadline - time.monotonic() if remaining <= 0: self._interrupt_and_recover(msg_id) - outputs.append( + capture.add( f"\n[execution timed out after {timeout}s and was interrupted]" ) break @@ -622,18 +664,20 @@ def _execute_locked(self, code: str, timeout: int | None) -> str: msg_type = msg["msg_type"] content = msg["content"] + text = "" if msg_type == "stream": - outputs.append(content["text"]) + text = content["text"] elif msg_type == "execute_result": text = content.get("data", {}).get("text/plain", "") if text: - outputs.append(text + "\n") + text += "\n" elif msg_type == "error": tb = "\n".join(content.get("traceback", [])) tb = _ANSI_RE.sub("", tb) - outputs.append(tb) + text = tb elif msg_type == "status" and content["execution_state"] == "idle": break + capture.add(text) finally: try: timeout = 0.1 if self._interrupt_requested.is_set() else 5 @@ -642,13 +686,15 @@ def _execute_locked(self, code: str, timeout: int | None) -> str: except Exception: pass - return "".join(outputs) + text, self.omitted = capture.result() + return text def shutdown(self): """Stop the kernel.""" if self._kc: self._kc.stop_channels() self._kc = None + self._kill_kernel_commands() if self._km and self._km.has_kernel: self._km.shutdown_kernel(now=True) if self._km: diff --git a/src/rlm/tools/output.py b/src/rlm/tools/output.py new file mode 100644 index 0000000..5390d4f --- /dev/null +++ b/src/rlm/tools/output.py @@ -0,0 +1,97 @@ +"""Bounded capture of a command's output.""" + +from __future__ import annotations + +from dataclasses import dataclass + +OUTPUT_KEEP_BYTES = 1024 * 1024 +"""The most of one tool result's output that is kept, half from its start and half from +its end. The conversation shows far less of it; the rest stays in the session log, and +a tool result always fits one execution frame.""" + +KEPT_PER_END = OUTPUT_KEEP_BYTES // 4 - 4096 +"""Bytes of a long output's start, and of its end, that a tool result holds at least: +bash keeps a quarter of the bound from each end of each stream, less what a cut at a +character boundary or stripped whitespace takes.""" + + +@dataclass(frozen=True) +class OutputSize: + """Text measured as the conversation's truncation measures it.""" + + nbytes: int = 0 + """UTF-8 bytes.""" + nchars: int = 0 + newlines: int = 0 + + @classmethod + def of(cls, text: str) -> OutputSize: + return cls( + len(text.encode("utf-8", "surrogatepass")), len(text), text.count("\n") + ) + + def __add__(self, other: OutputSize) -> OutputSize: + return OutputSize( + self.nbytes + other.nbytes, + self.nchars + other.nchars, + self.newlines + other.newlines, + ) + + def __sub__(self, other: OutputSize) -> OutputSize: + return OutputSize( + self.nbytes - other.nbytes, + self.nchars - other.nchars, + self.newlines - other.newlines, + ) + + +def _char_start(data: bytes | bytearray, at: int) -> int: + """`at`, moved back to the start of the UTF-8 character it falls in.""" + while 0 < at < len(data) and data[at] & 0xC0 == 0x80: + at -= 1 + return at + + +class OutputCapture: + """Output as it is produced, kept to its first and last `keep // 2` bytes. The middle + is dropped but measured, so a long command runs to its end in bounded memory and the + whole output's size is still known.""" + + def __init__(self, keep: int = OUTPUT_KEEP_BYTES): + self._half = keep // 2 + self._head = bytearray() + self._head_full = False + self._tail = bytearray() + self._dropped = False + self.size = OutputSize() + """The whole output's size.""" + + def add(self, text: str) -> None: + if not text: + return + self.size += OutputSize.of(text) + data = text.encode("utf-8", "surrogatepass") + if not self._head_full: + room = self._half - len(self._head) + if len(data) <= room: + self._head += data + return + cut = _char_start(data, room) + self._head += data[:cut] + self._head_full = True + data = data[cut:] + self._tail += data + if len(self._tail) > self._half: + del self._tail[: _char_start(self._tail, len(self._tail) - self._half)] + self._dropped = True + + def result(self) -> tuple[str, OutputSize | None]: + """The kept text and, when its middle was dropped, how much more the whole + output measures than the text: the middle, less the note that replaces it.""" + head = self._head.decode("utf-8", "surrogatepass") + tail = self._tail.decode("utf-8", "surrogatepass") + if not self._dropped: + return head + tail, None + dropped = self.size.nbytes - len(self._head) - len(self._tail) + text = f"{head}\n[... {dropped} bytes of output omitted ...]\n{tail}" + return text, self.size - OutputSize.of(text) diff --git a/src/rlm/workspace_host.py b/src/rlm/workspace_host.py new file mode 100644 index 0000000..64a6c4c --- /dev/null +++ b/src/rlm/workspace_host.py @@ -0,0 +1,397 @@ +"""Shell jobs, services and path watches under split execution. + +The supervisor stays with the agent loop, but every process the agent starts must run +in the task workspace, beside its kernel. The execution worker hosts them there with the +same job, service and watch code a local supervisor uses (`WorkspaceHost`); the +supervisor keeps its records and reaches the worker through `RemoteHost`. +""" + +from __future__ import annotations + +import asyncio +import base64 +import contextlib +import dataclasses +import os +import threading +import time +import uuid +from collections.abc import Callable +from pathlib import Path +from typing import Any + +from rlm.service_manager import ServiceManager +from rlm.shell import JobInfo +from rlm.shell_jobs import JobRecord, ShellJobs +from rlm.subscriptions import PathChanges, watch_path +from rlm.tools.ipython import build_kernel_env + +EVENT_BYTES = 4 * 1024 * 1024 +"""Job output one events call carries, so its reply stays well under a frame.""" +EVENT_WAIT = 1.0 +"""Seconds an events call waits for something to report.""" +EVENT_LINGER = 0.05 +"""Seconds an events call holds a running job's new output for the job to finish, so +a short command's output and exit cost the supervisor one round trip, not two.""" +_JOB_FIELDS = [field.name for field in dataclasses.fields(JobInfo)] + + +def _env(layers: dict) -> dict[str, str]: + """A command's environment, built where it runs from the supervisor's layers.""" + return { + **build_kernel_env(layers["kernel_env"]), + **layers["shell_env"], + **layers["env"], + } + + +class WorkspaceHost: + """Runs the supervisor's shell jobs, services and path watches in this workspace, + on an event loop of its own beside the worker's kernel.""" + + def __init__(self, directory: Path) -> None: + self.directory = directory + self.loop = asyncio.new_event_loop() + threading.Thread( + target=self.loop.run_forever, name="rlm-host", daemon=True + ).start() + self.jobs = ShellJobs(self._finished, self._output) + self.services = ServiceManager( + Path(os.environ.get("RLM_SERVICES_DIR", "/tmp/rlm-services")) + ) + self.watches: dict[str, asyncio.Task] = {} + self._changed = threading.Condition() + self._jobs_changed: dict[str, bool] = {} # job id -> finished + self._sent: dict[str, int] = {} # job id -> output bytes reported + self._paths: list[dict] = [] + + def handle(self, method: str, params: dict) -> Any: + if method == "host.events": + return self._events(params["wait"]) + if method == "jobs.start": + return self._run(self._start(params)) + if method == "jobs.cancel": + return self._run(self.jobs.cancel(self.jobs.jobs[params["id"]])) + if method == "jobs.close": + return self._run(self.jobs.close(params["owner_id"])) + if method == "services.start": + return self.services.start( + name=params["name"], + command=params["command"], + cwd=os.path.realpath(params["cwd"]), + env=_env(params["env"]), + owner_id=params["owner_id"], + ) + if method == "services.list": + return self.services.list() + if method == "services.get": + return self.services.get(params["name"]) + if method == "services.tail": + return self.services.tail(params["name"], params["lines"]) + if method == "services.stop": + return self._run(self.services.stop(params["name"], params["grace"])) + if method == "watch.path": + return self._run(self._watch(params)) + if method == "watch.cancel": + if (task := self.watches.pop(params["id"], None)) is not None: + self.loop.call_soon_threadsafe(task.cancel) + return None + raise ValueError(f"unknown execution method: {method}") + + def close(self) -> None: + """End the session's jobs and watches; services outlive it, as they do a + supervisor.""" + self._run(self.jobs.close(None)) + for task in self.watches.values(): + self.loop.call_soon_threadsafe(task.cancel) + self.loop.call_soon_threadsafe(self.loop.stop) + + def _run(self, coroutine) -> Any: + return asyncio.run_coroutine_threadsafe(coroutine, self.loop).result() + + async def _start(self, params: dict) -> dict: + info = self.jobs.start( + owner_id=params["owner_id"], + command=params["command"], + cwd=os.path.realpath(params["cwd"]), + directory=self.directory, + env=_env(params["env"]), + source_request_id=None, + timeout=params["timeout"], + ) + job = self.jobs.jobs[info["id"]] + while job.info.status == "starting" and job.finished is None: + await asyncio.sleep(0.01) + return job.snapshot() + + def _output(self, job: JobRecord, offset: int) -> None: + with self._changed: + self._jobs_changed.setdefault(job.info.id, False) + self._changed.notify_all() + + def _finished(self, job: JobRecord) -> None: + with self._changed: + self._jobs_changed[job.info.id] = True + self._changed.notify_all() + + def _events(self, wait: float) -> list[dict]: + """What changed since the last call: each changed job's state and new output, + in all at most EVENT_BYTES of it, and the path watches' changes.""" + with self._changed: + self._changed.wait_for( + lambda: self._jobs_changed or self._paths, timeout=wait + ) + self._changed.wait_for( + lambda: all(self._jobs_changed.values()), timeout=EVENT_LINGER + ) + changed, self._jobs_changed = self._jobs_changed, {} + paths, self._paths = self._paths, [] + events, budget = [], EVENT_BYTES + for job_id, finished in changed.items(): + job = self.jobs.jobs[job_id] + sent = self._sent.get(job_id, 0) + with open(job.info.output_path, "rb") as output: + output.seek(sent) + data = output.read(budget) + budget -= len(data) + self._sent[job_id] = sent + len(data) + more = self._sent[job_id] < job.info.output_bytes + if more: # the rest goes with the next call + with self._changed: + self._jobs_changed[job_id] = ( + self._jobs_changed.get(job_id, False) or finished + ) + events.append( + { + "type": "job", + "job": job.snapshot(), + "data": base64.b64encode(data).decode(), + "finished": finished and not more, + } + ) + return events + paths + + async def _watch(self, params: dict) -> str: + target, changes = await watch_path(Path(params["path"]), params["recursive"]) + self.watches[params["id"]] = asyncio.create_task( + self._forward(params["id"], changes) + ) + return str(target) + + async def _forward(self, watch_id: str, changes: PathChanges) -> None: + first = True + try: + async for paths, exists in changes: + # The first report says the watch is ready; later ones only changes. + if first or paths or not exists: + self._push_path( + { + "type": "path", + "id": watch_id, + "paths": sorted(paths), + "exists": exists, + } + ) + first = False + if not exists: + return + except asyncio.CancelledError: + raise + except Exception as exc: + self._push_path({"type": "path", "id": watch_id, "error": str(exc)}) + finally: + await changes.aclose() + + def _push_path(self, event: dict) -> None: + with self._changed: + self._paths.append(event) + self._changed.notify_all() + + +class RemoteShellJobs(ShellJobs): + """Shell jobs that run in the execution workspace. The worker owns each process and + its output file; the supervisor keeps the job's record and a copy of its output.""" + + def __init__( + self, + host: RemoteHost, + publish: Callable[[JobRecord], None], + output: Callable[[JobRecord, int], None] | None = None, + ): + super().__init__(publish, output) + self.host = host + self._done: dict[str, asyncio.Future] = {} + self._copies: dict[str, Path] = {} + self._early: dict[str, list[dict]] = {} + + async def start_remote( + self, + *, + owner_id: str, + command: str, + cwd: str, + directory: Path, + env: dict, + source_request_id: str | None, + timeout: float | None = None, + notify: bool = True, + ) -> dict: + snapshot = await self.host.call( + "jobs.start", + { + "owner_id": owner_id, + "command": command, + "cwd": cwd, + "env": env, + "timeout": timeout, + }, + ) + job = JobRecord( + JobInfo(**{name: snapshot[name] for name in _JOB_FIELDS}), + source_request_id, + notify=notify, + ) + job.update(output_bytes=0) # counts the copy, which fills from events + copy = directory / "jobs" / job.info.id / "output.bin" + copy.parent.mkdir(parents=True) + copy.touch() + self._copies[job.info.id] = copy + self._done[job.info.id] = asyncio.get_running_loop().create_future() + self.jobs[job.info.id] = job + job.task = asyncio.create_task(self._follow(job)) + job.task.add_done_callback(self._observe_failure) + for event in self._early.pop(job.info.id, []): + self.apply(event) + return job.snapshot() + + def output_file(self, job: JobRecord) -> str: + return str(self._copies[job.info.id]) + + async def cancel(self, job: JobRecord) -> dict: + job.cancel_requested = True + if job.finished is None: + await self.host.call("jobs.cancel", {"id": job.info.id}) + await asyncio.shield(job.task) + return job.snapshot() + + async def close(self, owner_id: str | None = None) -> None: + jobs = [ + job + for job in self.jobs.values() + if job.finished is None + and (owner_id is None or job.info.owner_id == owner_id) + ] + for job in jobs: + job.cancel_requested = True + if jobs: + try: + await self.host.call("jobs.close", {"owner_id": owner_id}) + except Exception as exc: + self.disconnect(str(exc)) + await asyncio.gather( + *(asyncio.shield(job.task) for job in jobs), return_exceptions=True + ) + + def apply(self, event: dict) -> None: + job = self.jobs.get(event["job"]["id"]) + if job is None: + # A fast job's events can arrive before the reply that registers it. + self._early.setdefault(event["job"]["id"], []).append(event) + return + if job.finished is not None: + return + data = base64.b64decode(event["data"]) + start = job.info.output_bytes + if data: + with open(self._copies[job.info.id], "ab") as copy: + copy.write(data) + job.info = JobInfo( + **{ + **{name: event["job"][name] for name in _JOB_FIELDS}, + "output_bytes": start + len(data), + } + ) + if data and self.output is not None: + self.output(job, start) + if event["finished"]: + self._done[job.info.id].set_result(None) + + def disconnect(self, error: str) -> None: + """Fail every running job: the worker that ran them is gone.""" + for job_id, done in self._done.items(): + if not done.done(): + self.jobs[job_id].update( + status="failed", error=f"execution worker disconnected: {error}" + ) + done.set_result(None) + + async def _follow(self, job: JobRecord) -> None: + try: + await self._done[job.info.id] + finally: + job.finished = time.monotonic() + self.publish(job) + + +class RemoteHost: + """The supervisor's handle on the WorkspaceHost behind the root agent's worker.""" + + def __init__(self, call: Callable[[str, dict], Any]) -> None: + self._call = call + self.jobs: RemoteShellJobs | None = None + self._paths: dict[str, asyncio.Queue] = {} + self._events: asyncio.Task | None = None + + async def call(self, method: str, params: dict) -> Any: + return await asyncio.to_thread(self._call, method, params) + + def follow(self, jobs: RemoteShellJobs) -> None: + self.jobs = jobs + self._events = asyncio.create_task(self._pump()) + + async def close(self) -> None: + if self._events is not None: + self._events.cancel() + await asyncio.gather(self._events, return_exceptions=True) + + async def watch(self, target: Path, recursive: bool) -> tuple[Path, PathChanges]: + watch_id = uuid.uuid4().hex + queue: asyncio.Queue = asyncio.Queue() + self._paths[watch_id] = queue + try: + resolved = await self.call( + "watch.path", + {"id": watch_id, "path": str(target), "recursive": recursive}, + ) + except BaseException: + del self._paths[watch_id] + raise + return Path(resolved), self._changes(watch_id, queue) + + async def _changes(self, watch_id: str, queue: asyncio.Queue) -> PathChanges: + try: + while True: + event = await queue.get() + if "error" in event: + raise RuntimeError(event["error"]) + yield set(event["paths"]), event["exists"] + finally: + self._paths.pop(watch_id, None) + with contextlib.suppress(Exception): + await self.call("watch.cancel", {"id": watch_id}) + + async def _pump(self) -> None: + assert self.jobs is not None + try: + while True: + for event in await self.call("host.events", {"wait": EVENT_WAIT}): + if event["type"] == "job": + self.jobs.apply(event) + elif (queue := self._paths.get(event["id"])) is not None: + queue.put_nowait(event) + except asyncio.CancelledError: + raise + except Exception as exc: # the worker is gone + self.jobs.disconnect(str(exc)) + for queue in self._paths.values(): + queue.put_nowait({"error": f"execution worker disconnected: {exc}"}) diff --git a/tests/test_acp.py b/tests/test_acp.py index 9e62d47..eb97a92 100644 --- a/tests/test_acp.py +++ b/tests/test_acp.py @@ -855,6 +855,10 @@ async def test_acp_session_reuses_engine(monkeypatch, tmp_path): assert initialized.field_meta == { CONTRACT_METADATA_KEY: True, "steering": {"supported": True}, + "ai.prime.rlm/execution-v1": { + "transports": ["stdio", "tcp"], + "tools": ["ipython", "bash", "edit", "fetch"], + }, } created = await agent.new_session( @@ -1113,6 +1117,10 @@ async def test_acp_stdio_lifecycle(tmp_path): assert initialized.field_meta == { CONTRACT_METADATA_KEY: True, "steering": {"supported": True}, + "ai.prime.rlm/execution-v1": { + "transports": ["stdio", "tcp"], + "tools": ["ipython", "bash", "edit", "fetch"], + }, } assert created.field_meta is None assert closed.field_meta[SESSION_METADATA_KEY]["session_id"] == "wire-session" diff --git a/tests/test_compaction.py b/tests/test_compaction.py index 1b4c8f1..67fb7f6 100644 --- a/tests/test_compaction.py +++ b/tests/test_compaction.py @@ -1,7 +1,10 @@ from __future__ import annotations import json +import re +import sys from copy import deepcopy +from pathlib import Path from types import SimpleNamespace from typing import Any @@ -256,7 +259,18 @@ async def test_compaction_retries_reasoning_without_final_content(session, conte engine.close() -async def test_tool_result_overflow_compacts_and_retries(session): +@pytest.mark.parametrize("split", [False, True]) +async def test_tool_result_overflow_compacts_and_retries(session, tmp_path, split): + """Under split execution the kernel reads the worker's synced copy of the history, + and the compaction notice names that copy.""" + config = _config() + if split: + config = config.model_copy( + update={ + "execution_command": (sys.executable, "-m", "rlm.execution"), + "builtin_tools": ("ipython",), + } + ) client = _ScriptedClient( [ _response( @@ -294,13 +308,14 @@ async def test_tool_result_overflow_compacts_and_retries(session): engine = RLMEngine( client=client, # type: ignore[arg-type] session=session, - runtime_config=_config(), + runtime_config=config, + cwd=str(tmp_path), ) try: result = await engine.run("produce a large tool result") finally: - engine.close() + await engine.aclose() assert result.answer == "done" assert engine._metrics.num_compactions == 1 @@ -310,7 +325,9 @@ async def test_tool_result_overflow_compacts_and_retries(session): assert retry_messages[1]["content"].startswith( '\n' + SUMMARY_FRAMING ) - assert str(session.dir / "messages.jsonl") in retry_messages[1]["content"] + named = re.search(r"available in (\S+)\. ", retry_messages[1]["content"]) + assert named is not None + assert (Path(named.group(1)) != session.dir / "messages.jsonl") is split records = [ json.loads(line) for line in (session.dir / "messages.jsonl").read_text().splitlines() diff --git a/tests/test_kernel_recovery.py b/tests/test_kernel_recovery.py index f9803fe..601bad8 100644 --- a/tests/test_kernel_recovery.py +++ b/tests/test_kernel_recovery.py @@ -1,7 +1,10 @@ from __future__ import annotations import asyncio +import contextlib import json +import os +import uuid import pytest @@ -263,3 +266,365 @@ async def started(): repl._interrupt_requested.set() await pending await asyncio.to_thread(repl.shutdown) + + +async def test_remote_execution_preserves_state_and_brokers_children(session, tmp_path): + import sys + + # One CPU, like a sandbox VM: a fast job then finishes before the worker + # acknowledges its start, which is when its events used to get lost. + cpu = str(min(os.sched_getaffinity(0))) + config = _config(max_depth=1).model_copy( + update={ + "execution_command": ( + "taskset", + "-c", + cpu, + sys.executable, + "-m", + "rlm.execution", + ), + "skills": ("bash", "edit"), + "builtin_tools": ("ipython",), + "policy": _config(max_depth=1).policy.model_copy( + update={"compaction": False} + ), + } + ) + + def factory(**kwargs): + return RLMEngine( + client=DummyClient( + [ + _tool( + "from pathlib import Path; Path('child.txt').write_text('child'); print('CHILD')" + ), + DummyMessage(content="child complete"), + ] + ), + **kwargs, + ) + + supervisor = SessionTreeSupervisor( + root_session=session, + runtime_config=config, + cwd=str(tmp_path), + engine_factory=factory, + ) + client = DummyClient( + [ + _tool("value = 41; print(await bash('pwd'))"), + _tool( + """ +from pathlib import Path +assert value == 41 +child = await rlm.agent.spawn('write child.txt') +await child.wait(timeout=30) +assert (await child.result()).answer == 'child complete' +assert 'write child.txt' in (await child.history()).user_messages()[0]['content'] +assert Path('child.txt').read_text() == 'child' +h = await rlm.history() +assert h.user_messages()[0]['content'] == 'run remotely' +import asyncio +# The supervisor's jobs, services and watches run in the workspace, beside this kernel. +job = await rlm.shell.run('echo $PPID; touch from-shell') +assert job.ok and int(job.text.split()[0]) != SUPERVISOR_PID, job.text +assert Path('from-shell').exists() +import time +started = time.monotonic() +assert (await rlm.shell.run('true')).ok and time.monotonic() - started < 5 +svc = await rlm.services.start(SERVICE, 'echo service-up; sleep 30') +for _ in range(50): + if 'service-up' in await svc.tail(): + break + await asyncio.sleep(0.1) +assert 'service-up' in await svc.tail() +assert not (await svc.stop(grace=1)).running +await rlm.watch.path('.') +Path('watched.txt').write_text('x') +for _ in range(50): + if any(e['type'] == 'watch.path' for e in await rlm.inbox.list()): + break + await asyncio.sleep(0.1) +assert any(e['type'] == 'watch.path' for e in await rlm.inbox.list()) +print('REMOTE_OK') +""".replace("SUPERVISOR_PID", str(os.getpid())).replace( + "SERVICE", repr(f"split-{uuid.uuid4().hex[:8]}") + ) + ), + DummyMessage(content="done"), + ] + ) + engine = RLMEngine( + client=client, + session=session, + runtime_config=config, + cwd=str(tmp_path), + supervisor=supervisor, + ) + try: + await engine.prompt("run remotely") + logs = (session.dir / "messages.jsonl").read_text() + assert any( + record.get("type") == "tool_result" + and "REMOTE_OK" in record.get("content", "") + for record in map(json.loads, logs.splitlines()) + ) + finally: + await engine.aclose() + await supervisor.aclose() + + +async def test_remote_execution_interrupt_and_recovery(session, tmp_path): + import sys + from rlm.execution import RemoteREPL + + repl = RemoteREPL( + command=(sys.executable, "-m", "rlm.execution"), + cwd=str(tmp_path), + session=session, + ) + try: + await repl.run_in_thread(repl.start) + task = asyncio.create_task( + repl.run_in_thread(repl.execute, "import time; time.sleep(60)", 120) + ) + await asyncio.sleep(1) + await asyncio.to_thread(repl.interrupt) + await asyncio.wait_for(task, 10) + repl.finish_interrupt() + result = await repl.run_in_thread(repl.execute, "print('STILL_ALIVE')", 10) + assert "STILL_ALIVE" in result + await repl.run_in_thread(repl.execute, "import os; os._exit(7)", 10) + assert repl.take_recovery_notices() + assert "RECOVERED" in await repl.run_in_thread( + repl.execute, "print('RECOVERED')", 10 + ) + finally: + await asyncio.to_thread(repl.shutdown) + + +async def test_endless_output_is_bounded_until_the_timeout(session): + from rlm.tools.bash import run_bash + from rlm.tools.output import OUTPUT_KEEP_BYTES + + repl = IPythonREPL(session=session, cwd=str(session.dir)) + try: + await repl.run_in_thread(repl.start) + output = await asyncio.wait_for( + repl.run_in_thread( + repl.execute, "while True: print('x' * 65536, flush=True)", 5 + ), + timeout=25, + ) + assert len(output.encode()) < OUTPUT_KEEP_BYTES + 256 + assert output.endswith("[execution timed out after 5s and was interrupted]") + assert repl.omitted.nbytes > 0 + continued = await repl.run_in_thread(repl.execute, "print('STILL_ALIVE')", 10) + assert "STILL_ALIVE" in continued + finally: + await repl.run_in_thread(repl.shutdown) + output = await asyncio.to_thread(run_bash, "yes", 3) + assert len(output.encode()) < OUTPUT_KEEP_BYTES + 256 + assert output.endswith("Error: command timed out after 3s") + + +@pytest.mark.parametrize("budget", [None, 3 << 20], ids=["default", "large"]) +@pytest.mark.parametrize("split", [False, True], ids=["local", "split"]) +@pytest.mark.parametrize("tool", ["bash", "ipython"]) +async def test_long_output_runs_to_its_end_and_reads_as_whole( + session, tmp_path, split, tool, budget +): + """A command that prints far more than a tool result keeps runs to its end, and the + model sees the truncation of its whole output: the same start, end and counts. A + budget larger than what was kept shows no more of either end than that, once.""" + import sys + + from rlm.compaction import truncate_tool_output + from rlm.tools.output import OUTPUT_KEEP_BYTES + + lines = "".join(f"{i}\n" for i in range(1, 300001)) + if tool == "bash": + call = DummyToolCall( + "bash", + { + "command": "printf 'a\\r\\nb\\rc\\n'; seq 300000; echo DONE; " + "seq 200000 >&2; echo ERR_END >&2" + }, + ) + errors = "".join(f"{i}\n" for i in range(1, 200001)) + whole = f"a\nb\nc\n{lines}DONE\n\n{errors}ERR_END\n".strip() + else: + call = DummyToolCall( + "ipython", {"code": "for i in range(1, 300001): print(i)\nprint('DONE')"} + ) + whole = f"{lines}DONE\n" + config = _config() + config = config.model_copy( + update={ + "builtin_tools": ("ipython", "bash"), + "policy": config.policy.model_copy( + update={"max_tool_output_bytes": budget} + ), + } + ) + if split: + config = config.model_copy( + update={"execution_command": (sys.executable, "-m", "rlm.execution")} + ) + client = DummyClient( + [DummyMessage(tool_calls=[call]), DummyMessage(content="done")] + ) + engine = RLMEngine( + client=client, session=session, runtime_config=config, cwd=str(tmp_path) + ) + try: + await engine.prompt("print a lot") + finally: + await engine.aclose() + seen = [m for m in client.calls[1]["messages"] if m["role"] == "tool"][-1] + if budget is None: + assert seen["content"] == truncate_tool_output(whole) + else: + assert seen["content"].count(whole[:200]) == 1 + assert seen["content"].count(whole[-200:]) == 1 + assert len(seen["content"].encode()) <= budget + logged = [ + record["content"] + for record in map(json.loads, (session.dir / "messages.jsonl").open()) + if record.get("type") == "tool_result" + ][-1] + assert "bytes of output omitted" in logged + assert len(logged.encode()) < OUTPUT_KEEP_BYTES + 512 + assert logged[:10_000] == whole[:10_000] + assert logged[-10_000:] == whole[-10_000:] + + +@pytest.mark.parametrize("split", [False, True], ids=["local", "split"]) +async def test_processes_live_as_on_a_terminal(session, tmp_path, split): + """A cancelled turn stops its running bash command at once. The session's end stops + what is attached to the kernel, like a closing terminal, and leaves a daemon and a + bash command's background jobs running, among them one that held the command's + output until its timeout.""" + import signal + import sys + + def loop(name): + return f"echo $$ > {name}.pid; while :; do touch {name}.alive; sleep 0.1; done" + + cell = ( + "import subprocess\n" + f"subprocess.Popen(['sh', '-c', {loop('attached')!r}])\n" + f"subprocess.Popen(['sh', '-c', {loop('daemon')!r}], start_new_session=True)\n" + ) + background = f"(sh -c {loop('job')!r}) >/dev/null 2>&1 &" + held = f"sh -c {loop('held')!r} &" + client = DummyClient( + [ + _tool(cell), + DummyMessage(tool_calls=[DummyToolCall("bash", {"command": background})]), + DummyMessage(tool_calls=[DummyToolCall("bash", {"command": held})]), + DummyMessage(tool_calls=[DummyToolCall("bash", {"command": "sleep 60"})]), + ] + ) + config = _config() + config = config.model_copy( + update={ + "builtin_tools": ("ipython", "bash"), + "policy": config.policy.model_copy(update={"exec_timeout": 2}), + } + ) + if split: + config = config.model_copy( + update={"execution_command": (sys.executable, "-m", "rlm.execution")} + ) + engine = RLMEngine( + client=client, session=session, runtime_config=config, cwd=str(tmp_path) + ) + names = ("attached", "daemon", "job", "held") + try: + try: + turn = asyncio.create_task(engine.prompt("run the commands")) + while len(client.calls) < 4: + await asyncio.sleep(0.1) + await asyncio.sleep(0.5) + turn.cancel() + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(turn, 5) + assert all((tmp_path / f"{name}.alive").exists() for name in names) + finally: + await engine.aclose() + for name in names: + (tmp_path / f"{name}.alive").unlink(missing_ok=True) + await asyncio.sleep(0.5) + assert not (tmp_path / "attached.alive").exists() + assert (tmp_path / "daemon.alive").exists() + assert (tmp_path / "job.alive").exists() + assert (tmp_path / "held.alive").exists() + finally: + for name in names: + pid = tmp_path / f"{name}.pid" + if pid.exists(): + with contextlib.suppress(ProcessLookupError): + os.kill(int(pid.read_text()), signal.SIGKILL) + + +@pytest.mark.parametrize("split", [False, True], ids=["local", "split"]) +async def test_history_past_the_frame_limit_reads_whole(session, tmp_path, split): + """A history larger than one frame reads whole: the agent's own, which reaches a + worker in frame-sized chunks, and a child's, which the supervisor returns in pages.""" + import sys + + from rlm.execution import MAX_FRAME + + prompt = "x" * (MAX_FRAME + 1024) + config = _config(max_depth=1).model_copy(update={"builtin_tools": ("ipython",)}) + if split: + config = config.model_copy( + update={"execution_command": (sys.executable, "-m", "rlm.execution")} + ) + + def factory(**kwargs): + return RLMEngine(client=DummyClient([DummyMessage(content="ok")]), **kwargs) + + supervisor = SessionTreeSupervisor( + root_session=session, + runtime_config=config, + cwd=str(tmp_path), + engine_factory=factory, + ) + client = DummyClient( + [ + _tool( + f""" +import json +from rlm import history +assert len((await history()).user_messages()[0]['content']) == {len(prompt)} +child = await rlm.agent.spawn('answer') +await child.wait(timeout=30) +# A long child conversation: five records of a quarter frame each. +with (child.session_dir / 'messages.jsonl').open('a') as stream: + for _ in range(5): + stream.write(json.dumps({{'type': 'padding', 'text': 'y' * {MAX_FRAME // 4}}}) + '\\n') +h = await child.history() +assert 'answer' in h.user_messages()[0]['content'] +print(f"PADDED {{sum(event['type'] == 'padding' for event in h.events)}}") +""" + ), + DummyMessage(content="done"), + ] + ) + engine = RLMEngine( + client=client, + session=session, + runtime_config=config, + cwd=str(tmp_path), + supervisor=supervisor, + ) + try: + await engine.prompt(prompt) + finally: + await engine.aclose() + tool = [m for m in client.calls[1]["messages"] if m["role"] == "tool"][-1] + # Printed, not asserted in the cell: a traceback would show the cell's source. + assert "PADDED 5" in tool["content"], tool["content"][-2000:]