From b56eaeb67c796b745d23cdc78bc8372be1085520 Mon Sep 17 00:00:00 2001 From: Dizhan Xue Date: Tue, 27 Oct 2026 08:18:08 +0000 Subject: [PATCH 01/19] perf(*): let raven agents share memory with their host and with each other Each product agent under agents/ ran as a whole Raven process of its own, about 240 MiB PSS, nearly all of it private: a host running the five shipped products paid about 1.2 GiB for them. Three mechanisms, each switchable live under subagents.runtime in the host's config: - idleCloseS (default 600): the ACP pool closes an agent's process once nothing is in flight either way and nothing has been said for the budget. Only an agent that advertises loadSession, or one that never opened a session, is closed, so a stateful instance's next turn reopens its session in a fresh process. A held wake is a prompt still waiting for its answer, so it is never idle. - llmRelay (default on): the host makes a Raven agent's LiteLLM calls, so the agent never imports LiteLLM. LiteLLMProvider's LiteLLM-free construction and answers move to LiteLLMShape, which RelayProvider shares; only the single chat / chat_stream attempt and the catalogue lookups that need LiteLLM's tables (the relayed decorator in rates, capabilities, common_models and the Anthropic adapter) cross a local socket with a token. Retry, classification, usage records and traces stay in the agent; the host's calls run under recorded_by_caller. An agent that cannot reach its host loads LiteLLM as before. - forkTemplate (default on, Linux only): a template process that has already imported what agents import forks each Raven agent the host launches, so the loaded pages are shared. It preloads the modules at least two products import, learned per product from the agents themselves (one product's own while only one has been seen), plus the token estimator's table. A launcher's os.execv of python -m raven runs Raven in the same process. Anything else, and any failure, launches the ordinary way. The relay and the template are per event loop, since a CLI turn is one asyncio.run. raven acp takes the relay offer out of its environment at startup, and the token estimator drops the duplicate copy of its ranks that LiteLLM's import used to drop. Measured with five product agents, one turn each (PSS): 1414 -> 724 MiB in total, the agents 1259 -> 547 MiB. One agent alone: 402 -> 371 MiB. --- CONTEXT.md | 28 + docs-site/docs/agent-integrations.md | 22 + docs-site/docs/agent-integrations.zh.md | 18 + raven/acp_client/client.py | 56 +- raven/acp_client/fork_template.py | 843 ++++++++++++++++++ raven/acp_client/pool.py | 102 +++ raven/cli/acp_commands.py | 4 + raven/config/live.py | 27 +- raven/config/schema.py | 32 + .../providers/anthropic_messages_provider.py | 2 + raven/providers/capabilities.py | 25 +- raven/providers/common_models.py | 3 + raven/providers/factory.py | 8 +- raven/providers/litellm_provider.py | 333 +------ raven/providers/litellm_shape.py | 354 ++++++++ raven/providers/per_model_provider.py | 9 +- raven/providers/rates.py | 22 +- raven/providers/relay.py | 274 ++++++ raven/providers/relay_provider.py | 186 ++++ raven/providers/relay_server.py | 288 ++++++ raven/providers/relay_wire.py | 302 +++++++ raven/utils/tokens.py | 21 +- tests/test_acp_fork_template.py | 703 +++++++++++++++ tests/test_acp_pool_idle_close.py | 226 +++++ tests/test_provider_auth_method.py | 3 + tests/test_provider_catalog.py | 7 +- tests/test_provider_prompt_cache.py | 2 +- tests/test_provider_relay.py | 776 ++++++++++++++++ tests/test_provider_resolution_invariants.py | 16 +- 29 files changed, 4341 insertions(+), 351 deletions(-) create mode 100644 raven/acp_client/fork_template.py create mode 100644 raven/providers/litellm_shape.py create mode 100644 raven/providers/relay.py create mode 100644 raven/providers/relay_provider.py create mode 100644 raven/providers/relay_server.py create mode 100644 raven/providers/relay_wire.py create mode 100644 tests/test_acp_fork_template.py create mode 100644 tests/test_acp_pool_idle_close.py create mode 100644 tests/test_provider_relay.py diff --git a/CONTEXT.md b/CONTEXT.md index 24bba0e85..ff453caa6 100644 --- a/CONTEXT.md +++ b/CONTEXT.md @@ -3392,6 +3392,34 @@ and every tool result. _Avoid_: calling it a transcript - a reader drawing a delegated run wants the Instance Log or `transcript.jsonl`, not this. +**Idle close** (`AcpConnectionPool.close_idle`, `subagents.runtime.idleCloseS`): +The pool ending a pooled ACP agent's process once the connection has been quiet for the +configured budget: no request of the host's unanswered (a held wake is one), none of the +agent's being answered, no frame from the agent and no lease for the whole budget. Only an +agent that advertises `loadSession`, or one that never opened a session, is closed, so a +stateful instance's next turn reopens its session in a fresh process with `session/load`. +_Avoid_: calling it a timeout - nothing in flight is cut short. + +**LLM relay** (`raven/providers/relay.py`, `relay_provider.py`, `relay_server.py`, `subagents.runtime.llmRelay`): +The host making a Raven agent's LiteLLM calls for it. The host offers a local socket and a +token in the launch environment (`RAVEN_LLM_RELAY`), and an agent that reads them builds a +`RelayProvider` where it would have built a `LiteLLMProvider`: the same construction and the +same answers (both stand on `LiteLLMShape`), with only the single `chat` / `chat_stream` +attempt and the `relayed` catalogue lookups crossing to the host. Retry, classification, usage +records and traces stay in the agent; the host's calls run under `recorded_by_caller` so a call +is billed once. An agent that cannot reach the host loads LiteLLM and works as before. + +**Fork template** (`raven/acp_client/fork_template.py`, `subagents.runtime.forkTemplate`): +On Linux, one process per host that has imported what Raven agents import, frozen out of the +collector's reach, and forks each Raven agent the host launches, so the agents share those +pages. It starts through the parent-death wrapper like any agent; each forked agent leads a +session of its own and arms the parent-death signal against the template. A product launcher's +`os.execv` of `python -m raven` runs Raven in the same process. Agents write the modules they +imported beside the template, per product, and the next template loads the modules at least two +products import (one product's own while only one has been seen): a module one product alone imports +is shared by no other process, so it is left for that product to import after the fork. Only the host's own interpreter starting +`-m raven` or a product launcher is forked; anything else launches the ordinary way. + **Leftover agent** (`raven/acp_client/leftovers.py`): An ACP agent process group whose launching host is no longer running. An ACP agent runs in a session of its own, so it gets none of the signals sent to its host's process group or diff --git a/docs-site/docs/agent-integrations.md b/docs-site/docs/agent-integrations.md index e390557d4..35cc3ce1e 100644 --- a/docs-site/docs/agent-integrations.md +++ b/docs-site/docs/agent-integrations.md @@ -73,6 +73,28 @@ and Raven-Code's CLI hosting start a native product launcher explicitly. A successful readiness probe or delegation receipt is not evidence that the task completed; inspect the result, files, and tests or report artifacts. +### Sharing memory with the host { #sharing-memory-with-the-host } + +Each native agent runs as a Raven process of its own. Three settings under +`subagents.runtime` in the host's `config.json` keep that from costing a whole +Raven per agent. All three are on by default and are read live, so a change +applies to the next launch. + +| Key | Default | What it does | +| --- | --- | --- | +| `idleCloseS` | `600` | Closes an agent's process once it has been idle this many seconds: no request in flight either way and nothing said. Only an agent that can reload its sessions is closed, so the next turn resumes the conversation in a fresh process. `0` or `null` keeps agents until the host exits. | +| `llmRelay` | `true` | The host makes a native agent's LiteLLM calls, so the agent never loads LiteLLM. The agent still builds each request, retries and records usage itself. | +| `forkTemplate` | `true` | On Linux, native agents are forked from one process that has already loaded Raven, and share those pages instead of each loading its own copy. Other systems launch agents as before. | + +The template loads ahead only the modules at least two native agents share -- +one agent's own modules are imported by that agent after the fork -- so it +grows with what the agents have in common, not with how many kinds there are. +It is a process too, so it saves memory once two or more native agents run at +the same time; with a single agent it costs somewhat more than it saves. Set +`forkTemplate` to `false` on a device that only ever runs one. An agent the +relay or the template cannot serve -- the host cannot be reached, or the agent +runs on another interpreter -- starts the ordinary way. + ### Minimal ACP invocation { #minimal-acp-invocation } Each shipped entry is started by its launcher. Use the corresponding command: diff --git a/docs-site/docs/agent-integrations.zh.md b/docs-site/docs/agent-integrations.zh.md index a61bd03ba..ca2ce5382 100644 --- a/docs-site/docs/agent-integrations.zh.md +++ b/docs-site/docs/agent-integrations.zh.md @@ -62,6 +62,24 @@ roster。不要假定修改显示名就创建了新身份。可通过 `raven onb launcher。readiness probe 成功或收到委派回执，都不代表任务已完成；仍需检查结果、文件、 测试或报告产物。 +### 与宿主共享内存 { #sharing-memory-with-the-host } + +每个原生 Agent 都是一个独立的 Raven 进程。宿主 `config.json` 中 `subagents.runtime` +下的三个设置让它不必每个 Agent 都占用一整份 Raven 的内存。三者默认开启，并实时读取， +修改后从下一次启动起生效。 + +| 键 | 默认值 | 作用 | +| --- | --- | --- | +| `idleCloseS` | `600` | Agent 空闲达到该秒数（双方都没有在途请求，也没有任何消息）后关闭其进程。只关闭能重新加载会话的 Agent，因此下一轮会在新进程中续上原对话。`0` 或 `null` 表示一直保留到宿主退出。 | +| `llmRelay` | `true` | 由宿主代为发出原生 Agent 的 LiteLLM 调用，Agent 自身不再加载 LiteLLM。请求组装、重试和用量记录仍由 Agent 自己完成。 | +| `forkTemplate` | `true` | 在 Linux 上，原生 Agent 从一个已加载 Raven 的进程 fork 出来，共享这些内存页，而不是各自加载一份。其他系统仍按原方式启动。 | + +模板只预加载至少两个原生 Agent 共用的模块，某个 Agent 独有的模块由它在 fork 之后自行 +导入，因此模板的大小取决于各 Agent 的共同部分，而不是 Agent 的种类数。模板本身也是一个 +进程，因此同时运行两个及以上原生 Agent 时才会省内存；只有一个 Agent 时，它的占用会略多于 +它省下的部分。只会运行单个 Agent 的设备可将 `forkTemplate` 设为 `false`。中转或模板无法服务的 Agent（宿主不可达，或 Agent 使用另一个解释器） +会按原方式启动。 + ### 最小 ACP 调用 { #minimal-acp-invocation } 随 Raven 发布的每个 Agent 都由自己的 launcher 启动。使用对应命令： diff --git a/raven/acp_client/client.py b/raven/acp_client/client.py index 0b5db96f2..4c108fdcd 100644 --- a/raven/acp_client/client.py +++ b/raven/acp_client/client.py @@ -142,6 +142,44 @@ async def _end(pgid: int, proc: Any) -> None: pass +async def _share_with_raven_agent( + name: str, argv: list[str], env: dict[str, str], cwd: str | None, base_env: dict[str, str] +) -> Any: + """What this host shares with a Raven agent it launches: its model calls, and its loaded code. + + Returns the environment to launch with -- the relay's address added when + the host's ``subagents.runtime.llmRelay`` is on -- or, when the fork + template started the agent, the endpoint it is already running behind. + Anything that is not a Raven agent on this interpreter is returned its + environment unchanged, and every failure here falls back to that. + """ + from raven.acp_client import fork_template, leftovers + from raven.agent.subagent.backends.env import host_identity_env + from raven.agent.subagent.role import subagent_role_env + from raven.config.live import default_live, subagent_runtime + + if not fork_template.raven_agent_argv(argv): + return env + runtime = subagent_runtime(default_live()) + relayed = False + if runtime.llm_relay: + from raven.providers.relay_server import relay_env + + try: + env = {**env, **(await relay_env())} + relayed = True + except Exception as exc: # noqa: BLE001 - the agent loads LiteLLM itself, as before + logger.warning("acp agent {!r}: no llm relay for it: {}", name, exc) + if runtime.fork_template: + template_env = {**base_env, **host_identity_env(), **subagent_role_env(), **leftovers.launch_env()} + endpoint = await fork_template.fork_launch( + argv, env, cwd, base_env=template_env, skip=("litellm",) if relayed else () + ) + if endpoint is not None: + return endpoint + return env + + def _retrieved(task: asyncio.Task) -> None: """Read the outcome of a task nobody waits on, so a failure is not reported as never retrieved.""" if not task.cancelled(): @@ -494,6 +532,16 @@ async def launch( # The agent's own argv, for what this logs and reports: the spawned one # may lead with the wrapper that ties the agent to this host. program_argv = argv + shared = await _share_with_raven_agent(name, argv, child_env, cwd, base_env) + if isinstance(shared, dict): + child_env = shared + else: + client = cls( + name=name, endpoint=shared, on_request=on_request, on_notification=on_notification, journal=journal + ) + client._start_loops() + logger.info("acp agent {!r}: forked {} from the template (pid {})", name, program_argv[:2], shared.pid) + return client argv = leftovers.bind_to_host(argv, child_env, cwd) job = None try: @@ -594,7 +642,7 @@ async def close(self) -> None: if self._closed: return self._closed = True - if isinstance(self._endpoint, _ProcessEndpoint): + if hasattr(self._endpoint, "grace"): self._endpoint.grace = self._granted_grace() try: await self._endpoint.end() @@ -853,6 +901,12 @@ def prompting(self) -> bool: """ return any(method == "session/prompt" for method in self._request_methods.values()) + @property + def busy(self) -> bool: + """Whether either side is waiting on the other: a request of raven's own + unanswered (a turn, a held wake included) or one of the agent's being answered.""" + return bool(self._pending) or bool(self._answer_tasks) + @property def idle_seconds(self) -> float: """Seconds since the agent last put any frame on the wire. diff --git a/raven/acp_client/fork_template.py b/raven/acp_client/fork_template.py new file mode 100644 index 000000000..72bab171a --- /dev/null +++ b/raven/acp_client/fork_template.py @@ -0,0 +1,843 @@ +"""Raven agents started by forking a process that has already loaded their code. + +A Raven agent is a whole Raven: launched fresh, each one imports the same few +thousand modules and holds its own copy of them, about 120 MiB before it has +done anything. On Linux a process that has imported them once can ``fork``, and +the children share those pages with it until they write to them. So the host +keeps one such process -- the template -- and asks it for each Raven agent it +launches; measured on the product agents, what each one costs on its own falls +by about two thirds. + +**The template.** ``python -m raven.acp_client.fork_template``, started by the +host through the same parent-death wrapper as any agent (``leftovers``), so it +ends with the host. It imports a seed of the modules every agent needs plus +the modules agents were seen to import before -- each agent writes its module +list beside the template after a while, under its product, and the next +template loads what at least two products share (one product's own list while +only one has been seen; see :func:`preload_names`) -- then collects, freezes the heap out of the collector's reach -- a collection writes +to every object it visits, which would copy the shared pages back apart -- and +then answers fork requests on a socket in a private directory. It runs no +thread, and refuses to fork if an import started one. + +**A fork.** The host hands over the agent's argv, environment and working +directory, and the three pipe ends that become its stdin, stdout and stderr. +The child starts a session of its own (so the host ends it with ``killpg`` as +it would any agent), arms the parent-death signal against the template, and +runs the agent's program in place of the ``exec`` a fresh interpreter would +have done: a product launcher's own ``os.execv`` of ``python -m raven`` is +answered the same way, by running Raven in the same process. The template +reaps the child and tells the host its exit status over the request's +connection, which stays open for that. + +**What is eligible.** Only an agent this interpreter would run as Raven: the +host's own interpreter starting ``-m raven`` or a product launcher (a ``.py`` +beside a ``subagent.json``). Anything else, any failure to start the template +or to fork, and every system but Linux, launches the ordinary way. +""" + +from __future__ import annotations + +import argparse +import asyncio +import collections +import contextlib +import gc +import importlib +import json +import os +import signal +import socket +import sys +import tempfile +import threading +import time +from collections.abc import Sequence +from pathlib import Path +from typing import Any + +#: Modules every Raven agent imports before it serves, imported by every +#: template whatever it has learned. +SEED_MODULES = ( + "raven.config.product_render", + "raven.cli.commands", + "raven.acp.server", + "raven.acp.methods", + "raven.rpc.bootstrap", + "raven.agent.loop", +) + +#: How long after starting an agent writes the modules it has imported. Long +#: enough to cover its first session and turn, where nearly all of them load. +LEARN_AFTER_S = 30.0 +_LEARNED_FILE = "modules.json" +_LEARNED_CAP = 20_000 + +_READY_TIMEOUT_S = 60.0 +_REQUEST_TIMEOUT_S = 30.0 +_GRACE_POLL_S = 0.02 +_TEARDOWN_GRACE_S = 1.0 +_READER_LIMIT = 8 * 1024 * 1024 + + +def raven_agent_argv(argv: Sequence[str]) -> bool: + """Whether ``argv`` starts a Raven agent on this very interpreter. + + The interpreter is compared by path, not by file: two virtual environments + link the same base interpreter, and an agent installed in another one must + run with its own packages, not with the ones this process has loaded. + """ + if len(argv) < 2: + return False + try: + same = os.path.normcase(os.path.abspath(argv[0])) == os.path.normcase(os.path.abspath(sys.executable)) + except (TypeError, ValueError): + return False + if not same: + return False + if argv[1] == "-m": + return len(argv) >= 3 and argv[2] == "raven" + script = Path(argv[1]) + try: + return script.suffix == ".py" and (script.parent / "subagent.json").is_file() + except OSError: + return False + + +# ---- the host's side --------------------------------------------------------- + + +class ForkedEndpoint: + """An agent the template forked: its pipes, its group, and the exit status the template reports. + + The same surface as a launched child's endpoint (``client.Endpoint``). The + agent is not this process's child, so its status cannot be waited for here; + the template reaps it and says how it ended on the connection the request + came in on. + """ + + def __init__( + self, + pid: int, + *, + stdin: asyncio.StreamWriter, + stdout: asyncio.StreamReader, + stderr: asyncio.StreamReader, + control: tuple[asyncio.StreamReader, asyncio.StreamWriter], + template: ForkTemplate, + ) -> None: + self.pid = pid + self.grace: float | None = None + self._stdin = stdin + self._stdout = stdout + self._stderr = stderr + self._control_reader, self._control_writer = control + self._template = template + self._exited: asyncio.Future[int] = asyncio.get_running_loop().create_future() + self._watch = asyncio.get_running_loop().create_task(self._await_exit(), name=f"forked-agent-{pid}") + + async def _await_exit(self) -> None: + from raven.providers import relay_wire + + code = -signal.SIGKILL + try: + frame = await relay_wire.read_frame(self._control_reader) + if isinstance(frame, dict) and isinstance(frame.get("exit"), int): + code = frame["exit"] + except (ConnectionError, relay_wire.FrameError, asyncio.IncompleteReadError): + pass + finally: + # No report means the template went, and the parent-death signal + # it armed took the agent with it. + if not self._exited.done(): + self._exited.set_result(code) + self._template.forget(self) + + @property + def returncode(self) -> int | None: + return self._exited.result() if self._exited.done() else None + + def describe(self) -> str: + return f"pid {self.pid} (forked)" + + async def write(self, data: bytes) -> None: + if self._stdin.is_closing(): + raise BrokenPipeError("stdin is closed") + self._stdin.write(data) + await self._stdin.drain() + + async def readline(self) -> bytes: + return await self._stdout.readline() + + async def readerr(self) -> bytes: + return await self._stderr.readline() + + async def end(self) -> None: + try: + if self.returncode is None and not self._stdin.is_closing(): + self._stdin.close() + loop = asyncio.get_running_loop() + deadline = loop.time() + (_TEARDOWN_GRACE_S if self.grace is None else self.grace) + while self.returncode is None and loop.time() < deadline: + await asyncio.sleep(_GRACE_POLL_S) + finally: + with contextlib.suppress(ProcessLookupError, PermissionError): + os.killpg(self.pid, signal.SIGKILL) + + async def wait(self) -> int | None: + return await asyncio.shield(self._exited) + + def ended(self) -> str | None: + return None + + def release(self) -> None: + from raven.acp_client import leftovers + + for stream in (self._stdin, self._control_writer): + with contextlib.suppress(Exception): + stream.close() + for reader in (self._stdout, self._stderr): + transport = getattr(reader, "_transport", None) + if transport is not None: + with contextlib.suppress(Exception): + transport.close() + self._watch.cancel() + leftovers.forget(self.pid) + + +class ForkTemplate: + """The host's handle on one template process: starts it, asks it for agents, closes it when idle.""" + + def __init__(self, *, skip: Sequence[str] = ()) -> None: + self._skip = tuple(skip) + self._proc: asyncio.subprocess.Process | None = None + self._dir: Path | None = None + self._socket = "" + self._children: set[ForkedEndpoint] = set() + self._forking = 0 + self._idle_task: asyncio.Task | None = None + self._log_task: asyncio.Task | None = None + self._lock = asyncio.Lock() + self.loop = asyncio.get_running_loop() + """The loop its process and pipes were opened under; the handle is no use to another.""" + self.broken = "" + """Why this template will not be used again, once it has failed.""" + + @property + def running(self) -> bool: + return self._proc is not None and self._proc.returncode is None + + @property + def children(self) -> int: + return len(self._children) + + def forget(self, child: ForkedEndpoint) -> None: + self._children.discard(child) + if not self._children: + from raven.config.live import default_live, subagent_runtime + + budget = subagent_runtime(default_live()).idle_close_s + self.close_when_idle(budget if budget is not None and budget > 0 else None) + + async def _start(self, env: dict[str, str]) -> None: + from raven.acp_client import leftovers + from raven.config.paths import get_runtime_subdir + + self._dir = Path(tempfile.mkdtemp(prefix="raven-fork-template-")) + self._socket = str(self._dir / "template.sock") + learned = get_runtime_subdir("acp-template") / _LEARNED_FILE + argv = [sys.executable, "-P", "-m", __name__, "--socket", self._socket, "--learned", str(learned)] + argv += ["--learn-after", str(LEARN_AFTER_S)] + for prefix in self._skip: + argv += ["--skip", prefix] + argv = leftovers.bind_to_host(argv, env, str(self._dir)) + self._proc = await asyncio.create_subprocess_exec( + *argv, + stdin=asyncio.subprocess.PIPE, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + cwd=str(self._dir), + env=env, + start_new_session=True, + ) + leftovers.record(self._proc.pid, "fork-template") + self._log_task = asyncio.get_running_loop().create_task(self._log_stderr(self._proc), name="fork-template-log") + ready = await asyncio.wait_for(self._ready_report(self._proc), timeout=_READY_TIMEOUT_S) + if not ready.get("ready"): + raise RuntimeError(str(ready.get("reason") or "the template exited before it was ready")) + from loguru import logger + + logger.info( + "fork template: ready (pid {}, {} modules, {:.1f}s)", + self._proc.pid, + ready.get("modules"), + ready.get("seconds", 0.0), + ) + + @staticmethod + async def _ready_report(proc: asyncio.subprocess.Process) -> dict[str, Any]: + """The template's one line saying whether it is ready; ``{}`` if it ended without one.""" + if proc.stdout is None: + return {} + while line := await proc.stdout.readline(): + try: + report = json.loads(line.decode("utf-8")) + except ValueError: + continue + if isinstance(report, dict) and "ready" in report: + return report + return {} + + @staticmethod + async def _log_stderr(proc: asyncio.subprocess.Process) -> None: + from loguru import logger + + if proc.stderr is None: + return + while line := await proc.stderr.readline(): + logger.debug("fork template: {}", line.decode("utf-8", "replace").rstrip()) + + async def spawn( + self, argv: list[str], env: dict[str, str], cwd: str | None, base_env: dict[str, str] + ) -> ForkedEndpoint: + """An agent forked from the template, starting the template first if needed. Raises when it cannot.""" + if self.broken: + raise RuntimeError(self.broken) + async with self._lock: + if not self.running: + if self._proc is not None: + # A template that ended on its own: clear what it left before the next. + await self.close() + try: + await self._start(base_env) + except BaseException as exc: + self.broken = f"the fork template did not start: {exc}" + await self.close() + raise RuntimeError(self.broken) from exc + if self._idle_task is not None: + self._idle_task.cancel() + self._idle_task = None + self._forking += 1 + try: + child = await self._fork(argv, env, cwd) + finally: + self._forking -= 1 + self._children.add(child) + return child + + async def _fork(self, argv: list[str], env: dict[str, str], cwd: str | None) -> ForkedEndpoint: + from raven.acp_client import leftovers + from raven.providers import relay_wire + + loop = asyncio.get_running_loop() + stdin_r, stdin_w = os.pipe() + stdout_r, stdout_w = os.pipe() + stderr_r, stderr_w = os.pipe() + sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) + try: + sock.settimeout(_REQUEST_TIMEOUT_S) + + def ask() -> dict[str, Any]: + sock.connect(self._socket) + socket.send_fds(sock, [b"F"], [stdin_r, stdout_w, stderr_w]) + relay_wire.send_frame(sock, {"argv": argv, "env": env, "cwd": cwd or os.getcwd()}) + return relay_wire.recv_frame(sock) + + reply = await asyncio.to_thread(ask) + except BaseException: + for fd in (stdin_w, stdout_r, stderr_r): + os.close(fd) + sock.close() + raise + finally: + for fd in (stdin_r, stdout_w, stderr_w): + os.close(fd) + pid = reply.get("pid") if isinstance(reply, dict) else None + if not isinstance(pid, int): + for fd in (stdin_w, stdout_r, stderr_r): + os.close(fd) + sock.close() + raise RuntimeError(f"the fork template refused: {reply.get('error') if isinstance(reply, dict) else reply}") + leftovers.record(pid, os.path.basename(argv[1]) if len(argv) > 1 else "agent") + sock.setblocking(False) + control = await asyncio.open_unix_connection(sock=sock) + stdout = asyncio.StreamReader(limit=_READER_LIMIT, loop=loop) + await loop.connect_read_pipe( + lambda: asyncio.StreamReaderProtocol(stdout, loop=loop), os.fdopen(stdout_r, "rb", 0) + ) + stderr = asyncio.StreamReader(limit=_READER_LIMIT, loop=loop) + await loop.connect_read_pipe( + lambda: asyncio.StreamReaderProtocol(stderr, loop=loop), os.fdopen(stderr_r, "rb", 0) + ) + transport, protocol = await loop.connect_write_pipe( + lambda: asyncio.streams.FlowControlMixin(loop=loop), os.fdopen(stdin_w, "wb", 0) + ) + stdin = asyncio.StreamWriter(transport, protocol, None, loop) + return ForkedEndpoint(pid, stdin=stdin, stdout=stdout, stderr=stderr, control=control, template=self) + + def close_when_idle(self, after: float | None) -> None: + """Close the template once it has had no agent for ``after`` seconds; ``None`` keeps it.""" + if after is None or self._children or not self.running: + return + if self._idle_task is not None and not self._idle_task.done(): + return + + async def wait_then_close() -> None: + await asyncio.sleep(after) + async with self._lock: + if not self._children and not self._forking: + await self.close() + + self._idle_task = asyncio.get_running_loop().create_task(wait_then_close(), name="fork-template-idle") + + def abandon(self) -> None: + """End the template without the loop it was started on, which may be closed by now. + + The agents it forked end with it (they armed its death signal); a loop + that has ended has no use for them either, their pipes being its. + """ + from raven.acp_client import leftovers + + proc, self._proc = self._proc, None + if proc is not None: + with contextlib.suppress(ProcessLookupError, PermissionError): + os.killpg(proc.pid, signal.SIGKILL) + leftovers.forget(proc.pid) + if self._dir is not None: + import shutil + + shutil.rmtree(self._dir, ignore_errors=True) + self._dir = None + + async def close(self) -> None: + from raven.acp_client import leftovers + + proc, self._proc = self._proc, None + if proc is not None and proc.returncode is None: + # Letting go of its stdin is how the template is asked to stop; the + # group kill after the grace is for one that does not. + if proc.stdin is not None: + with contextlib.suppress(Exception): + proc.stdin.close() + with contextlib.suppress(asyncio.TimeoutError, ProcessLookupError): + await asyncio.wait_for(proc.wait(), timeout=_TEARDOWN_GRACE_S) + with contextlib.suppress(ProcessLookupError, PermissionError): + os.killpg(proc.pid, signal.SIGKILL) + if proc is not None: + leftovers.forget(proc.pid) + transport = getattr(proc, "_transport", None) + if transport is not None: + with contextlib.suppress(Exception): + transport.close() + if self._log_task is not None: + self._log_task.cancel() + self._log_task = None + if self._dir is not None: + import shutil + + shutil.rmtree(self._dir, ignore_errors=True) + self._dir = None + + +_TEMPLATE: ForkTemplate | None = None + + +def template_available() -> bool: + """Whether this system can fork agents from a template at all.""" + return sys.platform.startswith("linux") and hasattr(socket, "send_fds") and hasattr(os, "fork") + + +async def fork_launch( + argv: list[str], env: dict[str, str], cwd: str | None, *, base_env: dict[str, str], skip: Sequence[str] = () +) -> ForkedEndpoint | None: + """The agent ``argv`` names, forked from this host's template; ``None`` to launch it the ordinary way. + + ``base_env`` is what the template itself starts with: the environment every + agent shares, before any one agent's own entries. ``skip`` names module + prefixes the template does not load (LiteLLM, when the agents' model calls + go to the host). A template that cannot start is not tried again in this + process; one fork that fails is logged and this launch goes the ordinary way. + """ + from loguru import logger + + global _TEMPLATE + if not template_available() or not raven_agent_argv(argv): + return None + if cwd is not None and not os.path.isdir(cwd): + # The ordinary launch refuses a missing directory in the host, where the + # caller hears why; a fork would only see the agent die. + return None + if _TEMPLATE is not None and _TEMPLATE.loop is not asyncio.get_running_loop(): + # Started under a loop that has ended (a CLI turn is one asyncio.run). + _TEMPLATE.abandon() + _TEMPLATE = None + if _TEMPLATE is None or (_TEMPLATE._skip != tuple(skip) and not _TEMPLATE.children): + if _TEMPLATE is not None: + await _TEMPLATE.close() + _TEMPLATE = ForkTemplate(skip=skip) + if _TEMPLATE.broken: + return None + try: + return await _TEMPLATE.spawn(argv, env, cwd, base_env) + except Exception as exc: # noqa: BLE001 - the ordinary launch is always there to fall back to + logger.warning("fork template: launching {} the ordinary way: {}", argv[1:2], exc) + return None + + +async def close_template() -> None: + """Close the template. For process shutdown and for tests.""" + global _TEMPLATE + template, _TEMPLATE = _TEMPLATE, None + if template is not None and template.loop is not asyncio.get_running_loop(): + template.abandon() + elif template is not None: + await template.close() + + +# ---- the template's side ------------------------------------------------------- + + +#: Not loaded ahead whatever was learned: the entry point itself, and this +#: module under its import name (it already runs as ``__main__``). +_NEVER_LOAD = frozenset({"__main__", "raven.acp_client.fork_template"}) + + +def product_of(argv: Sequence[str]) -> str: + """The name an agent's modules are learned under: its product folder, or ``raven`` for ``-m raven``.""" + if len(argv) > 1 and argv[1] != "-m": + return Path(argv[1]).parent.name or "raven" + return "raven" + + +def preload_names(learned: Any) -> list[str]: + """What a template loads ahead of the agents: the modules they share, or one product's own. + + A module only one product imports is shared by nobody -- a host normally + runs one process per product -- so loading it ahead would hold it in the + template, unused, while any other product runs. Measured on the five + shipped products: they share about two thousand modules, each adds a few + dozen of its own, and the deck engine adds about a thousand. With only one + product learned yet, its own set is all there is to share. + """ + if not isinstance(learned, dict) or learned.get("python") != list(sys.version_info[:2]): + return [] + products = learned.get("products") + if not isinstance(products, dict): + return [] + sets = [{n for n in names if isinstance(n, str)} for names in products.values() if isinstance(names, list)] + if len(sets) == 1: + return sorted(sets[0]) + counts = collections.Counter(name for names in sets for name in names) + return sorted(name for name, seen in counts.items() if seen >= 2) + + +def _load(seed: Sequence[str], learned: Path, skip: Sequence[str]) -> int: + """Import what the agents will import; the module count afterwards.""" + names = list(seed) + try: + names += preload_names(json.loads(learned.read_text(encoding="utf-8"))) + except (OSError, ValueError): + pass + for name in names: + if name in sys.modules or name in _NEVER_LOAD or name.endswith(".__main__") or name.startswith(tuple(skip)): + continue + try: + importlib.import_module(name) + except BaseException: # noqa: BLE001 - an import that fails here fails the same way in the agent + continue + # Not an import: every agent loads the token estimator's table on its first + # count, and the table's core is read-only once built, so loaded here it is + # shared by every agent for as long as they run. + with contextlib.suppress(Exception): + from raven.utils.tokens import encoding + + encoding() + return len(sys.modules) + + +def _learn(path: Path, skip: Sequence[str], product: str) -> None: + """Add the modules this agent imported to what its product is known to import.""" + try: + import fcntl + + mine = { + name + for name, module in list(sys.modules.items()) + if getattr(module, "__spec__", None) is not None + and name != "__main__" + and not name.endswith(".__main__") + and not name.startswith(tuple(skip)) + } + path.parent.mkdir(parents=True, exist_ok=True) + with open(path.with_suffix(".lock"), "a") as lock: + fcntl.flock(lock, fcntl.LOCK_EX) + try: + data = json.loads(path.read_text(encoding="utf-8")) + except (OSError, ValueError): + data = {} + products = data.get("products") if isinstance(data, dict) else None + if not isinstance(products, dict) or data.get("python") != list(sys.version_info[:2]): + products = {} + known = products.get(product) + known = {n for n in known if isinstance(n, str)} if isinstance(known, list) else set() + merged = sorted(known | mine)[:_LEARNED_CAP] + if merged != sorted(known): + products[product] = merged + tmp = path.with_suffix(f".{os.getpid()}.tmp") + tmp.write_text( + json.dumps({"python": list(sys.version_info[:2]), "products": products}), encoding="utf-8" + ) + os.replace(tmp, path) + except Exception: # noqa: BLE001 - learning is an optimisation for the next template + pass + + +def _arm_parent_death( + parent: int, +) -> None: # pragma: no cover - runs only in a forked agent; see tests/test_acp_fork_template.py + """End this process when the template does; it ends with the host the same way.""" + try: + import ctypes + + armed = ctypes.CDLL(None).prctl(1, signal.SIGKILL, 0, 0, 0) == 0 + except Exception: # noqa: BLE001 - without the signal the agent still ends on its closed stdin + armed = False + if armed and os.getppid() != parent: + os._exit(1) + + +def _finish(code: int) -> None: # pragma: no cover - runs only in a forked agent; see tests/test_acp_fork_template.py + """End the agent the way an interpreter ends: threads joined, exit handlers run, streams flushed. + + Never by unwinding: the frames below the agent's program are the template's, + and a launcher's own exit handling must not run after Raven, which it would + never have reached past a real ``exec``. + """ + import atexit + + with contextlib.suppress(Exception): + threading._shutdown() # type: ignore[attr-defined] + with contextlib.suppress(Exception): + atexit._run_exitfuncs() + for stream in (sys.stdout, sys.stderr): + with contextlib.suppress(Exception): + stream.flush() + os._exit(code) + + +def _exit_code(exc: BaseException) -> int: + """What an interpreter exits with when ``exc`` ends its program, printing what it would print.""" + if isinstance(exc, SystemExit): + if exc.code is None: + return 0 + if isinstance(exc.code, int): + return exc.code + with contextlib.suppress(Exception): + print(exc.code, file=sys.stderr) + return 1 + with contextlib.suppress(Exception): + sys.excepthook(type(exc), exc, exc.__traceback__) + return 1 + + +def _run_as_raven(args: Sequence[str]) -> None: + """``python -m raven `` in this process, then the end of it.""" + import runpy + + sys.argv = ["raven", *args] + sys.path.insert(0, os.getcwd()) + try: + runpy.run_module("raven", run_name="__main__", alter_sys=True) + code = 0 + except BaseException as exc: # noqa: BLE001 - the program's end, however it came + code = _exit_code(exc) + _finish(code) + + +def _run_agent(argv: list[str]) -> None: + """Run the agent's program here, as a fresh interpreter would have run it, and end the process.""" + import runpy + + real_execv = os.execv + + def execv(path: Any, args: Any) -> None: + args = [str(a) for a in args] + if os.path.abspath(str(path)) == os.path.abspath(sys.executable) and args[1:3] == ["-m", "raven"]: + os.execv = real_execv + _run_as_raven(args[3:]) + real_execv(path, args) + + os.execv = execv + if argv[1] == "-m": + _run_as_raven(argv[3:]) + script = os.path.abspath(argv[1]) + sys.argv = [argv[1], *argv[2:]] + sys.path.insert(0, os.path.dirname(script)) + try: + runpy.run_path(script, run_name="__main__") + code = 0 + except BaseException as exc: # noqa: BLE001 - the program's end, however it came + code = _exit_code(exc) + _finish(code) + + +def _become( # pragma: no cover - runs only in a forked agent; see tests/test_acp_fork_template.py + request: dict[str, Any], fds: list[int], template: int, learned: Path, skip: Sequence[str], learn_after: float +) -> None: + """Turn this freshly forked child into the agent the request names. Does not return.""" + os.setsid() + _arm_parent_death(template) + for signum in (signal.SIGCHLD, signal.SIGTERM): + signal.signal(signum, signal.SIG_DFL) + with contextlib.suppress(ValueError): + signal.set_wakeup_fd(-1) + for target, fd in zip((0, 1, 2), fds, strict=True): + os.dup2(fd, target) + os.close(fd) + os.chdir(request["cwd"]) + os.environ.clear() + os.environ.update({str(k): str(v) for k, v in (request.get("env") or {}).items()}) + timer = threading.Timer(learn_after, _learn, args=(learned, skip, product_of(request["argv"]))) + timer.daemon = True + timer.start() + _run_agent([str(a) for a in request["argv"]]) + + +def _serve(path: str, *, lifeline: int) -> tuple[dict[str, Any], list[int]] | None: + """Fork agents on request until the host lets go of ``lifeline``. + + In a forked child this returns the request it is to become; in the template + it returns ``None`` once ``lifeline`` -- the template's stdin, whose other end + only the host holds -- reads as closed. + """ + import selectors + + from raven.providers import relay_wire + + template = os.getpid() + listener = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) + listener.bind(path) + listener.listen(64) + wake_r, wake_w = os.pipe() + os.set_blocking(wake_w, False) + signal.set_wakeup_fd(wake_w) + signal.signal(signal.SIGCHLD, lambda *_: None) + selector = selectors.DefaultSelector() + selector.register(listener, selectors.EVENT_READ) + selector.register(wake_r, selectors.EVENT_READ) + selector.register(lifeline, selectors.EVENT_READ) + reporting: dict[int, socket.socket] = {} + + def let_go(*, child: bool) -> None: + selector.close() + listener.close() + for conn in reporting.values(): + conn.close() + if not child: + signal.set_wakeup_fd(-1) + signal.signal(signal.SIGCHLD, signal.SIG_DFL) + os.close(wake_r) + os.close(wake_w) + + print( + json.dumps( + {"ready": True, "pid": template, "modules": len(sys.modules), "seconds": time.monotonic() - _STARTED} + ), + flush=True, + ) + while True: + for key, _ in selector.select(): + if key.fileobj == lifeline: + if not os.read(lifeline, 4096): + let_go(child=False) + return None + continue + if key.fileobj == wake_r: + with contextlib.suppress(BlockingIOError): + os.read(wake_r, 4096) + _report_exits(reporting) + continue + conn, _ = listener.accept() + conn.settimeout(_REQUEST_TIMEOUT_S) + fds: list[int] = [] + try: + _, fds, _, _ = socket.recv_fds(conn, 1, 3) + request = relay_wire.recv_frame(conn) + if len(fds) != 3 or not isinstance(request.get("argv"), list) or not request.get("cwd"): + raise ValueError("a fork request carries an argv, a cwd and three descriptors") + if threading.active_count() != 1: + raise RuntimeError("the template is running a thread and cannot fork safely") + pid = os.fork() + except Exception as exc: # noqa: BLE001 - one bad request must not end the template + for fd in fds: + os.close(fd) + with contextlib.suppress(OSError): + relay_wire.send_frame(conn, {"error": f"{type(exc).__name__}: {exc}"}) + conn.close() + continue + if pid == 0: # pragma: no cover - the forked agent's side of the fork + conn.close() + let_go(child=True) + return request, fds + for fd in fds: + os.close(fd) + try: + relay_wire.send_frame(conn, {"pid": pid}) + reporting[pid] = conn + except OSError: + conn.close() + + +def _report_exits(reporting: dict[int, socket.socket]) -> None: + """Reap every child that has ended and tell its requester how.""" + from raven.providers import relay_wire + + while True: + try: + pid, status = os.waitpid(-1, os.WNOHANG) + except ChildProcessError: + return + if pid == 0: + return + conn = reporting.pop(pid, None) + if conn is not None: + with contextlib.suppress(OSError): + relay_wire.send_frame(conn, {"exit": os.waitstatus_to_exitcode(status)}) + conn.close() + + +_STARTED = time.monotonic() + + +def main(argv: Sequence[str] | None = None) -> int: + parser = argparse.ArgumentParser(prog="raven-fork-template") + parser.add_argument("--socket", required=True) + parser.add_argument("--learned", required=True) + parser.add_argument("--skip", action="append", default=[]) + parser.add_argument("--learn-after", type=float, default=LEARN_AFTER_S) + args = parser.parse_args(argv) + learned = Path(args.learned) + # stdout is the channel the host reads the ready line from; whatever an + # import prints goes to stderr instead, where the host only logs it. + sys.stdout.flush() + ready_channel = os.dup(1) + os.dup2(2, 1) + _load(SEED_MODULES, learned, args.skip) + sys.stdout.flush() + os.dup2(ready_channel, 1) + os.close(ready_channel) + gc.collect() + gc.freeze() + if threading.active_count() != 1: + names = ", ".join(t.name for t in threading.enumerate() if t is not threading.current_thread()) + print(json.dumps({"ready": False, "reason": f"an import started a thread ({names})"}), flush=True) + return 3 + forked = _serve(args.socket, lifeline=0) + if forked is None: + return 0 + request, fds = forked + _become(request, fds, os.getppid(), learned, args.skip, args.learn_after) + return 1 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/raven/acp_client/pool.py b/raven/acp_client/pool.py index f119947fc..e4765808f 100644 --- a/raven/acp_client/pool.py +++ b/raven/acp_client/pool.py @@ -24,6 +24,7 @@ import asyncio import hashlib import json +import time from collections.abc import Awaitable, Callable from typing import Any @@ -57,6 +58,13 @@ # transport, but no longer reaps the child. _CLOSE_TIMEOUT_S = 10.0 +# How often the idle check runs, as a share of the idle budget and within +# bounds: often enough that a closed agent's memory comes back close to the +# budget the operator set, rarely enough that the check itself is free. +_IDLE_CHECK_SHARE = 4 +_IDLE_CHECK_MIN_S = 1.0 +_IDLE_CHECK_MAX_S = 60.0 + # How long a connection may sit silent before a reuse pays for a probe. Under # this, the connection spoke recently and is handed out on that evidence alone; # past it, one round trip settles the question a 120s session open would @@ -334,11 +342,19 @@ def __init__( actually serving requests, where a stored snapshot only describes the one that was verified.""" self._session_locks: dict[str, asyncio.Lock] = {} + self.leased_at = time.monotonic() + """When the pool last handed this connection out. A caller sends its + request right after taking it, so the idle check leaves a connection + alone for a full budget after a lease even if the agent has said nothing.""" @property def alive(self) -> bool: return self.client.alive + def idle_for(self, now: float) -> float: + """How long nothing has happened here: no frame from the agent and no lease, as of ``now``.""" + return min(self.client.idle_seconds, now - self.leased_at) + def session_lock(self, session_id: str) -> asyncio.Lock: """Serialises turns on one session id. @@ -367,6 +383,7 @@ class AcpConnectionPool: def __init__(self) -> None: self._connections: dict[str, dict[str, _Connection]] = {} self._locks: dict[str, asyncio.Lock] = {} + self._idle_task: asyncio.Task | None = None def connections(self, name: str) -> list[_Connection]: """Every connection held under ``name``, live or not.""" @@ -444,6 +461,7 @@ async def acquire( existing = held.get(bkey) if existing is not None and existing.alive and existing.launch_key == key: if await self._responsive(existing): + existing.leased_at = time.monotonic() return existing # Alive by its process and unchanged by its key, yet not answering: # a wedged worker is replaced rather than hit again (see _responsive). @@ -533,6 +551,7 @@ async def acquire( config=config, ) held[bkey] = connection + self._watch_idle() return connection async def adopt( @@ -582,6 +601,7 @@ async def adopt( if existing is not None and existing.alive: return connection, False held[bkey] = connection + self._watch_idle() return connection, True async def drop(self, name: str) -> None: @@ -619,6 +639,78 @@ async def _responsive(existing: _Connection) -> bool: ) return ok + @staticmethod + def _idle_budget() -> float | None: + """``subagents.runtime.idleCloseS`` as the host's config has it now; ``None`` when closing is off.""" + from raven.config.live import default_live, subagent_runtime + + budget = subagent_runtime(default_live()).idle_close_s + return budget if budget is not None and budget > 0 else None + + @staticmethod + def _closable(connection: _Connection, budget: float, now: float) -> bool: + """Whether closing this connection now loses nothing a later turn could want. + + Nothing in flight either way, quiet for the whole budget, and an agent + that can reopen what it held: a session this connection served is + resumed in the next process with ``session/load``, so only an agent + that advertises it, or one that never opened a session, is closed. + """ + client = connection.client + if not connection.alive or getattr(client, "busy", True): + return False + if connection.idle_for(now) < budget: + return False + return handshake_of(connection.initialize).can_load or not connection.router.seen + + async def close_idle(self, budget: float | None = None) -> list[str]: + """Close every connection idle for ``budget`` seconds; the names closed. + + ``budget`` defaults to the host's live setting. Each candidate is + judged again under its agent's lock, the one ``acquire`` holds while it + hands a connection out, so a connection being reused is never closed + underneath its caller. + """ + budget = self._idle_budget() if budget is None else budget + if budget is None: + return [] + closing: list[tuple[str, _Connection]] = [] + for name in list(self._connections): + async with self._lock_for(name): + held = self._connections.get(name, {}) + now = time.monotonic() + for bkey, connection in list(held.items()): + # A connection whose agent already ended goes too: nothing + # would hand it out again, and holding it keeps this check running. + if not connection.alive or self._closable(connection, budget, now): + held.pop(bkey, None) + closing.append((name, connection)) + if not held: + self._connections.pop(name, None) + for name, connection in closing: + if connection.alive: + logger.info("acp agent {!r}: idle for {:.0f}s, closing its process", name, budget) + await connection.client.close() + return [name for name, _ in closing] + + def _watch_idle(self) -> None: + """Start the idle check if it is not running; it ends once nothing is held.""" + loop = asyncio.get_running_loop() + task = self._idle_task + if task is not None and not task.done() and task.get_loop() is loop: + return + self._idle_task = loop.create_task(self._idle_loop(), name="acp-idle-close") + + async def _idle_loop(self) -> None: + while self._connections: + budget = self._idle_budget() + share = budget / _IDLE_CHECK_SHARE if budget is not None else _IDLE_CHECK_MAX_S + await asyncio.sleep(min(max(share, _IDLE_CHECK_MIN_S), _IDLE_CHECK_MAX_S)) + try: + await self.close_idle(budget) + except Exception as exc: # noqa: BLE001 - a failed check must not end the next one + logger.warning("acp idle check failed: {}", exc) + async def close_all(self) -> None: """Close every connection. For process shutdown and for tests. @@ -628,6 +720,9 @@ async def close_all(self) -> None: request runs. Serially and unbounded, one such connection held every other one -- and the exit waiting behind them -- for that whole time. """ + task, self._idle_task = self._idle_task, None + if task is not None and not task.get_loop().is_closed(): + task.cancel() names = list(self._connections) if not names: return @@ -712,6 +807,13 @@ async def close_pool() -> None: if _POOL is not None: await _POOL.close_all() _POOL = None + # What the pool's agents shared, started with the first of them and of no + # use once they are gone. + from raven.acp_client.fork_template import close_template + from raven.providers.relay_server import close_relay + + await close_template() + await close_relay() __all__ = ["AcpConnectionPool", "SessionSink", "close_pool", "get_pool"] diff --git a/raven/cli/acp_commands.py b/raven/cli/acp_commands.py index ba0fcb08f..b8849a4b6 100644 --- a/raven/cli/acp_commands.py +++ b/raven/cli/acp_commands.py @@ -22,6 +22,7 @@ from raven.acp.stdio import open_stdin as _open_stdin from raven.cli._log_file import redirect_loguru_to_file from raven.config.loader import set_config_path +from raven.providers.relay import target as relay_target from raven.utils import asyncio_runner as bounded_asyncio from raven.utils.heap import start_trimmer @@ -119,6 +120,9 @@ def acp( """ if ctx.invoked_subcommand is not None: return + # Taken out of the environment before anything this agent starts -- an MCP + # server, a shell command -- can inherit the host's LLM relay offer with it. + relay_target() if config: _use_config(config) try: diff --git a/raven/config/live.py b/raven/config/live.py index 43b324385..adb0562bf 100644 --- a/raven/config/live.py +++ b/raven/config/live.py @@ -37,11 +37,14 @@ from contextvars import ContextVar from functools import lru_cache from pathlib import Path -from typing import Any, Callable, Iterator +from typing import TYPE_CHECKING, Any, Callable, Iterator from loguru import logger from pydantic.alias_generators import to_camel +if TYPE_CHECKING: + from raven.config.schema import SubagentRuntimeConfig + __all__ = [ "LiveConfig", "hold_for_this_turn", @@ -65,6 +68,7 @@ "permissions_config", "routing_profile", "skill_gate_pin", + "subagent_runtime", "web_provider_keys", "web_search_key", ] @@ -582,6 +586,27 @@ def memory_top_k(live: LiveConfig) -> int | None: return None +def subagent_runtime(live: LiveConfig) -> SubagentRuntimeConfig: + """``subagents.runtime`` validated, keeping the last valid answer; the defaults without one. + + Read by the host when it launches an agent and on each idle check, so the + next launch and the next check follow an edit. A rejected edit keeps the + last admitted answer rather than falling back to the defaults, which would + switch sharing on or off behind the operator's back. + """ + from raven.config.schema import SubagentRuntimeConfig + + section = live.get("subagents.runtime") + present = isinstance(section, dict) + value = None + if present: + try: + value = SubagentRuntimeConfig.model_validate(section) + except ValueError: + value = None + return _admit(live, "subagent_runtime", present=present, value=value) or SubagentRuntimeConfig() + + def exec_timeout(live: LiveConfig) -> int | None: """``tools.exec.timeout`` as the file has it, or None for "no answer".""" # One spelling: ``exec`` and ``timeout`` are the same word either way. diff --git a/raven/config/schema.py b/raven/config/schema.py index dd911cfd0..49c15bfb1 100644 --- a/raven/config/schema.py +++ b/raven/config/schema.py @@ -2555,6 +2555,37 @@ class PlaybookConfig(Base): that send no mode needs those clients to send one.""" +class SubagentRuntimeConfig(Base): + """How the agent processes this host launches share its memory. + + Read live by the host that launches them, so a change reaches the next + launch (and the next idle check) without a restart. None of it changes what + an agent does; it changes how many processes are kept around and what each + one has to load. + """ + + idle_close_s: float | None = 600.0 + """Seconds a pooled ACP agent may sit with no request in flight and no frame + either way before its process is closed. ``None`` or ``0`` keeps every agent + until the host exits. Only an agent that can reopen its sessions + (``loadSession``) is closed, so a later turn resumes the conversation in a + fresh process; an agent holding a wake is still answering a prompt and is + never idle.""" + + llm_relay: bool = True + """Serve a Raven agent's LiteLLM calls from this process instead of loading + LiteLLM in each agent. The agent builds the request, retries and records + usage itself; only the call to the model, and the catalogue questions that + need LiteLLM's tables, cross to the host. An agent that cannot reach the host + loads LiteLLM as before.""" + + fork_template: bool = True + """On Linux, start Raven agents by forking a process that has already loaded + their code, so the loaded modules are shared between them instead of loaded + once per agent. Any agent the template cannot start is launched the ordinary + way. Has no effect on other systems.""" + + class SubagentsConfig(Base): """The one table of agents raven can dispatch to. @@ -2572,6 +2603,7 @@ class SubagentsConfig(Base): default_factory=list, validation_alias=AliasChoices("agents", "thirdParty", "third_party"), ) + runtime: SubagentRuntimeConfig = Field(default_factory=SubagentRuntimeConfig) @model_validator(mode="after") def _dedupe_names(self) -> "SubagentsConfig": diff --git a/raven/providers/anthropic_messages_provider.py b/raven/providers/anthropic_messages_provider.py index 1d93aabe3..f41f175cc 100644 --- a/raven/providers/anthropic_messages_provider.py +++ b/raven/providers/anthropic_messages_provider.py @@ -39,6 +39,7 @@ ) from raven.providers.prompt_cache import accepts_cache_control from raven.providers.rates import DEFAULT_MAX_OUTPUT_TOKENS +from raven.providers.relay import relayed from raven.providers.tool_names import normalized_tool_name from raven.providers.usage import reported_cost, token_count @@ -421,6 +422,7 @@ def _visible_reserve(max_tokens: int) -> int: return max(1024, max_tokens // 2) +@relayed(keep=bool) def _litellm_budgets() -> dict[str, int]: """litellm's own effort-to-budget table, or empty when it cannot be imported. diff --git a/raven/providers/capabilities.py b/raven/providers/capabilities.py index d05cb9724..c4f48a088 100644 --- a/raven/providers/capabilities.py +++ b/raven/providers/capabilities.py @@ -26,6 +26,8 @@ from loguru import logger +from raven.providers.relay import relayed + if TYPE_CHECKING: from raven.providers.registry import ProviderSpec @@ -105,23 +107,32 @@ def supports_image_tool_result(provider: Any, model: str, spec: "ProviderSpec | # admits ResponseInputImageContentParam. return True - from raven.providers.litellm_provider import LiteLLMProvider + from raven.providers.relay_provider import is_litellm_transport - if not isinstance(provider, LiteLLMProvider): + if not is_litellm_transport(provider): # Azure direct and any other bespoke transport speak Chat Completions. return False + target = _litellm_target(model) + if target is None: + return False + + if target in GATEWAY_TARGETS: + return _gateway_route(model).startswith(GATEWAY_IMAGE_TOOL_RESULT_PREFIXES) + return target in IMAGE_TOOL_RESULT_TARGETS + + +@relayed() +def _litellm_target(model: str) -> str | None: + """The upstream LiteLLM routes ``model`` to, or None when it cannot say.""" try: import litellm _, target, _, _ = litellm.get_llm_provider(model=model) except Exception as e: logger.debug("supports_image_tool_result: cannot resolve target for {}: {}", model, e) - return False - - if target in GATEWAY_TARGETS: - return _gateway_route(model).startswith(GATEWAY_IMAGE_TOOL_RESULT_PREFIXES) - return target in IMAGE_TOOL_RESULT_TARGETS + return None + return target def _model_id_is_caller_chosen(model: str, provider: Any, spec: "ProviderSpec | None") -> bool: diff --git a/raven/providers/common_models.py b/raven/providers/common_models.py index 3b3fdd4e5..edeb8d218 100644 --- a/raven/providers/common_models.py +++ b/raven/providers/common_models.py @@ -19,6 +19,8 @@ from functools import lru_cache +from raven.providers.relay import relayed + def common_models_for(slug: str) -> list[str]: """The curated shortlist for ``slug``, as ids a caller can store and send. @@ -48,6 +50,7 @@ def _is_priced_template(model: str) -> bool: @lru_cache(maxsize=1) +@relayed(keep=bool, shape=lambda index: {provider: tuple(models) for provider, models in (index or {}).items()}) def _cached_chat_models_by_provider() -> dict[str, tuple[str, ...]]: """LiteLLM's own catalogue, indexed by the provider it belongs to. diff --git a/raven/providers/factory.py b/raven/providers/factory.py index 6546ae7e0..6fb429d58 100644 --- a/raven/providers/factory.py +++ b/raven/providers/factory.py @@ -165,7 +165,7 @@ def make_native(ep=None): else: from raven.providers.capabilities import wire_overrides from raven.providers.endpoints import provider_endpoints - from raven.providers.litellm_provider import LiteLLMProvider + from raven.providers.relay_provider import litellm_provider if provider_name and configured_protocol(p, model) is None and inferred_protocol(model) != "chat": # The stepping-aside itself, logged where the chat client is built: @@ -183,7 +183,7 @@ def make_native(ep=None): from raven.providers.endpoint_rotor import EndpointRotorProvider def make_inner(ep): - return LiteLLMProvider( + return litellm_provider( api_key=ep.api_key, # ``ep.api_base`` already carries the section's flat address # when the endpoint named none of its own (see @@ -208,7 +208,7 @@ def make_inner(ep): ) elif eps: extra_body = wire_overrides(provider_name, model) or None - provider = LiteLLMProvider( + provider = litellm_provider( api_key=eps[0].api_key, # Same fallback as ``make_inner`` above: only reached when the # flat address is empty too, for a provider's usable spec @@ -223,7 +223,7 @@ def make_inner(ep): ) else: extra_body = wire_overrides(provider_name, model) or None - provider = LiteLLMProvider( + provider = litellm_provider( api_key=p.effective_api_key if p else None, api_base=config.get_api_base(model), default_model=model, diff --git a/raven/providers/litellm_provider.py b/raven/providers/litellm_provider.py index e0e718413..556bf024f 100644 --- a/raven/providers/litellm_provider.py +++ b/raven/providers/litellm_provider.py @@ -3,10 +3,8 @@ import asyncio import hashlib import json -import os import secrets import string -import uuid import warnings from collections.abc import AsyncIterator from typing import Any @@ -34,13 +32,16 @@ stream_first_byte_budget, ) from raven.providers.litellm_setup import import_litellm +from raven.providers.litellm_shape import ( # noqa: F401 - re-exported for callers of this module + _ANTHROPIC_EXTRA_KEYS, + _OPENROUTER_ATTRIBUTION, + LiteLLMShape, + session_affinity_headers, +) from raven.providers.reasoning import split_orphan_think from raven.providers.registry import ( - canonical_provider_name, find_by_keywords, find_by_model, - find_by_name, - find_gateway, ) from raven.providers.tool_names import normalized_tool_name from raven.providers.transport_failure import ( @@ -52,7 +53,6 @@ without_sentinels, ) from raven.providers.usage import merge_usage, reported_cost, token_count -from raven.providers.wire import wire_model from raven.security.tls import is_untrusted_certificate litellm = import_litellm() @@ -75,19 +75,8 @@ # Standard chat-completion message keys. _ALLOWED_MSG_KEYS = frozenset({"role", "content", "tool_calls", "tool_call_id", "name", "reasoning_content"}) -_ANTHROPIC_EXTRA_KEYS = frozenset({"thinking_blocks"}) _ALNUM = string.ascii_letters + string.digits -# LiteLLM defaults to X-Title="liteLLM" / HTTP-Referer="https://litellm.ai" for OpenRouter -# requests, which would credit traffic to liteLLM instead of Raven on openrouter.ai/apps. -# Explicit headers here override those defaults; user-supplied extra_headers win over these. -_OPENROUTER_ATTRIBUTION: dict[str, str] = { - "HTTP-Referer": "https://raven.evermind.ai", - "X-Title": "Raven Agent", - "X-OpenRouter-Title": "Raven Agent", - "X-OpenRouter-Categories": "cli-agent,personal-agent", -} - def _short_tool_id() -> str: """Generate a 9-char alphanumeric ID compatible with all providers (incl. Mistral).""" @@ -131,15 +120,6 @@ def _merge_extra_body(kwargs: dict[str, Any], wire_extra_body: dict[str, Any]) - kwargs["extra_body"] = wire_extra_body -def session_affinity_headers() -> dict[str, str]: - """Headers pinning one caller to one backend replica. - - Self-hosted OpenAI-compatible backends (vLLM and friends) route by this - header, so a stable value per provider instance keeps prefix-cache hits warm. - """ - return {"x-session-affinity": uuid.uuid4().hex} - - def _usage_field(value: Any, key: str) -> Any: return value.get(key) if isinstance(value, dict) else getattr(value, key, None) @@ -195,7 +175,7 @@ def _reasoning_tokens(usage: Any) -> int | None: return None -class LiteLLMProvider(LLMProvider): +class LiteLLMProvider(LiteLLMShape): """ LLM provider using LiteLLM for multi-provider support. @@ -204,258 +184,18 @@ class LiteLLMProvider(LLMProvider): (see providers/registry.py) — no if-elif chains needed here. """ - api_protocol = "chat" - - def __init__( - self, - api_key: str | None = None, - api_base: str | None = None, - default_model: str = "anthropic/claude-opus-4-5", - extra_headers: dict[str, str] | None = None, - provider_name: str | None = None, - disable_auto_cache_control: bool = False, - extra_body: dict[str, Any] | None = None, - model_overrides: dict[str, dict[str, Any]] | None = None, - tls_verify: bool = True, - *, - unparsed_reasoning: bool | None = None, - ): - addressed_spec = find_by_name(provider_name) or find_by_model(default_model) - if api_base and addressed_spec and addressed_spec.strip_api_base_trailing_slash: - api_base = api_base.rstrip("/") - super().__init__(api_key, api_base) - # None: derive from the resolved spec, as emits_unparsed_reasoning always - # did. An explicit bool overrides that derivation outright -- for a - # caller that already knows the answer and for which the spec would - # guess wrong, e.g. a per-model routing endpoint built with - # provider_name="custom" for its api_base/api_key shape alone, not - # because the backend behind it is a self-hosted inference server. - self._unparsed_reasoning = unparsed_reasoning - self.default_model = default_model - self.extra_headers = extra_headers or {} - # When a TokenStrategy (e.g. CacheOptimizer) handles cache_control - # placement upstream, turn this on so the provider doesn't also - # stamp its own breakpoints on top. - self.disable_auto_cache_control = disable_auto_cache_control - # Provider-specific request body extras forwarded verbatim to LiteLLM. - # Common use: OpenRouter routing affinity to keep prompt-cache hits warm, - # extra_body={"provider": {"order": ["Anthropic"], "allow_fallbacks": False}} - self.extra_body = extra_body or {} - # User-configured per-model parameter overrides; win over the registry's. - self.model_overrides = model_overrides or {} - # A section that opted out of certificate verification sends its calls - # through a client built here with verification off; True (the default) - # lets LiteLLM build its own verifying client exactly as before. - self.tls_verify = tls_verify + def __init__(self, *args: Any, **kwargs: Any) -> None: + super().__init__(*args, **kwargs) # Built lazily on the event loop the first call runs on, and shared by - # every call this provider makes afterwards. + # every call this provider makes afterwards (``tls_verify`` off only). self._verify_off_client: Any | None = None self._verify_off_http_handler: Any | None = None # Whether LiteLLM serves this provider's route through its own HTTP # handler rather than the OpenAI SDK; learned, see `_learn_route_from`. self._on_litellm_http_route = False - - # Detect gateway / local deployment. - # provider_name (from config key) is the primary signal; - # api_key / api_base are fallback for auto-detection. - # Kept because the id alone cannot say where a request goes: a bare - # `anthropic/claude-...` sent through this client reads as Anthropic's - # wire, which is not the wire it will travel on. - self._provider_name = provider_name or "" - self._gateway = find_gateway(provider_name, api_key, api_base) - if self._gateway and self._gateway.name == "openrouter": - self.extra_headers = {**_OPENROUTER_ATTRIBUTION, **self.extra_headers} - - if api_key: - self._setup_env(api_key, api_base, default_model) - # Drop unsupported parameters for providers (e.g., gpt-5 rejects some params) litellm.drop_params = True - def _setup_env(self, api_key: str, api_base: str | None, model: str) -> None: - """Set environment variables based on detected provider.""" - spec = self._gateway or find_by_model(model) - if not spec: - # No spec: the key still reaches LiteLLM as an explicit api_key - # kwarg on every call, so nothing needs to go into the environment. - return - if not spec.env_key: - # OAuth/provider-only specs (for example: openai_codex) - return - - # Gateway/local overrides existing env; standard provider doesn't - if self._gateway: - os.environ[spec.env_key] = api_key - else: - os.environ.setdefault(spec.env_key, api_key) - - # Resolve env_extras placeholders: - # {api_key} → user's API key - # {api_base} → user's api_base, falling back to spec.default_api_base - effective_base = api_base or spec.default_api_base - for env_name, env_val in spec.env_extras: - resolved = env_val.replace("{api_key}", api_key) - resolved = resolved.replace("{api_base}", effective_base) - os.environ.setdefault(env_name, resolved) - - def _strip_gateway_prefix(self, model: str) -> str: - """Drop this gateway's own prefix, leaving the upstream vendor's id.""" - if not self._gateway: - return model - prefix = f"{self._gateway.model_prefix}/" - return model[len(prefix) :] if model.startswith(prefix) else model - - def _resolve_model(self, model: str) -> str: - """The id this request is sent under. See ``providers.wire``.""" - return wire_model(model, gateway=self._gateway) - - def wire_model_id(self, model: str) -> str: - """See ``LLMProvider.wire_model_id``.""" - return self._resolve_model(model) - - def can_serve(self, model: str) -> bool: - """See ``LLMProvider.can_serve``. - - A gateway instance answers for any model -- it is the one deciding - which upstream vendor actually serves it, and its credentials are the - gateway's own, not tied to one vendor. - - For a direct instance, this only vetoes the one case both sides are - certain about: this instance's own provider_name resolves to a known, - non-OAuth spec, the model resolves to a *different* known spec, and - the two disagree -- that is one vendor's key answering for another - vendor's model, rejected outright. Every other case is let through - rather than guessed away here: - - this instance's own identity does not resolve to a spec (empty - provider_name, "auto", or a custom passthrough name LiteLLM - recognizes natively but Raven has no ProviderSpec for, e.g. - nebius/fireworks/together) -- there is nothing to compare against; - - the resolved spec is OAuth-based (e.g. github_copilot): one OAuth - grant can serve several upstream vendors, so a spec mismatch there - says nothing about whether this instance can serve the model; - - the model resolves to no spec at all (custom endpoints, bare ids - only LiteLLM itself recognizes). - In all of those, the model is not known to be wrong for this - instance, so it fails loudly at the wire instead of being guessed - away here. - """ - if self._gateway is not None: - return True - mine = find_by_name(canonical_provider_name(self._provider_name)) - if mine is None or mine.is_oauth: - return True - theirs = find_by_model(model) - if theirs is None: - return True - return theirs.name == mine.name - - def emits_unparsed_reasoning(self) -> bool: - """See ``LLMProvider.emits_unparsed_reasoning``. - - ``self._unparsed_reasoning``, when set explicitly at construction, wins - outright: it exists for a caller that already knows the answer and for - which the spec-based guess below is wrong -- a per-model routing - endpoint is built with ``provider_name="custom"`` for its api_base / - api_key shape alone, not because the backend behind it is known to be a - self-hosted inference server, so ``custom`` there would falsely claim - every one of its responses. - - Otherwise, ``self._gateway`` already answers this for both shapes it - can hold: a real network gateway (OpenRouter, AiHubMix) fronts one of - the large hosted vendors below it, so a bare ```` in content - is just content; the generic ``custom`` endpoint and a local spec - (hosted_vllm, ollama_chat) *are* the self-hosted inference server this - normalization exists for. When nothing was auto-detected, fall back to - whatever spec ``provider_name`` resolves to. - - An identity that resolves to nothing answers False, the same reading - ``can_serve`` settled on: an unresolved name says nothing about the - backend, and several production constructors (the proactive planner, - the evolver) build direct big-vendor connections with no - ``provider_name`` at all -- guessing "self-hosted" there re-opens the - false-positive cut on ordinary content this gate exists to close. A - genuinely self-hosted backend is reached through ``custom`` or a - local spec, which is where the parser-less sglang/vLLM shape comes - from; a resolved direct big vendor (anthropic, openai, ...) never - produces it behind its own API. - """ - if self._unparsed_reasoning is not None: - return self._unparsed_reasoning - spec = self._gateway or find_by_name(canonical_provider_name(self._provider_name)) - return spec is not None and (spec.is_local or spec.name == "custom") - - def supports_prompt_caching(self, model: str) -> bool: - """See ``LLMProvider.supports_prompt_caching``.""" - return self._supports_cache_control(model) - - def _supports_cache_control(self, model: str) -> bool: - """Return True when this request may carry cache_control blocks. - - Decided by ``providers.prompt_cache``, which the token strategies ask too - -- three copies of this question disagreed, and the one here could not - have answered for the marks they place. - - The address falls back to the auto-detected gateway when no - ``provider_name`` was given: several production constructors (the - evolver's launch models, the sentinel planner) pass only an - ``api_base``, and answering from the model id alone reads - ``anthropic/claude-...`` as Anthropic's wire while the request actually - travels through whatever gateway that base names -- a wire that may - have nowhere honest to put the field. - """ - from raven.providers.prompt_cache import accepts_cache_control - - addressed = self._provider_name or (self._gateway.name if self._gateway else "") - return accepts_cache_control(model, addressed_to=addressed) - - def _apply_cache_control( - self, - messages: list[dict[str, Any]], - tools: list[dict[str, Any]] | None, - ) -> tuple[list[dict[str, Any]], list[dict[str, Any]] | None]: - """Return copies of messages and tools with cache_control injected. - - Two breakpoints land on the system message when it declares a stable - prefix (``prompt_cache.STABLE_PREFIX_KEY``), one when it does not. - - The second one is the one that survives a turn. A breakpoint keys its - cache on every block up to and including itself, and the tail of this - message -- the memory recall, the skill router's hits, the Curator's - working state -- is rebuilt from whatever the user just said. So the - end-of-message breakpoint changes key on - every new turn and re-bills the identity and bootstrap text in front of - it that never changed at all. Splitting at the boundary gives that head - a key of its own. - - The end-of-message one is kept beside it rather than moved, because it - is what holds the whole message across the iterations *within* a turn: - the assembler builds this message once per turn and the tool loop only - appends after it. - """ - new_messages = [] - for msg in messages: - if msg.get("role") == "system": - content = msg["content"] - blocks = prompt_cache.split_stable_prefix(content, int(msg.get(prompt_cache.STABLE_PREFIX_KEY) or 0)) - if blocks is not None: - blocks[0] = {**blocks[0], "cache_control": prompt_cache.cache_control()} - new_content = blocks - elif isinstance(content, str): - new_content = [{"type": "text", "text": content}] - else: - new_content = list(content) - new_content[-1] = {**new_content[-1], "cache_control": prompt_cache.cache_control()} - new_messages.append({**msg, "content": new_content}) - else: - new_messages.append(msg) - - new_tools = tools - if tools: - new_tools = list(tools) - new_tools[-1] = {**new_tools[-1], "cache_control": prompt_cache.cache_control()} - - return new_messages, new_tools - def _apply_model_overrides(self, model: str, kwargs: dict[str, Any]) -> None: """Layer per-model parameter overrides: registry defaults, config on top. @@ -478,45 +218,6 @@ def _apply_model_overrides(self, model: str, kwargs: dict[str, Any]) -> None: if matches: kwargs.update(max(matches, key=lambda item: len(item[0]))[1]) - @staticmethod - def _is_anthropic_family(original_model: str, resolved_model: str) -> bool: - """Whether this request is bound for an Anthropic model, by any spelling. - - One rule for the two decisions that must agree: whether - ``thinking_blocks`` ride on the wire (``_extra_msg_keys``) and whether - the request may end on an assistant message - (``supports_assistant_prefill``). The request Anthropic rejects is - exactly those blocks plus a trailing assistant message, so deciding the - two separately would let them drift apart. - """ - spec = find_by_model(original_model) or find_by_model(resolved_model) - return bool( - (spec and spec.name == "anthropic") - or "claude" in original_model.lower() - or resolved_model.startswith("anthropic/") - ) - - @staticmethod - def _extra_msg_keys(original_model: str, resolved_model: str) -> frozenset[str]: - """Return provider-specific extra keys to preserve in request messages.""" - if LiteLLMProvider._is_anthropic_family(original_model, resolved_model): - return _ANTHROPIC_EXTRA_KEYS - return frozenset() - - def supports_assistant_prefill(self, model: str | None = None) -> bool: - """Anthropic rejects a trailing assistant message while thinking is on. - - See ``LLMProvider.supports_assistant_prefill``. The family test is the - one ``_extra_msg_keys`` applies, on purpose. - """ - original = model or self.default_model - return not self._is_anthropic_family(original, self._resolve_model(original)) - - @staticmethod - def _is_deepseek_family(model: str) -> bool: - """Whether a Wire Model names a DeepSeek model, whichever provider serves it.""" - return model.rsplit("/", 1)[-1].lower().startswith("deepseek") - @staticmethod def _normalize_tool_call_id(tool_call_id: Any) -> Any: """Normalize tool_call_id to a provider-safe 9-char alphanumeric form.""" @@ -1260,17 +961,3 @@ def _parse_response(self, response: Any, *, sent_chars: int | None = None) -> LL thinking_blocks=thinking_blocks, call_record=record, ) - - @property - def provider_name(self) -> str: - """The config section this provider was built for, or ``""``. - - Read by callers deciding whether the model string is the operator's - own naming (a ``custom`` gateway serves whatever its endpoint calls - the model) -- see ``capabilities._model_id_is_caller_chosen``. - """ - return self._provider_name - - def get_default_model(self) -> str: - """Get the default model.""" - return self.default_model diff --git a/raven/providers/litellm_shape.py b/raven/providers/litellm_shape.py new file mode 100644 index 000000000..1a025bdd4 --- /dev/null +++ b/raven/providers/litellm_shape.py @@ -0,0 +1,354 @@ +"""What a LiteLLM-routed provider answers without calling LiteLLM. + +A provider that sends through LiteLLM answers a handful of questions before and +after it sends: which id goes on the wire, whether the request may carry cache +breakpoints, whether it may end on an assistant message, whether the backend +leaks bare think tags -- and where the cache breakpoints go when it may. Every +one of those is decided from the provider's own construction arguments and +Raven's registry, never by LiteLLM itself, so it lives here, apart from the +module that imports LiteLLM. + +Two classes stand on it: :class:`raven.providers.litellm_provider.LiteLLMProvider`, +which sends, and :class:`raven.providers.relay_provider.RelayProvider`, which asks the host +that launched this process to send. Sharing the base is what keeps a relayed +agent's answers identical to a direct one's: there is one copy of each rule. +""" + +import os +import uuid +from typing import Any + +from raven.providers import prompt_cache +from raven.providers.base import LLMProvider +from raven.providers.registry import ( + canonical_provider_name, + find_by_model, + find_by_name, + find_gateway, +) +from raven.providers.wire import wire_model + +_ANTHROPIC_EXTRA_KEYS = frozenset({"thinking_blocks"}) + +# LiteLLM defaults to X-Title="liteLLM" / HTTP-Referer="https://litellm.ai" for OpenRouter +# requests, which would credit traffic to liteLLM instead of Raven on openrouter.ai/apps. +# Explicit headers here override those defaults; user-supplied extra_headers win over these. +_OPENROUTER_ATTRIBUTION: dict[str, str] = { + "HTTP-Referer": "https://raven.evermind.ai", + "X-Title": "Raven Agent", + "X-OpenRouter-Title": "Raven Agent", + "X-OpenRouter-Categories": "cli-agent,personal-agent", +} + + +def session_affinity_headers() -> dict[str, str]: + """Headers pinning one caller to one backend replica. + + Self-hosted OpenAI-compatible backends (vLLM and friends) route by this + header, so a stable value per provider instance keeps prefix-cache hits warm. + """ + return {"x-session-affinity": uuid.uuid4().hex} + + +class LiteLLMShape(LLMProvider): + """The construction and the answers of a LiteLLM-routed provider; the sending is a subclass's. + + Supports OpenRouter, Anthropic, OpenAI, Gemini, MiniMax, and many other providers through + a unified interface. Provider-specific logic is driven by the registry + (see providers/registry.py) -- no if-elif chains needed here. + """ + + api_protocol = "chat" + + def __init__( + self, + api_key: str | None = None, + api_base: str | None = None, + default_model: str = "anthropic/claude-opus-4-5", + extra_headers: dict[str, str] | None = None, + provider_name: str | None = None, + disable_auto_cache_control: bool = False, + extra_body: dict[str, Any] | None = None, + model_overrides: dict[str, dict[str, Any]] | None = None, + tls_verify: bool = True, + *, + unparsed_reasoning: bool | None = None, + ): + addressed_spec = find_by_name(provider_name) or find_by_model(default_model) + if api_base and addressed_spec and addressed_spec.strip_api_base_trailing_slash: + api_base = api_base.rstrip("/") + super().__init__(api_key, api_base) + # None: derive from the resolved spec, as emits_unparsed_reasoning always + # did. An explicit bool overrides that derivation outright -- for a + # caller that already knows the answer and for which the spec would + # guess wrong, e.g. a per-model routing endpoint built with + # provider_name="custom" for its api_base/api_key shape alone, not + # because the backend behind it is a self-hosted inference server. + self._unparsed_reasoning = unparsed_reasoning + self.default_model = default_model + self.extra_headers = extra_headers or {} + # When a TokenStrategy (e.g. CacheOptimizer) handles cache_control + # placement upstream, turn this on so the provider doesn't also + # stamp its own breakpoints on top. + self.disable_auto_cache_control = disable_auto_cache_control + # Provider-specific request body extras forwarded verbatim to LiteLLM. + # Common use: OpenRouter routing affinity to keep prompt-cache hits warm, + # extra_body={"provider": {"order": ["Anthropic"], "allow_fallbacks": False}} + self.extra_body = extra_body or {} + # User-configured per-model parameter overrides; win over the registry's. + self.model_overrides = model_overrides or {} + # A section that opted out of certificate verification has its calls sent + # through a client built with verification off (``LiteLLMProvider``); + # True, the default, lets LiteLLM build its own verifying client. + self.tls_verify = tls_verify + + # Detect gateway / local deployment. + # provider_name (from config key) is the primary signal; + # api_key / api_base are fallback for auto-detection. + # Kept because the id alone cannot say where a request goes: a bare + # `anthropic/claude-...` sent through this client reads as Anthropic's + # wire, which is not the wire it will travel on. + self._provider_name = provider_name or "" + self._gateway = find_gateway(provider_name, api_key, api_base) + if self._gateway and self._gateway.name == "openrouter": + self.extra_headers = {**_OPENROUTER_ATTRIBUTION, **self.extra_headers} + + if api_key: + self._setup_env(api_key, api_base, default_model) + + def _setup_env(self, api_key: str, api_base: str | None, model: str) -> None: + """Set environment variables based on detected provider.""" + spec = self._gateway or find_by_model(model) + if not spec: + # No spec: the key still reaches LiteLLM as an explicit api_key + # kwarg on every call, so nothing needs to go into the environment. + return + if not spec.env_key: + # OAuth/provider-only specs (for example: openai_codex) + return + + # Gateway/local overrides existing env; standard provider doesn't + if self._gateway: + os.environ[spec.env_key] = api_key + else: + os.environ.setdefault(spec.env_key, api_key) + + # Resolve env_extras placeholders: + # {api_key} → user's API key + # {api_base} → user's api_base, falling back to spec.default_api_base + effective_base = api_base or spec.default_api_base + for env_name, env_val in spec.env_extras: + resolved = env_val.replace("{api_key}", api_key) + resolved = resolved.replace("{api_base}", effective_base) + os.environ.setdefault(env_name, resolved) + + def _strip_gateway_prefix(self, model: str) -> str: + """Drop this gateway's own prefix, leaving the upstream vendor's id.""" + if not self._gateway: + return model + prefix = f"{self._gateway.model_prefix}/" + return model[len(prefix) :] if model.startswith(prefix) else model + + def _resolve_model(self, model: str) -> str: + """The id this request is sent under. See ``providers.wire``.""" + return wire_model(model, gateway=self._gateway) + + def wire_model_id(self, model: str) -> str: + """See ``LLMProvider.wire_model_id``.""" + return self._resolve_model(model) + + def can_serve(self, model: str) -> bool: + """See ``LLMProvider.can_serve``. + + A gateway instance answers for any model -- it is the one deciding + which upstream vendor actually serves it, and its credentials are the + gateway's own, not tied to one vendor. + + For a direct instance, this only vetoes the one case both sides are + certain about: this instance's own provider_name resolves to a known, + non-OAuth spec, the model resolves to a *different* known spec, and + the two disagree -- that is one vendor's key answering for another + vendor's model, rejected outright. Every other case is let through + rather than guessed away here: + - this instance's own identity does not resolve to a spec (empty + provider_name, "auto", or a custom passthrough name LiteLLM + recognizes natively but Raven has no ProviderSpec for, e.g. + nebius/fireworks/together) -- there is nothing to compare against; + - the resolved spec is OAuth-based (e.g. github_copilot): one OAuth + grant can serve several upstream vendors, so a spec mismatch there + says nothing about whether this instance can serve the model; + - the model resolves to no spec at all (custom endpoints, bare ids + only LiteLLM itself recognizes). + In all of those, the model is not known to be wrong for this + instance, so it fails loudly at the wire instead of being guessed + away here. + """ + if self._gateway is not None: + return True + mine = find_by_name(canonical_provider_name(self._provider_name)) + if mine is None or mine.is_oauth: + return True + theirs = find_by_model(model) + if theirs is None: + return True + return theirs.name == mine.name + + def emits_unparsed_reasoning(self) -> bool: + """See ``LLMProvider.emits_unparsed_reasoning``. + + ``self._unparsed_reasoning``, when set explicitly at construction, wins + outright: it exists for a caller that already knows the answer and for + which the spec-based guess below is wrong -- a per-model routing + endpoint is built with ``provider_name="custom"`` for its api_base / + api_key shape alone, not because the backend behind it is known to be a + self-hosted inference server, so ``custom`` there would falsely claim + every one of its responses. + + Otherwise, ``self._gateway`` already answers this for both shapes it + can hold: a real network gateway (OpenRouter, AiHubMix) fronts one of + the large hosted vendors below it, so a bare ```` in content + is just content; the generic ``custom`` endpoint and a local spec + (hosted_vllm, ollama_chat) *are* the self-hosted inference server this + normalization exists for. When nothing was auto-detected, fall back to + whatever spec ``provider_name`` resolves to. + + An identity that resolves to nothing answers False, the same reading + ``can_serve`` settled on: an unresolved name says nothing about the + backend, and several production constructors (the proactive planner, + the evolver) build direct big-vendor connections with no + ``provider_name`` at all -- guessing "self-hosted" there re-opens the + false-positive cut on ordinary content this gate exists to close. A + genuinely self-hosted backend is reached through ``custom`` or a + local spec, which is where the parser-less sglang/vLLM shape comes + from; a resolved direct big vendor (anthropic, openai, ...) never + produces it behind its own API. + """ + if self._unparsed_reasoning is not None: + return self._unparsed_reasoning + spec = self._gateway or find_by_name(canonical_provider_name(self._provider_name)) + return spec is not None and (spec.is_local or spec.name == "custom") + + def supports_prompt_caching(self, model: str) -> bool: + """See ``LLMProvider.supports_prompt_caching``.""" + return self._supports_cache_control(model) + + def _supports_cache_control(self, model: str) -> bool: + """Return True when this request may carry cache_control blocks. + + Decided by ``providers.prompt_cache``, which the token strategies ask too + -- three copies of this question disagreed, and the one here could not + have answered for the marks they place. + + The address falls back to the auto-detected gateway when no + ``provider_name`` was given: several production constructors (the + evolver's launch models, the sentinel planner) pass only an + ``api_base``, and answering from the model id alone reads + ``anthropic/claude-...`` as Anthropic's wire while the request actually + travels through whatever gateway that base names -- a wire that may + have nowhere honest to put the field. + """ + from raven.providers.prompt_cache import accepts_cache_control + + addressed = self._provider_name or (self._gateway.name if self._gateway else "") + return accepts_cache_control(model, addressed_to=addressed) + + def _apply_cache_control( + self, + messages: list[dict[str, Any]], + tools: list[dict[str, Any]] | None, + ) -> tuple[list[dict[str, Any]], list[dict[str, Any]] | None]: + """Return copies of messages and tools with cache_control injected. + + Two breakpoints land on the system message when it declares a stable + prefix (``prompt_cache.STABLE_PREFIX_KEY``), one when it does not. + + The second one is the one that survives a turn. A breakpoint keys its + cache on every block up to and including itself, and the tail of this + message -- the memory recall, the skill router's hits, the Curator's + working state -- is rebuilt from whatever the user just said. So the + end-of-message breakpoint changes key on + every new turn and re-bills the identity and bootstrap text in front of + it that never changed at all. Splitting at the boundary gives that head + a key of its own. + + The end-of-message one is kept beside it rather than moved, because it + is what holds the whole message across the iterations *within* a turn: + the assembler builds this message once per turn and the tool loop only + appends after it. + """ + new_messages = [] + for msg in messages: + if msg.get("role") == "system": + content = msg["content"] + blocks = prompt_cache.split_stable_prefix(content, int(msg.get(prompt_cache.STABLE_PREFIX_KEY) or 0)) + if blocks is not None: + blocks[0] = {**blocks[0], "cache_control": prompt_cache.cache_control()} + new_content = blocks + elif isinstance(content, str): + new_content = [{"type": "text", "text": content}] + else: + new_content = list(content) + new_content[-1] = {**new_content[-1], "cache_control": prompt_cache.cache_control()} + new_messages.append({**msg, "content": new_content}) + else: + new_messages.append(msg) + + new_tools = tools + if tools: + new_tools = list(tools) + new_tools[-1] = {**new_tools[-1], "cache_control": prompt_cache.cache_control()} + + return new_messages, new_tools + + @staticmethod + def _is_anthropic_family(original_model: str, resolved_model: str) -> bool: + """Whether this request is bound for an Anthropic model, by any spelling. + + One rule for the two decisions that must agree: whether + ``thinking_blocks`` ride on the wire (``_extra_msg_keys``) and whether + the request may end on an assistant message + (``supports_assistant_prefill``). The request Anthropic rejects is + exactly those blocks plus a trailing assistant message, so deciding the + two separately would let them drift apart. + """ + spec = find_by_model(original_model) or find_by_model(resolved_model) + return bool( + (spec and spec.name == "anthropic") + or "claude" in original_model.lower() + or resolved_model.startswith("anthropic/") + ) + + @staticmethod + def _extra_msg_keys(original_model: str, resolved_model: str) -> frozenset[str]: + """Return provider-specific extra keys to preserve in request messages.""" + if LiteLLMShape._is_anthropic_family(original_model, resolved_model): + return _ANTHROPIC_EXTRA_KEYS + return frozenset() + + def supports_assistant_prefill(self, model: str | None = None) -> bool: + """Anthropic rejects a trailing assistant message while thinking is on. + + See ``LLMProvider.supports_assistant_prefill``. The family test is the + one ``_extra_msg_keys`` applies, on purpose. + """ + original = model or self.default_model + return not self._is_anthropic_family(original, self._resolve_model(original)) + + @staticmethod + def _is_deepseek_family(model: str) -> bool: + """Whether a Wire Model names a DeepSeek model, whichever provider serves it.""" + return model.rsplit("/", 1)[-1].lower().startswith("deepseek") + + @property + def provider_name(self) -> str: + """The config section this provider was built for, or ``""``. + + Read by callers deciding whether the model string is the operator's + own naming (a ``custom`` gateway serves whatever its endpoint calls + the model) -- see ``capabilities._model_id_is_caller_chosen``. + """ + return self._provider_name + + def get_default_model(self) -> str: + """Get the default model.""" + return self.default_model diff --git a/raven/providers/per_model_provider.py b/raven/providers/per_model_provider.py index cd7eba94a..f5d85a949 100644 --- a/raven/providers/per_model_provider.py +++ b/raven/providers/per_model_provider.py @@ -14,13 +14,14 @@ from loguru import logger from raven.providers.base import ChatDelta, LLMProvider, LLMResponse -from raven.providers.litellm_provider import LiteLLMProvider, session_affinity_headers +from raven.providers.litellm_shape import LiteLLMShape, session_affinity_headers +from raven.providers.relay_provider import litellm_provider if TYPE_CHECKING: from raven.config.schema import ModelEndpoint -def _endpoint_provider(endpoint: "ModelEndpoint") -> LiteLLMProvider: +def _endpoint_provider(endpoint: "ModelEndpoint") -> LiteLLMShape: """Build a LiteLLM provider pinned to one endpoint. ``provider_name="custom"`` selects the generic OpenAI-compatible gateway @@ -39,7 +40,7 @@ def _endpoint_provider(endpoint: "ModelEndpoint") -> LiteLLMProvider: # api.openai.com), silently shipping the prompt and this endpoint's key # to a third party. Routing config must name the endpoint it routes to. raise ValueError(f"routing.models entry for {endpoint.model!r} has no api_base") - return LiteLLMProvider( + return litellm_provider( api_key=endpoint.api_key, api_base=endpoint.api_base, default_model=endpoint.model, @@ -55,7 +56,7 @@ class PerModelProvider(LLMProvider): def __init__(self, models: "Sequence[ModelEndpoint]", fallback: LLMProvider): super().__init__() self._fallback = fallback - self._by_model: dict[str, LiteLLMProvider] = {m.model: _endpoint_provider(m) for m in models if m.model} + self._by_model: dict[str, LiteLLMShape] = {m.model: _endpoint_provider(m) for m in models if m.model} self._default = next(iter(self._by_model), None) or fallback.get_default_model() # Per-model sub-providers are built here without generation settings; # inherit the fallback's (already configured from AgentDefaults) and push diff --git a/raven/providers/rates.py b/raven/providers/rates.py index 5e1136ce3..ba23ddb66 100644 --- a/raven/providers/rates.py +++ b/raven/providers/rates.py @@ -30,6 +30,8 @@ from loguru import logger from raven.providers import model_catalog_cache +from raven.providers.relay import active as relay_active +from raven.providers.relay import relayed #: One home for the window ladder's documented fallback: an unknown model gets #: this many tokens of headroom rather than a number invented at the call site. @@ -212,6 +214,7 @@ def _litellm_price_table() -> dict: return {} +@relayed() def _table_entry(model: str) -> dict | None: """The table row keyed exactly by this model id, or None. @@ -229,11 +232,17 @@ def _table_entry(model: str) -> dict | None: @lru_cache(maxsize=1) def _drivers_dir() -> pathlib.Path | None: - """Where the installed LiteLLM keeps its per-provider drivers.""" + """Where the installed LiteLLM keeps its per-provider drivers. + + Located, not imported: the package's directory is all this needs, and a + process whose model calls go through its host's relay never loads LiteLLM. + """ try: - from raven.providers.litellm_setup import import_litellm + import importlib.util - return pathlib.Path(import_litellm().__file__).parent / "llms" + spec = importlib.util.find_spec("litellm") + locations = list(spec.submodule_search_locations or []) if spec is not None else [] + return pathlib.Path(locations[0]) / "llms" if locations else None except Exception: return None @@ -325,6 +334,7 @@ def _endpoint_metadata_model(model: str) -> str: return model +@relayed() def _may_query_endpoint(model: str) -> bool: """Would ``get_model_info`` ask this model's own server rather than a table? @@ -395,6 +405,7 @@ def _candidates(model: str) -> list[str]: return metadata_candidates(model) +@relayed(shape=lambda value: tuple(value) if isinstance(value, list) else value) def _try_litellm_rates(model: str, input_tokens: int, output_tokens: int) -> tuple[float, float] | None: """Ask LiteLLM for per-token rates. Returns (prompt_rate, completion_rate) or None.""" try: @@ -894,12 +905,16 @@ def _litellm_ready() -> bool: import machinery reads to tell a live import from a finished one; a module carrying no spec is one no import is still running. """ + if relay_active(): + # The tables are the host's, already loaded there; asking costs a local round trip. + return True module = sys.modules.get("litellm") if module is None: return False return not getattr(getattr(module, "__spec__", None), "_initializing", False) +@relayed() def _try_litellm_max_output(model: str, *, allow_import: bool = True, allow_fetch: bool = True) -> int | None: """The model's own output ceiling, from the same metadata as the window. @@ -1055,6 +1070,7 @@ def resolve_max_output_tokens(model: str | None, *, window: int | None = None, a return max(1, min(ceiling, room, MAX_OUTPUT_TOKENS_PER_ITERATION)) +@relayed() def _try_litellm_context_window(model: str, *, allow_import: bool = True, allow_fetch: bool = True) -> int | None: """LiteLLM's model metadata, covering most mapped providers. diff --git a/raven/providers/relay.py b/raven/providers/relay.py new file mode 100644 index 000000000..f42eecfe0 --- /dev/null +++ b/raven/providers/relay.py @@ -0,0 +1,274 @@ +"""A Raven agent's LiteLLM calls, made by the host that launched it. + +Every Raven process that talks to a model through LiteLLM loads LiteLLM, its +OpenAI SDK and its catalogue -- measured at about 120 MiB of a product agent's +240. A host that launches several Raven agents already holds all of that once, +so it can make their calls for them: the host serves a relay on a local socket +and names it in the launch environment (``RAVEN_LLM_RELAY``), and the agent +reads it here. + +What crosses is deliberately narrow. The agent keeps everything a call is +wrapped in -- request assembly, the retry ladder, classification, usage +records, traces -- and hands over only the single attempt: +:class:`RelayProvider` stands where a :class:`LiteLLMProvider` would, answers the +same questions from the same code (:mod:`raven.providers.litellm_shape`), and +sends ``chat`` and ``chat_stream`` to a :class:`LiteLLMProvider` the host builds +from the very arguments this one was built with (:mod:`raven.providers.relay_provider`). +The catalogue questions that need LiteLLM's tables -- a context window, a price, +an output ceiling -- cross the same way, through :func:`relayed`. + +This module is the transport and the offer, and imports nothing of the provider +stack: the catalogue modules that use :func:`relayed` sit under that stack, and +an import back up from here would tie them into one cycle. + +Nothing here is required. An agent started without the variable, or one whose +host stops answering, loads LiteLLM and does it all itself, as before. +""" + +from __future__ import annotations + +import asyncio +import functools +import os +import socket +import threading +from collections import OrderedDict +from collections.abc import Callable +from typing import Any, TypeVar + +from loguru import logger + +from raven.providers import relay_wire + +RELAY_ENV = "RAVEN_LLM_RELAY" +"""Where the host's relay listens: ``unix:`` or ``tcp::``.""" +RELAY_TOKEN_ENV = "RAVEN_LLM_RELAY_TOKEN" # noqa: S105 - the variable's name, not its value +"""The secret every request carries, so only the agents the host launched are served.""" + +#: The modules whose :func:`relayed` functions a host will run for an agent. +#: The host imports a module by name from a request, so only these are honoured. +RELAYED_MODULES = ( + "raven.providers.rates", + "raven.providers.common_models", + "raven.providers.capabilities", + "raven.providers.anthropic_messages_provider", +) + +#: How long a catalogue question waits for the host. Generous because the +#: host's first answer may wait out its own LiteLLM import. +_ASK_TIMEOUT_S = 60.0 +#: How long opening the socket may take before the host is judged gone. +_CONNECT_TIMEOUT_S = 10.0 +_ANSWER_CACHE_SIZE = 4096 + +_UNREAD = object() +_target: Any = _UNREAD +_target_lock = threading.Lock() +_given_up = False +_RELAYED: dict[str, Callable[..., Any]] = {} +_ANSWERS: list[OrderedDict[Any, Any]] = [] + +F = TypeVar("F", bound=Callable[..., Any]) + + +class RelayUnavailableError(ConnectionError): + """The host's relay could not be reached; the caller does the work itself. + + ``connecting`` tells the two ways apart: a host nobody answers for is gone + for the rest of this process, while a question that timed out on a host that + did answer leaves the next question to it. + """ + + def __init__(self, message: str, *, connecting: bool = False) -> None: + super().__init__(message) + self.connecting = connecting + + +class RelayFaultError(RuntimeError): + """The host refused a request as malformed: a bug on one side, not a provider failure.""" + + +def target() -> tuple[str, str] | None: + """``(address, token)`` from the launch environment, or ``None`` when there is no relay. + + Read once and taken out of ``os.environ`` on the way: every process this + agent starts -- a shell command, a nested Raven -- inherits its environment, + and a relay is an offer to the agent the host launched, not to everything + downstream of it. + """ + global _target + with _target_lock: + if _target is _UNREAD: + address = os.environ.pop(RELAY_ENV, "").strip() + token = os.environ.pop(RELAY_TOKEN_ENV, "").strip() + _target = (address, token) if address and token else None + return _target + + +def active() -> bool: + """Whether this process sends its LiteLLM work to a host right now.""" + return not _given_up and target() is not None + + +def give_up(reason: BaseException) -> None: + """Stop asking a host that cannot be reached; this process does the work from now on.""" + global _given_up + if not _given_up: + _given_up = True + logger.warning("llm relay: the host cannot be reached ({}); loading LiteLLM in this process", reason) + + +def _reset_for_tests() -> None: + global _target, _given_up + with _target_lock: + _target = _UNREAD + _given_up = False + for answers in _ANSWERS: + answers.clear() + + +def _split(address: str) -> tuple[str, str | tuple[str, int]]: + kind, _, rest = address.partition(":") + if kind == "unix" and rest: + return "unix", rest + if kind == "tcp": + host, _, port = rest.rpartition(":") + if host and port.isdigit(): + return "tcp", (host, int(port)) + raise RelayFaultError(f"unreadable relay address {address!r}") + + +def request_frame(op: str, **fields: Any) -> dict[str, Any]: + found = target() + if found is None: + raise RelayUnavailableError("no relay was offered to this process") + return {"token": found[1], "op": op, **fields} + + +def _connect_sync(timeout: float) -> socket.socket: + found = target() + if found is None: + raise RelayUnavailableError("no relay was offered to this process") + kind, where = _split(found[0]) + try: + if kind == "unix": + sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) + sock.settimeout(_CONNECT_TIMEOUT_S) + sock.connect(where) + else: + sock = socket.create_connection(where, timeout=_CONNECT_TIMEOUT_S) + except OSError as exc: + raise RelayUnavailableError(str(exc), connecting=True) from exc + sock.settimeout(timeout) + return sock + + +async def connect() -> tuple[asyncio.StreamReader, asyncio.StreamWriter]: + found = target() + if found is None: + raise RelayUnavailableError("no relay was offered to this process") + kind, where = _split(found[0]) + limit = relay_wire.MAX_FRAME_BYTES + try: + if kind == "unix": + opening = asyncio.open_unix_connection(where, limit=limit) + else: + host, port = where # type: ignore[misc] + opening = asyncio.open_connection(host, port, limit=limit) + return await asyncio.wait_for(opening, timeout=_CONNECT_TIMEOUT_S) + except (OSError, asyncio.TimeoutError) as exc: + raise RelayUnavailableError(str(exc), connecting=True) from exc + + +def read_answer(reply: dict[str, Any] | None, *, sentinel: object | None = None) -> Any: + if reply is None: + raise relay_wire.FrameError("the host closed the relay connection without answering") + if "error" in reply: + raise relay_wire.decode_exception(reply["error"]) + if "fault" in reply: + raise RelayFaultError(str(reply["fault"])) + return relay_wire.decode(reply.get("result"), sentinel=sentinel) + + +def ask(op: str, *, timeout: float = _ASK_TIMEOUT_S, **fields: Any) -> Any: + """One blocking question to the host. Small, local and rare enough to make from the event loop.""" + sock = _connect_sync(timeout) + try: + relay_wire.send_frame(sock, request_frame(op, **fields)) + reply = relay_wire.recv_frame(sock) + except (OSError, relay_wire.FrameError) as exc: + raise RelayUnavailableError(str(exc)) from exc + finally: + sock.close() + return read_answer(reply) + + +def relayed( + *, keep: Callable[[Any], bool] = lambda value: value is not None, shape: Callable[[Any], Any] | None = None +): + """Run the decorated function on the host while this process has a relay. + + For a function whose answer comes out of LiteLLM's tables. The host runs the + same function -- the undecorated original, registered here by name -- so the + answer is the one this process would have computed with LiteLLM loaded. + + Answers ``keep`` accepts are remembered per arguments: a table row does not + change while a process runs, and these are asked hundreds of times a turn. An + unknown (``None`` by default) is asked again, since the host may learn it. + ``shape`` restores what JSON flattens (a tuple travels as a list). A host + that cannot be reached leaves the call to this process. + """ + + def wrap(fn: F) -> F: + name = f"{fn.__module__}:{fn.__qualname__}" + _RELAYED[name] = fn + answers: OrderedDict[Any, Any] = OrderedDict() + _ANSWERS.append(answers) + lock = threading.Lock() + + @functools.wraps(fn) + def call(*args: Any, **kwargs: Any) -> Any: + if not active(): + return fn(*args, **kwargs) + try: + key: Any = (args, tuple(sorted(kwargs.items()))) + hash(key) + except TypeError: + key = None + if key is not None: + with lock: + if key in answers: + answers.move_to_end(key) + return answers[key] + try: + value = ask("call", fn=name, args=relay_wire.encode(list(args)), kwargs=relay_wire.encode(kwargs)) + except RelayUnavailableError as exc: + if exc.connecting: + give_up(exc) + return fn(*args, **kwargs) + if shape is not None: + value = shape(value) + if key is not None and keep(value): + with lock: + answers[key] = value + while len(answers) > _ANSWER_CACHE_SIZE: + answers.popitem(last=False) + return value + + return call # type: ignore[return-value] + + return wrap + + +def run_relayed(name: str, args: list[Any], kwargs: dict[str, Any]) -> Any: + """The host's half of :func:`relayed`: run a registered function locally, by name.""" + import importlib + + module, _, _qualname = name.partition(":") + if module not in RELAYED_MODULES: + raise RelayFaultError(f"{module!r} is not a relayed module") + importlib.import_module(module) + fn = _RELAYED.get(name) + if fn is None: + raise RelayFaultError(f"{name!r} is not a relayed function") + return fn(*args, **kwargs) diff --git a/raven/providers/relay_provider.py b/raven/providers/relay_provider.py new file mode 100644 index 000000000..1c17d1b4f --- /dev/null +++ b/raven/providers/relay_provider.py @@ -0,0 +1,186 @@ +"""The provider half of the LLM relay: a LiteLLM provider whose model calls the host makes. + +See :mod:`raven.providers.relay` for the offer and the transport; this module is +what the provider factory builds in a process that has one. +""" + +from __future__ import annotations + +from collections.abc import AsyncIterator +from typing import Any + +from raven.contracts.llm_provider import ChatDelta, GenerationSettings, LLMResponse +from raven.providers import relay_wire +from raven.providers.base import LLMProvider +from raven.providers.litellm_shape import LiteLLMShape +from raven.providers.relay import ( + RelayFaultError, + RelayUnavailableError, + active, + connect, + give_up, + read_answer, + request_frame, +) + + +class RelayProvider(LiteLLMShape): + """A LiteLLM provider whose model calls the host makes. + + Constructed with exactly the arguments a :class:`LiteLLMProvider` takes, and + answering every question from the same shared code; what it sends to the + host is that construction, as it stands at the moment of the call (callers + adjust ``generation``, ``model_overrides`` and the cache flag after building + a provider), plus the call's own arguments. + """ + + def __init__(self, *args: Any, **kwargs: Any) -> None: + super().__init__(*args, **kwargs) + self._local: LLMProvider | None = None + + # Answered by the provider built here once the host could not be reached: + # it is the one sending then. Built from this construction, it answers the + # same, but a wrapper asks the provider that has the wire. + def wire_model_id(self, model: str) -> str: + return self._local.wire_model_id(model) if self._local is not None else super().wire_model_id(model) + + def can_serve(self, model: str) -> bool: + return self._local.can_serve(model) if self._local is not None else super().can_serve(model) + + def emits_unparsed_reasoning(self) -> bool: + return self._local.emits_unparsed_reasoning() if self._local is not None else super().emits_unparsed_reasoning() + + def spec(self) -> dict[str, Any]: + """What the host builds its provider from: this one's construction, as it stands now.""" + return { + "api_key": self.api_key, + "api_base": self.api_base, + "default_model": self.default_model, + "extra_headers": dict(self.extra_headers or {}), + "provider_name": self._provider_name or None, + "disable_auto_cache_control": bool(self.disable_auto_cache_control), + "extra_body": dict(self.extra_body or {}) or None, + "model_overrides": dict(self.model_overrides or {}) or None, + "unparsed_reasoning": self._unparsed_reasoning, + "generation": relay_wire.encode(getattr(self, "generation", None) or GenerationSettings()), + } + + def _locally(self) -> LLMProvider: + """The provider this process builds for itself once the host cannot be reached. + + Kept in step with this one on every call, for the same reason the spec + is read at the call: callers adjust these after building a provider. + """ + if self._local is None: + from raven.providers.litellm_provider import LiteLLMProvider + + spec = self.spec() + spec.pop("generation") + self._local = LiteLLMProvider(**spec) + local = self._local + local.generation = getattr(self, "generation", None) or GenerationSettings() + for attr in ("model_overrides", "disable_auto_cache_control", "extra_body", "extra_headers"): + setattr(local, attr, getattr(self, attr)) + return local + + def _call(self, **args: Any) -> dict[str, Any]: + return {"spec": self.spec(), "args": relay_wire.encode(args, sentinel=self._SENTINEL)} + + async def chat( + self, + messages: list[dict[str, Any]], + tools: list[dict[str, Any]] | None = None, + model: str | None = None, + max_tokens: int | None = None, + temperature: float = 0.7, + reasoning_effort: str | None = None, + tool_choice: str | dict[str, Any] | None = None, + ) -> LLMResponse: + args = { + "messages": messages, + "tools": tools, + "model": model, + "max_tokens": max_tokens, + "temperature": temperature, + "reasoning_effort": reasoning_effort, + "tool_choice": tool_choice, + } + if active(): + try: + reader, writer = await connect() + except RelayUnavailableError as exc: + give_up(exc) + else: + try: + await relay_wire.write_frame(writer, request_frame("chat", **self._call(**args))) + return read_answer(await relay_wire.read_frame(reader)) + finally: + writer.close() + return await self._locally().chat(**args) + + async def chat_stream( + self, + messages: list[dict[str, Any]], + tools: list[dict[str, Any]] | None = None, + model: str | None = None, + max_tokens: object = LLMProvider._SENTINEL, + temperature: object = LLMProvider._SENTINEL, + reasoning_effort: object = LLMProvider._SENTINEL, + tool_choice: str | dict[str, Any] | None = None, + ) -> AsyncIterator[ChatDelta]: + args = { + "messages": messages, + "tools": tools, + "model": model, + "max_tokens": max_tokens, + "temperature": temperature, + "reasoning_effort": reasoning_effort, + "tool_choice": tool_choice, + } + if active(): + try: + reader, writer = await connect() + except RelayUnavailableError as exc: + give_up(exc) + else: + # Closing the connection is how a stream abandoned here -- a + # cancelled turn, a consumer that stopped reading -- is called off + # on the host: its reader sees the end and cancels the call. + try: + await relay_wire.write_frame(writer, request_frame("stream", **self._call(**args))) + while True: + frame = await relay_wire.read_frame(reader) + if frame is not None and "delta" in frame: + yield relay_wire.decode(frame["delta"]) + continue + if frame is not None and frame.get("done"): + return + read_answer(frame) + raise RelayFaultError("the host ended a stream with an unknown frame") + finally: + writer.close() + async for delta in self._locally().chat_stream(**args): + yield delta + + +def litellm_provider(*args: Any, **kwargs: Any) -> LiteLLMShape: + """A LiteLLM-routed provider: relayed to the host when this process has a relay, else direct. + + Takes :class:`LiteLLMProvider`'s arguments. The direct class is imported only + when it is the one returned, which is what keeps a relayed process free of + LiteLLM. + """ + if active(): + return RelayProvider(*args, **kwargs) + from raven.providers.litellm_provider import LiteLLMProvider + + return LiteLLMProvider(*args, **kwargs) + + +def is_litellm_transport(provider: Any) -> bool: + """Whether ``provider`` sends through LiteLLM, here or on the host, without importing LiteLLM. + + Both classes stand on :class:`LiteLLMShape`, so the question is answered + from the module that does not import it. + """ + return isinstance(provider, LiteLLMShape) diff --git a/raven/providers/relay_server.py b/raven/providers/relay_server.py new file mode 100644 index 000000000..014d67ce2 --- /dev/null +++ b/raven/providers/relay_server.py @@ -0,0 +1,288 @@ +"""The host's half of the LLM relay: model calls made here for the Raven agents it launched. + +One server per host process, started the first time a Raven agent is launched +with the relay on, listening where only this user can reach it -- a socket in a +fresh ``0700`` directory on POSIX, loopback TCP on Windows -- and answering +only requests carrying the token it handed out in the launch environment. + +Each request is one connection: open, one frame in, the answer out, close. A +stream answers with a frame per delta and a closing frame, and a connection the +agent closes early is how a cancelled turn reaches the call: the host stops it. + +Calls are made by a :class:`LiteLLMProvider` built from the agent's own +construction arguments, kept per distinct construction. They are made inside +``usage_record.recorded_by_caller``: the agent records every call it makes, so +the host recording it too would bill it twice. + +What the agent asks besides calls -- a context window, a price -- runs a +registered :func:`raven.providers.relay.relayed` function in a worker thread, +because the first answer may import LiteLLM and the host's loop has its own +work to do meanwhile. +""" + +from __future__ import annotations + +import asyncio +import atexit +import contextlib +import hashlib +import hmac +import json +import secrets +import shutil +import sys +import tempfile +from collections import OrderedDict +from collections.abc import Awaitable, Callable +from pathlib import Path +from typing import Any + +from loguru import logger + +from raven.providers import relay_wire, usage_record +from raven.providers.base import LLMProvider +from raven.providers.relay import RELAY_ENV, RELAY_TOKEN_ENV, RelayFaultError, run_relayed + +#: Distinct constructions kept built. An agent normally has one or two (its +#: model, a fallback); the bound is for a host that serves many agents. +_PROVIDER_CACHE_SIZE = 32 + +ProviderBuilder = Callable[[dict[str, Any]], LLMProvider] + + +def build_litellm_provider(spec: dict[str, Any]) -> LLMProvider: + """The provider an agent would have built for itself from ``spec``.""" + from raven.providers.litellm_provider import LiteLLMProvider + + spec = dict(spec) + generation = relay_wire.decode(spec.pop("generation", None)) + provider = LiteLLMProvider(**spec) + if generation is not None: + provider.generation = generation + return provider + + +class RelayServer: + """Serves the relay for this host's agents. Start once; hand :meth:`env` to each launch.""" + + def __init__(self, *, build: ProviderBuilder = build_litellm_provider, tcp: bool | None = None) -> None: + self._build = build + # Loopback TCP where there is no Unix socket to put in a private directory. + self._tcp = sys.platform == "win32" if tcp is None else tcp + self._token = secrets.token_urlsafe(32) + self._server: asyncio.AbstractServer | None = None + self._address = "" + self._dir: Path | None = None + self._providers: OrderedDict[str, LLMProvider] = OrderedDict() + self._building: dict[str, asyncio.Future[LLMProvider]] = {} + self.loop: asyncio.AbstractEventLoop | None = None + """The loop that answers this server's connections; a server is no use to another.""" + + @property + def serving(self) -> bool: + return self._server is not None and self._server.is_serving() + + def env(self) -> dict[str, str]: + """The launch environment that offers this relay to an agent.""" + return {RELAY_ENV: self._address, RELAY_TOKEN_ENV: self._token} + + async def start(self) -> None: + self.loop = asyncio.get_running_loop() + if self._tcp: + self._server = await asyncio.start_server(self._serve, host="127.0.0.1", port=0) + host, port = self._server.sockets[0].getsockname()[:2] + self._address = f"tcp:{host}:{port}" + else: + # mkdtemp makes the directory 0700, which is what keeps another user + # from reaching the socket at all; the token is the second lock. + self._dir = Path(tempfile.mkdtemp(prefix="raven-llm-relay-")) + path = self._dir / "relay.sock" + self._server = await asyncio.start_unix_server(self._serve, path=str(path)) + self._address = f"unix:{path}" + atexit.register(shutil.rmtree, self._dir, True) + logger.debug("llm relay: serving at {}", self._address) + + def abandon(self) -> None: + """Stop listening without the loop it was started on, which may be closed by now. + + A connection made to a server whose loop no longer runs is accepted by + the system and then answered by nobody, so its socket has to go. + """ + server, self._server = self._server, None + if server is not None: + # Server.close is synchronous and closes the listening sockets + # itself; on a closed loop it skips only the reader removal. + with contextlib.suppress(Exception): + server.close() + if self._dir is not None: + shutil.rmtree(self._dir, ignore_errors=True) + self._dir = None + self._providers.clear() + + async def close(self) -> None: + if self._server is not None: + self._server.close() + with contextlib.suppress(Exception): + await self._server.wait_closed() + self._server = None + if self._dir is not None: + shutil.rmtree(self._dir, ignore_errors=True) + self._dir = None + self._providers.clear() + + async def _provider(self, spec: Any) -> LLMProvider: + if not isinstance(spec, dict): + raise RelayFaultError("a call carries no provider construction") + key = hashlib.sha256(json.dumps(spec, sort_keys=True, default=str).encode("utf-8")).hexdigest() + found = self._providers.get(key) + if found is not None: + self._providers.move_to_end(key) + return found + pending = self._building.get(key) + if pending is None: + # Built off the loop: the first build of a process imports LiteLLM. + pending = asyncio.ensure_future(asyncio.to_thread(self._build, spec)) + self._building[key] = pending + try: + provider = await asyncio.shield(pending) + finally: + self._building.pop(key, None) + self._providers[key] = provider + while len(self._providers) > _PROVIDER_CACHE_SIZE: + self._providers.popitem(last=False) + return provider + + async def _serve(self, reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None: + try: + request = await relay_wire.read_frame(reader) + if request is None: + return + if not hmac.compare_digest(str(request.get("token") or ""), self._token): + await relay_wire.write_frame(writer, {"fault": "this relay does not serve that token"}) + return + op = request.get("op") + if op == "stream": + await self._stream(request, reader, writer) + elif op == "chat": + await self._chat(request, reader, writer) + elif op == "call": + await self._run(request, writer) + else: + await relay_wire.write_frame(writer, {"fault": f"unknown relay op {op!r}"}) + except (ConnectionError, relay_wire.FrameError, asyncio.IncompleteReadError): + pass + except RelayFaultError as exc: + with contextlib.suppress(Exception): + await relay_wire.write_frame(writer, {"fault": str(exc)}) + except Exception as exc: # noqa: BLE001 - one agent's bad request must not end the server + logger.warning("llm relay: request failed: {}", exc) + with contextlib.suppress(Exception): + await relay_wire.write_frame(writer, {"fault": f"{type(exc).__name__}: {exc}"}) + finally: + writer.close() + + async def _watched(self, work: Awaitable[None], reader: asyncio.StreamReader) -> BaseException | None: + """Run ``work`` until it ends or the agent hangs up; what it raised, if anything. + + The agent sends nothing after its request, so any read returning is the + connection ending -- which is the agent calling the work off. + """ + task = asyncio.ensure_future(work) + hangup = asyncio.ensure_future(reader.read(1)) + try: + await asyncio.wait({task, hangup}, return_when=asyncio.FIRST_COMPLETED) + finally: + if not task.done(): + task.cancel() + with contextlib.suppress(BaseException): + await task + hangup.cancel() + with contextlib.suppress(BaseException): + await hangup + if task.cancelled(): + return None + return task.exception() + + async def _failed(self, writer: asyncio.StreamWriter, provider: LLMProvider, exc: BaseException) -> None: + status = provider._extract_status_code(exc) + await relay_wire.write_frame(writer, {"error": relay_wire.encode_exception(exc, status=status)}) + + async def _stream( + self, request: dict[str, Any], reader: asyncio.StreamReader, writer: asyncio.StreamWriter + ) -> None: + provider = await self._provider(request.get("spec")) + args = relay_wire.decode(request.get("args") or {}, sentinel=provider._SENTINEL) + + async def pump() -> None: + with usage_record.recorded_by_caller(): + async for delta in provider.chat_stream(**args): + await relay_wire.write_frame(writer, {"delta": relay_wire.encode(delta)}) + await relay_wire.write_frame(writer, {"done": True}) + + failure = await self._watched(pump(), reader) + if failure is not None and not isinstance(failure, (ConnectionError, relay_wire.FrameError)): + await self._failed(writer, provider, failure) + + async def _chat(self, request: dict[str, Any], reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None: + provider = await self._provider(request.get("spec")) + args = relay_wire.decode(request.get("args") or {}, sentinel=provider._SENTINEL) + + async def answer() -> None: + with usage_record.recorded_by_caller(): + response = await provider.chat(**args) + await relay_wire.write_frame(writer, {"result": relay_wire.encode(response)}) + + failure = await self._watched(answer(), reader) + if failure is not None and not isinstance(failure, (ConnectionError, relay_wire.FrameError)): + await self._failed(writer, provider, failure) + + async def _run(self, request: dict[str, Any], writer: asyncio.StreamWriter) -> None: + name = request.get("fn") + args = relay_wire.decode(request.get("args") or []) + kwargs = relay_wire.decode(request.get("kwargs") or {}) + if not isinstance(name, str) or not isinstance(args, list) or not isinstance(kwargs, dict): + raise RelayFaultError("a relayed call names no function") + result = await asyncio.to_thread(run_relayed, name, args, kwargs) + await relay_wire.write_frame(writer, {"result": relay_wire.encode(result)}) + + +_SERVER: RelayServer | None = None +_STARTING: asyncio.Lock | None = None + + +async def relay_env() -> dict[str, str]: + """The environment offering this host's relay to an agent, starting the relay on first use. + + Also starts LiteLLM's import in the background, once: the agents will ask + for it, and a first answer that waits out the import waits on this host. + """ + global _SERVER, _STARTING + loop = asyncio.get_running_loop() + if _SERVER is not None and _SERVER.loop is not loop: + # Started under a loop that has ended (a CLI turn is one asyncio.run). + _SERVER.abandon() + _SERVER, _STARTING = None, None + if _SERVER is not None and _SERVER.serving: + return _SERVER.env() + if _STARTING is None: + _STARTING = asyncio.Lock() + async with _STARTING: + if _SERVER is None or not _SERVER.serving: + server = RelayServer() + await server.start() + _SERVER = server + if "litellm" not in sys.modules: + from raven.providers.litellm_setup import warm_up_in_background + + warm_up_in_background() + return _SERVER.env() + + +async def close_relay() -> None: + """Stop serving. For process shutdown and for tests.""" + global _SERVER, _STARTING + server, _SERVER, _STARTING = _SERVER, None, None + if server is not None and server.loop is not asyncio.get_running_loop(): + server.abandon() + elif server is not None: + await server.close() diff --git a/raven/providers/relay_wire.py b/raven/providers/relay_wire.py new file mode 100644 index 000000000..76c6bbeb3 --- /dev/null +++ b/raven/providers/relay_wire.py @@ -0,0 +1,302 @@ +"""The frames a relayed model call travels in, and how its values survive the trip. + +A Raven agent whose LiteLLM calls are made by the host that launched it +(:mod:`raven.providers.relay`) talks to that host over a local socket. Every +message is one frame: a four-byte big-endian length, then that many bytes of +UTF-8 JSON. Length-prefixed rather than line-delimited because a request is a +whole conversation, images included, and a line reader would need a limit large +enough for the largest one. + +What crosses is the provider vocabulary of :mod:`raven.contracts.llm_provider` +-- deltas, responses, tool calls, classifications, call records -- encoded by +name from a fixed list, never by pickling: the socket is local, but neither +side should have to trust the other's bytes to build arbitrary objects. + +A failure crosses as the facts the agent's own retry and classification read +off an exception: its message, the class names along its type and cause chain, +the HTTP status, the response body and kept headers. The agent rebuilds an +exception that answers those reads the way the original did, so a rate limit +on the host is a rate limit in the agent's retry ladder, and a first-byte or +idle timeout is the same class it would have been without the relay. +""" + +from __future__ import annotations + +import asyncio +import dataclasses +import json +import socket +import struct +from types import SimpleNamespace +from typing import Any + +from raven.contracts.llm_provider import ( + CallRecord, + ChatDelta, + ErrorClassification, + GenerationSettings, + LLMResponse, + RunMeta, + ToolCallRequest, + TruncationInfo, +) + +#: Bounds one frame. A request carries the whole conversation including any +#: images, so this is far above any real call and exists only so a corrupt +#: length cannot make a reader allocate without limit. +MAX_FRAME_BYTES = 512 * 1024 * 1024 + +_HEADER = struct.Struct(">I") + +_DATACLASSES: dict[str, type] = { + cls.__name__: cls + for cls in ( + CallRecord, + ChatDelta, + ErrorClassification, + GenerationSettings, + LLMResponse, + RunMeta, + ToolCallRequest, + TruncationInfo, + ) +} +_TAG = "__relay__" +_SENTINEL_TAG = "sentinel" + + +class FrameError(ConnectionError): + """The other side sent something that is not a frame, or went away mid-frame.""" + + +def _pack(message: dict[str, Any]) -> bytes: + body = json.dumps(message, ensure_ascii=False, separators=(",", ":")).encode("utf-8") + if len(body) > MAX_FRAME_BYTES: + raise FrameError(f"relay frame of {len(body)} bytes exceeds {MAX_FRAME_BYTES}") + return _HEADER.pack(len(body)) + body + + +def _unpack(body: bytes) -> dict[str, Any]: + try: + message = json.loads(body.decode("utf-8")) + except (UnicodeDecodeError, ValueError) as exc: + raise FrameError(f"relay frame is not JSON: {exc}") from exc + if not isinstance(message, dict): + raise FrameError("relay frame is not an object") + return message + + +def _length(header: bytes) -> int: + (size,) = _HEADER.unpack(header) + if size > MAX_FRAME_BYTES: + raise FrameError(f"relay frame announces {size} bytes, over {MAX_FRAME_BYTES}") + return size + + +async def write_frame(writer: asyncio.StreamWriter, message: dict[str, Any]) -> None: + writer.write(_pack(message)) + await writer.drain() + + +async def read_frame(reader: asyncio.StreamReader) -> dict[str, Any] | None: + """The next frame, or ``None`` when the other side closed between frames.""" + try: + header = await reader.readexactly(_HEADER.size) + except asyncio.IncompleteReadError as exc: + if exc.partial: + raise FrameError("relay connection closed inside a frame header") from exc + return None + try: + body = await reader.readexactly(_length(header)) + except asyncio.IncompleteReadError as exc: + raise FrameError("relay connection closed inside a frame") from exc + return _unpack(body) + + +def send_frame(sock: socket.socket, message: dict[str, Any]) -> None: + sock.sendall(_pack(message)) + + +def _recv_exactly(sock: socket.socket, size: int) -> bytes: + chunks: list[bytes] = [] + remaining = size + while remaining: + chunk = sock.recv(min(remaining, 1 << 20)) + if not chunk: + raise FrameError("relay connection closed inside a frame") + chunks.append(chunk) + remaining -= len(chunk) + return b"".join(chunks) + + +def recv_frame(sock: socket.socket) -> dict[str, Any]: + return _unpack(_recv_exactly(sock, _length(_recv_exactly(sock, _HEADER.size)))) + + +# ---- values ---------------------------------------------------------------- + + +def encode(value: Any, *, sentinel: object | None = None) -> Any: + """``value`` as plain JSON, with the provider vocabulary tagged by class name. + + ``sentinel`` is the provider's "not given" marker, which travels as a tag so + the far side substitutes its own: the two processes hold different objects. + Anything outside the vocabulary that a client library leaves in a payload -- + a pydantic model in a usage block -- is reduced to its plain form. + """ + if sentinel is not None and value is sentinel: + return {_TAG: _SENTINEL_TAG} + if value is None or isinstance(value, (bool, int, float, str)): + return value + if dataclasses.is_dataclass(value) and not isinstance(value, type) and type(value).__name__ in _DATACLASSES: + fields = {f.name: encode(getattr(value, f.name), sentinel=sentinel) for f in dataclasses.fields(value)} + return {_TAG: type(value).__name__, "fields": fields} + if isinstance(value, dict): + return {str(key): encode(item, sentinel=sentinel) for key, item in value.items()} + if isinstance(value, (list, tuple, set, frozenset)): + return [encode(item, sentinel=sentinel) for item in value] + dump = getattr(value, "model_dump", None) + if callable(dump): + try: + return encode(dump(), sentinel=sentinel) + except Exception: # noqa: BLE001 - fall through to the generic forms below + pass + if dataclasses.is_dataclass(value) and not isinstance(value, type): + return encode(dataclasses.asdict(value), sentinel=sentinel) + return str(value) + + +def decode(value: Any, *, sentinel: object | None = None) -> Any: + """The inverse of :func:`encode`; ``sentinel`` is this side's "not given" marker.""" + if isinstance(value, list): + return [decode(item, sentinel=sentinel) for item in value] + if not isinstance(value, dict): + return value + tag = value.get(_TAG) + if tag == _SENTINEL_TAG: + return sentinel + if isinstance(tag, str) and tag in _DATACLASSES and isinstance(value.get("fields"), dict): + cls = _DATACLASSES[tag] + known = {f.name for f in dataclasses.fields(cls)} + fields = {k: decode(v, sentinel=sentinel) for k, v in value["fields"].items() if k in known} + if cls is CallRecord and isinstance(fields.get("headers"), dict): + fields["headers"] = {str(k): str(v) for k, v in fields["headers"].items()} + return cls(**fields) + return {key: decode(item, sentinel=sentinel) for key, item in value.items()} + + +# ---- failures -------------------------------------------------------------- + +_BUILTIN_BASES = {"object", "BaseException", "Exception", "TimeoutError", "OSError"} +_CARRIERS: dict[tuple[tuple[str, ...], bool], type[BaseException]] = {} + + +class RelayedProviderError(Exception): + """A provider failure raised on the host, standing in for the original in the agent. + + Its class is built per failure so that its MRO carries the original's class + names (see :func:`decode_exception`); this is the root every one of them shares. + """ + + +def _chain(exc: BaseException) -> list[BaseException]: + seen: set[int] = set() + found: list[BaseException] = [] + cur: BaseException | None = exc + while cur is not None and id(cur) not in seen: + seen.add(id(cur)) + found.append(cur) + cur = cur.__cause__ or cur.__context__ + return found + + +def encode_exception(exc: BaseException, *, status: int | None) -> dict[str, Any]: + """The facts the agent's classification and call record read off ``exc``.""" + from raven.providers import call_record + from raven.providers.first_byte import FirstByteTimeoutError, StreamIdleTimeoutError + + names: list[str] = [] + for link in _chain(exc): + for klass in type(link).__mro__: + if klass.__name__ not in names: + names.append(klass.__name__) + response = getattr(exc, "response", None) + data: dict[str, Any] = { + "message": str(exc), + "names": names, + "status": status, + "body": call_record.exception_body(exc), + "headers": call_record.keep_headers(getattr(response, "headers", None) if response is not None else None), + "attrs": { + attr: value for attr in ("text", "message") if isinstance(value := getattr(exc, attr, None), str) and value + }, + } + if isinstance(exc, FirstByteTimeoutError): + data["kind"] = "first_byte" + data["timeout"] = {"phase": exc.phase, "budget": exc.budget, "waited": exc.waited} + elif isinstance(exc, StreamIdleTimeoutError): + data["kind"] = "stream_idle" + data["timeout"] = {"idle": getattr(exc, "idle", None)} + elif isinstance(exc, TimeoutError): + data["kind"] = "timeout" + else: + data["kind"] = "error" + return data + + +def _carrier(names: tuple[str, ...], timeout: bool) -> type[BaseException]: + """A class whose MRO spells ``names``, rooted in :class:`RelayedProviderError`. + + Built from the most general name down, so the most specific -- the original + class's own -- is the class itself and reads as its name. + """ + key = (names, timeout) + cached = _CARRIERS.get(key) + if cached is not None: + return cached + roots: tuple[type, ...] = (RelayedProviderError, TimeoutError) if timeout else (RelayedProviderError,) + root_names = {klass.__name__ for root in roots for klass in root.__mro__} + cls: type = type("RelayedTimeoutError", roots, {"__module__": __name__}) if timeout else RelayedProviderError + for name in reversed([n for n in names if n not in _BUILTIN_BASES and n not in root_names]): + cls = type(name, (cls,), {"__module__": __name__}) + _CARRIERS[key] = cls + return cls + + +def decode_exception(data: dict[str, Any]) -> BaseException: + """An exception that answers the agent's reads the way the host's original did.""" + from raven.providers.first_byte import FirstByteTimeoutError, StreamIdleTimeoutError + + kind = data.get("kind") + message = str(data.get("message") or "") + timeout = data.get("timeout") if isinstance(data.get("timeout"), dict) else {} + bound: BaseException | None = None + try: + if kind == "first_byte": + bound = FirstByteTimeoutError( + phase=str(timeout.get("phase") or ""), + budget=float(timeout.get("budget") or 0.0), + waited=float(timeout.get("waited") or 0.0), + ) + elif kind == "stream_idle": + bound = StreamIdleTimeoutError(idle=float(timeout.get("idle") or 0.0)) + except (TypeError, ValueError): + bound = None + if bound is not None: + # The host's own sentence, not one rebuilt from rounded numbers. + bound.args = (message,) + return bound + names = tuple(str(n) for n in data.get("names") or () if isinstance(n, str)) + exc = _carrier(names, kind in {"timeout", "first_byte", "stream_idle"})(message) + status = data.get("status") + if isinstance(status, int): + exc.status_code = status # type: ignore[attr-defined] + body = data.get("body") + if isinstance(body, str): + exc.body = body # type: ignore[attr-defined] + headers = data.get("headers") if isinstance(data.get("headers"), dict) else {} + exc.response = SimpleNamespace(headers=headers, text=body if isinstance(body, str) else None) # type: ignore[attr-defined] + for attr, value in (data.get("attrs") or {}).items(): + if attr in ("text", "message") and isinstance(value, str): + setattr(exc, attr, value) + return exc diff --git a/raven/utils/tokens.py b/raven/utils/tokens.py index 9f4fb27c9..15d4a3e33 100644 --- a/raven/utils/tokens.py +++ b/raven/utils/tokens.py @@ -4,6 +4,7 @@ wants a safe filename does not pay for an encoding table. """ +import contextlib import json import re from collections import Counter @@ -32,6 +33,22 @@ _RUN = re.compile(r"[^\S\x1c-\x1f]+|[\S\x1c-\x1f]+") +def encoding() -> tiktoken.Encoding: + """The cl100k encoding the estimator counts with, without the Python copy of its ranks. + + tiktoken's registered encoding keeps that copy beside the one its core + encodes with and reads it only to pickle an encoding outside the registry -- + a registered one pickles by name -- so it is dropped once loaded, about + 20 MiB. LiteLLM's import drops the same copy (``litellm_setup``); this drops + it in a process that never imports LiteLLM. + """ + enc = tiktoken.get_encoding(_ENCODING_NAME) + if hasattr(enc, "_mergeable_ranks"): + with contextlib.suppress(AttributeError): + del enc._mergeable_ranks + return enc + + def _encoded_len(encoding: tiktoken.Encoding, text: str) -> int: """``len(encoding.encode(text))``, with any run past ``_RUN_SLICE`` counted in slices.""" if len(text) <= _RUN_SLICE: @@ -150,7 +167,7 @@ def estimate_prompt_tokens( cache.clear() return extra_tokens try: - enc = tiktoken.get_encoding(_ENCODING_NAME) + enc = encoding() text_tokens = _encoded_len(enc, payload) if cache is None else cache.count(payload, enc) except Exception: if cache is not None: @@ -195,7 +212,7 @@ def estimate_message_tokens(message: dict[str, Any]) -> int: if not payload: return max(1, extra_tokens) try: - enc = tiktoken.get_encoding(_ENCODING_NAME) + enc = encoding() text_tokens = _encoded_len(enc, payload) except Exception: text_tokens = len(payload) // 4 diff --git a/tests/test_acp_fork_template.py b/tests/test_acp_fork_template.py new file mode 100644 index 000000000..4b3e2fd63 --- /dev/null +++ b/tests/test_acp_fork_template.py @@ -0,0 +1,703 @@ +"""Raven agents forked from a template that has already loaded their code (Linux only).""" + +from __future__ import annotations + +import asyncio +import json +import os +import signal +import sys +import textwrap +import time +from pathlib import Path + +import pytest + +from raven.acp_client import fork_template +from raven.acp_client.fork_template import fork_launch, raven_agent_argv + +pytestmark = pytest.mark.skipif(not fork_template.template_available(), reason="forking agents is Linux-only") + + +@pytest.fixture(autouse=True) +async def _no_template(): + yield + await fork_template.close_template() + + +def _product(tmp_path: Path, body: str) -> Path: + folder = tmp_path / "product" + folder.mkdir(exist_ok=True) + (folder / "subagent.json").write_text('{"name": "P"}', encoding="utf-8") + script = folder / "run.py" + script.write_text(textwrap.dedent(body), encoding="utf-8") + return script + + +async def _launch(script: Path, *args: str, env: dict[str, str] | None = None, cwd: str | None = None): + base = dict(os.environ) + endpoint = await fork_launch( + [sys.executable, str(script), *args], {**base, **(env or {})}, cwd or str(script.parent), base_env=base + ) + assert endpoint is not None + return endpoint + + +async def _drain(endpoint) -> tuple[bytes, int]: + out = await asyncio.wait_for(endpoint._stdout.read(), timeout=60) + return out, await asyncio.wait_for(endpoint.wait(), timeout=60) + + +def test_only_a_raven_agent_on_this_interpreter_is_forked(tmp_path: Path) -> None: + script = _product(tmp_path, "pass\n") + loose = tmp_path / "loose.py" + loose.write_text("pass\n") + + assert raven_agent_argv([sys.executable, str(script), "--acp"]) + assert raven_agent_argv([sys.executable, "-m", "raven", "acp"]) + assert not raven_agent_argv([sys.executable, "-m", "other"]) + assert not raven_agent_argv([sys.executable, str(loose)]) + assert not raven_agent_argv(["/usr/bin/node", str(script)]) + # Another environment's interpreter runs with its own packages, however it links. + other = tmp_path / "venv" / "bin" / "python" + other.parent.mkdir(parents=True) + other.symlink_to(sys.executable) + assert not raven_agent_argv([str(other), str(script)]) + + +async def test_a_forked_agent_runs_as_a_fresh_interpreter_would_have(tmp_path: Path) -> None: + script = _product( + tmp_path, + """ + import os, sys + line = sys.stdin.readline().strip() + print(line, os.getcwd(), os.environ.get("PROBE"), sys.argv[1:], sys.path[0], os.getpgid(0) == os.getpid(), flush=True) + print("diagnostic", file=sys.stderr, flush=True) + sys.exit(7) + """, + ) + endpoint = await _launch(script, "--flag", env={"PROBE": "seen"}, cwd=str(tmp_path)) + try: + assert "forked" in endpoint.describe() + await endpoint.write(b"hello\n") + out, code = await _drain(endpoint) + assert out.decode().split() == ["hello", str(tmp_path), "seen", "['--flag']", str(script.parent), "True"] + assert await endpoint.readerr() == b"diagnostic\n" + assert code == 7 + assert endpoint.returncode == 7 + finally: + endpoint.release() + + +async def test_a_launchers_exec_of_raven_runs_raven_in_the_same_process(tmp_path: Path) -> None: + script = _product( + tmp_path, + """ + import os, sys + print(os.getpid(), flush=True) + os.execv(sys.executable, [sys.executable, "-m", "raven", "--version"]) + raise AssertionError("unreachable") + """, + ) + endpoint = await _launch(script) + try: + out, code = await _drain(endpoint) + pid, rest = out.decode().split("\n", 1) + assert int(pid) == endpoint.pid + assert "Raven v" in rest + assert code == 0 + finally: + endpoint.release() + + +async def test_ending_an_agent_kills_its_group(tmp_path: Path) -> None: + script = _product( + tmp_path, + """ + import subprocess, sys, time + child = subprocess.Popen([sys.executable, "-c", "import time; time.sleep(600)"]) + print(child.pid, flush=True) + time.sleep(600) + """, + ) + endpoint = await _launch(script) + grandchild = int((await asyncio.wait_for(endpoint.readline(), timeout=60)).decode()) + await endpoint.end() + code = await asyncio.wait_for(endpoint.wait(), timeout=30) + endpoint.release() + + assert code == -signal.SIGKILL + for _ in range(100): + try: + os.kill(grandchild, 0) + except ProcessLookupError: + break + time.sleep(0.05) + else: + pytest.fail("the agent's own child outlived the group kill") + + +async def test_an_agent_ends_with_the_template(tmp_path: Path) -> None: + script = _product(tmp_path, "import time\nprint('up', flush=True)\ntime.sleep(600)\n") + endpoint = await _launch(script) + assert await asyncio.wait_for(endpoint.readline(), timeout=60) == b"up\n" + template = fork_template._TEMPLATE + assert template is not None and template._proc is not None + + os.kill(template._proc.pid, signal.SIGKILL) + await asyncio.wait_for(endpoint.wait(), timeout=30) + endpoint.release() + + for _ in range(100): + try: + os.kill(endpoint.pid, 0) + except ProcessLookupError: + return + time.sleep(0.05) + pytest.fail("a forked agent outlived its template") + + +async def test_an_agent_writes_what_it_imported_for_the_next_template(tmp_path: Path, monkeypatch) -> None: + monkeypatch.setattr(fork_template, "LEARN_AFTER_S", 0.2) + script = _product(tmp_path, "import time, colorsys\ntime.sleep(2)\n") + endpoint = await _launch(script) + await asyncio.wait_for(endpoint.wait(), timeout=60) + endpoint.release() + + from raven.config.paths import get_runtime_subdir + + learned = json.loads((get_runtime_subdir("acp-template") / "modules.json").read_text()) + assert learned["python"] == list(sys.version_info[:2]) + assert list(learned["products"]) == ["product"], "learned under the folder the launcher sits in" + assert "colorsys" in learned["products"]["product"] + assert "__main__" not in learned["products"]["product"] + + +async def test_the_template_closes_once_it_has_had_no_agent_for_the_budget(tmp_path: Path, monkeypatch) -> None: + from raven.config import live as live_mod + + config = tmp_path / "config.json" + config.write_text(json.dumps({"subagents": {"runtime": {"idleCloseS": 0.2}}})) + monkeypatch.setattr(live_mod, "default_live", lambda: live_mod.LiveConfig(config)) + script = _product(tmp_path, "pass\n") + endpoint = await _launch(script) + await asyncio.wait_for(endpoint.wait(), timeout=60) + endpoint.release() + template = fork_template._TEMPLATE + assert template is not None + + for _ in range(100): + if not template.running: + break + await asyncio.sleep(0.05) + assert not template.running + + +async def test_anything_but_a_raven_agent_is_launched_the_ordinary_way(tmp_path: Path) -> None: + loose = tmp_path / "loose.py" + loose.write_text("pass\n") + base = dict(os.environ) + assert await fork_launch([sys.executable, str(loose)], base, str(tmp_path), base_env=base) is None + script = _product(tmp_path, "pass\n") + assert await fork_launch([sys.executable, str(script)], base, str(tmp_path / "missing"), base_env=base) is None + assert fork_template._TEMPLATE is None + + +async def test_a_template_that_cannot_start_is_not_tried_again(tmp_path: Path, monkeypatch) -> None: + script = _product(tmp_path, "pass\n") + started: list[int] = [] + + async def refuse(self, env): + started.append(1) + raise RuntimeError("no room") + + monkeypatch.setattr(fork_template.ForkTemplate, "_start", refuse) + base = dict(os.environ) + for _ in range(2): + assert await fork_launch([sys.executable, str(script)], base, str(tmp_path), base_env=base) is None + assert started == [1] + + +async def test_an_acp_launch_of_a_raven_agent_is_forked_and_offered_the_relay(tmp_path: Path, monkeypatch) -> None: + """Through the real launch: the agent runs from the template and its environment names the host's relay.""" + from raven.acp_client.client import AcpClient + from raven.acp_client.pool import close_pool + + stub = Path(__file__).with_name("acp_stub_server.py") + seen = tmp_path / "relay.txt" + script = _product( + tmp_path, + f""" + import os, runpy + open({str(seen)!r}, "w").write(os.environ.get("RAVEN_LLM_RELAY", "")) + runpy.run_path({str(stub)!r}, run_name="__main__") + """, + ) + client = await AcpClient.launch(name="forked", command=[sys.executable, str(script)], cwd=str(tmp_path)) + try: + assert "forked" in client.endpoint.describe() + result = await client.request("initialize", {"protocolVersion": 1, "clientCapabilities": {}}, timeout=60) + assert isinstance(result, dict) + assert seen.read_text().startswith(("unix:", "tcp:")) + finally: + await client.close() + await close_pool() + assert client.endpoint.returncode is not None + + +async def test_switching_both_off_launches_a_raven_agent_as_before(tmp_path: Path, monkeypatch) -> None: + from raven.acp_client.client import AcpClient + from raven.config import live as live_mod + + config = tmp_path / "config.json" + config.write_text(json.dumps({"subagents": {"runtime": {"llmRelay": False, "forkTemplate": False}}})) + monkeypatch.setattr(live_mod, "default_live", lambda: live_mod.LiveConfig(config)) + stub = Path(__file__).with_name("acp_stub_server.py") + seen = tmp_path / "relay.txt" + script = _product( + tmp_path, + f""" + import os, runpy + open({str(seen)!r}, "w").write(os.environ.get("RAVEN_LLM_RELAY", "none")) + runpy.run_path({str(stub)!r}, run_name="__main__") + """, + ) + client = await AcpClient.launch(name="plain", command=[sys.executable, str(script)], cwd=str(tmp_path)) + try: + assert "forked" not in client.endpoint.describe() + await client.request("initialize", {"protocolVersion": 1, "clientCapabilities": {}}, timeout=60) + assert seen.read_text() == "none" + finally: + await client.close() + + +# ---- the template's side, driven in this process --------------------------------- + + +class _Ended(Exception): + def __init__(self, code: int) -> None: + super().__init__(code) + self.code = code + + +@pytest.fixture +def ended(monkeypatch: pytest.MonkeyPatch) -> list[int]: + """Run the agent-side functions here: the exit becomes an exception, and what they change is put back. + + The codes are recorded because the real exit never returns, while this one + unwinds through the frames above it, which may read it as a second ending. + """ + codes: list[int] = [] + + def finish(code: int) -> None: + codes.append(code) + raise _Ended(code) + + monkeypatch.setattr(fork_template, "_finish", finish) + monkeypatch.setattr(os, "execv", os.execv) + monkeypatch.setattr(sys, "argv", list(sys.argv)) + monkeypatch.setattr(sys, "path", list(sys.path)) + return codes + + +def _learned(**products: list[str]) -> dict: + return {"python": list(sys.version_info[:2]), "products": products} + + +def test_the_template_loads_the_seed_and_what_agents_imported_but_not_what_it_skips( + tmp_path: Path, monkeypatch +) -> None: + for name in ("colorsys", "netrc", "this_module_does_not_exist"): + monkeypatch.delitem(sys.modules, name, raising=False) + learned = tmp_path / "modules.json" + learned.write_text(json.dumps(_learned(only=["netrc", "this_module_does_not_exist"]))) + + fork_template._load(("colorsys",), learned, ()) + assert "colorsys" in sys.modules and "netrc" in sys.modules + + monkeypatch.delitem(sys.modules, "netrc") + fork_template._load((), learned, ("net",)) + assert "netrc" not in sys.modules + + learned.write_text(json.dumps({"python": [2, 7], "products": {"only": ["netrc"]}})) + fork_template._load((), learned, ()) + assert "netrc" not in sys.modules + + +def test_one_products_own_modules_are_loaded_ahead_only_while_it_is_the_only_one_seen() -> None: + """What only one product imports is shared by no other process, so it is not loaded ahead for the rest.""" + assert fork_template.preload_names(_learned(code=["core", "code.only"])) == ["code.only", "core"] + two = _learned( + code=["core", "both", "code.only"], ppt=["core", "both", "deck.engine"], design=["core", "design.only"] + ) + assert fork_template.preload_names(two) == ["both", "core"] + assert fork_template.preload_names(_learned()) == [] + assert fork_template.preload_names({"python": [2, 7], "products": {"code": ["core"]}}) == [] + assert fork_template.preload_names({"python": list(sys.version_info[:2]), "modules": ["core"]}) == [] + assert fork_template.preload_names(["not", "a", "record"]) == [] + assert fork_template.preload_names(_learned(code=["core", 7], odd="not a list")) == ["core"] + + +def test_an_agent_is_learned_under_its_product_folder() -> None: + assert fork_template.product_of([sys.executable, "/x/agents/raven-code/run.py", "--acp"]) == "raven-code" + assert fork_template.product_of([sys.executable, "-m", "raven", "acp"]) == "raven" + assert fork_template.product_of([sys.executable, "run.py"]) == "raven" + + +def test_an_agent_adds_to_its_products_record_and_starts_over_on_another_python(tmp_path: Path) -> None: + learned = tmp_path / "t" / "modules.json" + learned.parent.mkdir() + learned.write_text(json.dumps(_learned(code=["kept.before"], ppt=["deck.engine"]))) + + fork_template._learn(learned, ("json",), "code") + products = json.loads(learned.read_text())["products"] + assert "kept.before" in products["code"] and "os" in products["code"] + assert not any(n.startswith("json") for n in products["code"]) and "__main__" not in products["code"] + assert products["ppt"] == ["deck.engine"], "another product's record is left as it was" + + learned.write_text(json.dumps({"python": [2, 7], "products": {"code": ["stale.one"]}})) + fork_template._learn(learned, (), "code") + assert "stale.one" not in json.loads(learned.read_text())["products"]["code"] + + +def test_an_agents_end_is_read_the_way_an_interpreter_reads_it(capsys) -> None: + assert fork_template._exit_code(SystemExit(None)) == 0 + assert fork_template._exit_code(SystemExit(4)) == 4 + assert fork_template._exit_code(SystemExit("bad config")) == 1 + assert "bad config" in capsys.readouterr().err + assert fork_template._exit_code(ValueError("boom")) == 1 + assert "ValueError: boom" in capsys.readouterr().err + + +def test_a_script_runs_with_its_own_directory_first_on_the_path(tmp_path: Path, ended) -> None: + script = _product( + tmp_path, "import sys\nsys.exit(len(sys.argv) * 10 + (sys.path[0] == __file__.rsplit('/', 1)[0]))\n" + ) + with pytest.raises(_Ended) as caught: + fork_template._run_agent([sys.executable, str(script), "a", "b"]) + assert caught.value.code == 31 + + +def test_a_launchers_exec_of_raven_becomes_raven_here_and_any_other_exec_is_real( + tmp_path: Path, ended, monkeypatch, capsys +) -> None: + script = _product( + tmp_path, "import os, sys\nos.execv(sys.executable, [sys.executable, '-m', 'raven', '--version'])\n" + ) + with pytest.raises(_Ended): + fork_template._run_agent([sys.executable, str(script)]) + assert ended[0] == 0 + assert "Raven v" in capsys.readouterr().out + + execs: list[list[str]] = [] + + def real(path, args): + execs.append([str(path), *args]) + raise _Ended(99) + + monkeypatch.setattr(os, "execv", real) + other = _product(tmp_path, "import os\nos.execv('/bin/true', ['/bin/true'])\n") + with pytest.raises(_Ended): + fork_template._run_agent([sys.executable, str(other)]) + assert execs == [["/bin/true", "/bin/true"]] + + +def test_a_module_launch_runs_raven_here(ended, capsys) -> None: + with pytest.raises(_Ended) as caught: + fork_template._run_agent([sys.executable, "-m", "raven", "--version"]) + assert caught.value.code == 0 + assert "Raven v" in capsys.readouterr().out + + +# The test process runs threads (the host side below, pytest's own); the forked +# child does nothing but exit, which is what the warning's deadlock cannot reach. +@pytest.mark.filterwarnings("ignore:This process .* is multi-threaded:DeprecationWarning") +def test_the_template_loop_forks_reports_the_exit_and_stops_when_let_go(tmp_path: Path, monkeypatch) -> None: + import socket + import threading + + from raven.providers import relay_wire + + counts = iter([2]) + monkeypatch.setattr(fork_template.threading, "active_count", lambda: next(counts, 1)) + path = str(tmp_path / "t.sock") + life_r, life_w = os.pipe() + os.write(life_w, b"x") + seen: dict[str, object] = {} + + def host() -> None: + while not os.path.exists(path): + time.sleep(0.01) + pipes = [os.pipe() for _ in range(3)] + with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as bad: + bad.connect(path) + socket.send_fds(bad, [b"F"], [pipes[0][0]]) + relay_wire.send_frame(bad, {"argv": ["x"], "cwd": "/"}) + seen["refused"] = relay_wire.recv_frame(bad) + with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as busy: + busy.connect(path) + socket.send_fds(busy, [b"F"], [pipes[0][0], pipes[1][1], pipes[2][1]]) + relay_wire.send_frame(busy, {"argv": [sys.executable, "x.py"], "env": {}, "cwd": "/"}) + seen["threaded"] = relay_wire.recv_frame(busy) + with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as sock: + sock.connect(path) + socket.send_fds(sock, [b"F"], [pipes[0][0], pipes[1][1], pipes[2][1]]) + relay_wire.send_frame(sock, {"argv": [sys.executable, "x.py"], "env": {}, "cwd": "/"}) + seen["forked"] = relay_wire.recv_frame(sock) + seen["exit"] = relay_wire.recv_frame(sock) + os.close(life_w) + + thread = threading.Thread(target=host) + thread.start() + forked = fork_template._serve(path, lifeline=life_r) + if forked is not None: + os._exit(9 if forked[0]["argv"] == [sys.executable, "x.py"] and len(forked[1]) == 3 else 1) + thread.join(timeout=30) + os.close(life_r) + + assert "error" in seen["refused"] + assert "thread" in seen["threaded"]["error"] + assert isinstance(seen["forked"].get("pid"), int) + assert seen["exit"] == {"exit": 9} + assert signal.getsignal(signal.SIGCHLD) == signal.SIG_DFL + + +def test_an_argv_the_check_cannot_read_is_not_forked(monkeypatch, tmp_path: Path) -> None: + script = _product(tmp_path, "pass\n") + assert not raven_agent_argv([None, str(script)]) # type: ignore[list-item] + + def unreadable(self) -> bool: + raise PermissionError("denied") + + monkeypatch.setattr(Path, "is_file", unreadable) + assert not raven_agent_argv([sys.executable, str(script)]) + + +async def test_a_closed_agent_takes_no_more_input_and_its_template_counts_none(tmp_path: Path) -> None: + script = _product(tmp_path, "import sys\nsys.stdin.read()\n") + endpoint = await _launch(script) + await endpoint.end() + await asyncio.wait_for(endpoint.wait(), timeout=30) + with pytest.raises(BrokenPipeError): + await endpoint.write(b"late\n") + endpoint.release() + template = fork_template._TEMPLATE + assert template is not None and template.children == 0 + + +async def test_a_template_that_ended_on_its_own_is_replaced_on_the_next_launch(tmp_path: Path) -> None: + script = _product(tmp_path, "pass\n") + first = await _launch(script) + await asyncio.wait_for(first.wait(), timeout=30) + first.release() + template = fork_template._TEMPLATE + assert template is not None and template._proc is not None + old = template._proc.pid + os.kill(old, signal.SIGKILL) + await asyncio.wait_for(template._proc.wait(), timeout=30) + + second = await _launch(script) + await asyncio.wait_for(second.wait(), timeout=30) + second.release() + assert template._proc is not None and template._proc.pid != old + + +async def test_a_launch_cancels_the_pending_idle_close(tmp_path: Path) -> None: + script = _product(tmp_path, "pass\n") + first = await _launch(script) + await asyncio.wait_for(first.wait(), timeout=30) + first.release() + template = fork_template._TEMPLATE + assert template is not None + template.close_when_idle(3600) + pending = template._idle_task + assert pending is not None + template.close_when_idle(3600) + assert template._idle_task is pending, "a second request does not stack a second timer" + + second = await _launch(script) + await asyncio.wait_for(second.wait(), timeout=30) + second.release() + await asyncio.sleep(0) + assert pending.cancelled() + template.close_when_idle(None) + assert template._idle_task is None or template._idle_task is not pending + + +async def test_a_template_that_refuses_or_cannot_be_reached_fails_the_fork_alone(tmp_path: Path) -> None: + script = _product(tmp_path, "pass\n") + first = await _launch(script) + await asyncio.wait_for(first.wait(), timeout=30) + first.release() + template = fork_template._TEMPLATE + assert template is not None + + with pytest.raises(RuntimeError, match="refused"): + await template._fork("not an argv", dict(os.environ), str(tmp_path)) # type: ignore[arg-type] + real = template._socket + template._socket = str(tmp_path / "nobody-listens.sock") + try: + with pytest.raises(OSError): + await template._fork([sys.executable, str(script)], dict(os.environ), str(tmp_path)) + finally: + template._socket = real + assert not template.broken + + +async def test_a_template_that_never_reports_ready_is_not_used(tmp_path: Path, monkeypatch) -> None: + async def silent(proc): + return {"ready": False, "reason": "an import started a thread (x)"} + + monkeypatch.setattr(fork_template.ForkTemplate, "_ready_report", staticmethod(silent)) + script = _product(tmp_path, "pass\n") + base = dict(os.environ) + assert await fork_launch([sys.executable, str(script)], base, str(tmp_path), base_env=base) is None + template = fork_template._TEMPLATE + assert template is not None and "thread" in template.broken + with pytest.raises(RuntimeError): + await template.spawn([sys.executable, str(script)], base, str(tmp_path), base) + + +async def test_the_ready_line_is_found_among_what_an_import_printed() -> None: + from types import SimpleNamespace + + def proc(*lines: bytes): + stream = asyncio.StreamReader() + for line in lines: + stream.feed_data(line) + stream.feed_eof() + return SimpleNamespace(stdout=stream, stderr=stream) + + report = await fork_template.ForkTemplate._ready_report(proc(b"noise\n", b'{"other": 1}\n', b'{"ready": true}\n')) + assert report == {"ready": True} + assert await fork_template.ForkTemplate._ready_report(proc(b"noise\n")) == {} + assert await fork_template.ForkTemplate._ready_report(SimpleNamespace(stdout=None)) == {} + await fork_template.ForkTemplate._log_stderr(proc(b"a warning\n")) + await fork_template.ForkTemplate._log_stderr(SimpleNamespace(stderr=None)) + + +async def test_a_template_built_for_other_skips_is_replaced_while_it_has_no_agent(tmp_path: Path) -> None: + script = _product(tmp_path, "pass\n") + base = dict(os.environ) + first = await fork_launch([sys.executable, str(script)], base, str(tmp_path), base_env=base, skip=("litellm",)) + assert first is not None + await asyncio.wait_for(first.wait(), timeout=30) + first.release() + old = fork_template._TEMPLATE + + second = await fork_launch([sys.executable, str(script)], base, str(tmp_path), base_env=base) + assert second is not None + await asyncio.wait_for(second.wait(), timeout=30) + second.release() + assert fork_template._TEMPLATE is not old and old is not None and not old.running + + +def test_learning_survives_an_unreadable_record_and_an_unwritable_place(tmp_path: Path) -> None: + fork_template._load((), tmp_path / "missing.json", ()) + corrupt = tmp_path / "modules.json" + corrupt.write_text("{not json") + fork_template._learn(corrupt, (), "code") + assert "os" in json.loads(corrupt.read_text())["products"]["code"] + corrupt.write_text(json.dumps(["not", "a", "record"])) + fork_template._learn(corrupt, (), "code") + assert "os" in json.loads(corrupt.read_text())["products"]["code"] + blocked = tmp_path / "file-not-dir" + blocked.write_text("") + fork_template._learn(blocked / "modules.json", (), "code") + + +def test_a_script_that_just_ends_exits_zero(tmp_path: Path, ended) -> None: + script = _product(tmp_path, "x = 1\n") + with pytest.raises(_Ended): + fork_template._run_agent([sys.executable, str(script)]) + assert ended == [0] + + +def test_the_template_entry_loads_freezes_and_serves(tmp_path: Path, monkeypatch) -> None: + import gc + + served: list[str] = [] + became: list[tuple] = [] + monkeypatch.setattr(gc, "freeze", lambda: None) + monkeypatch.setattr(fork_template, "_load", lambda seed, learned, skip: 0) + monkeypatch.setattr(fork_template.threading, "active_count", lambda: 1) + + def serve(path, *, lifeline): + served.append(path) + return None + + monkeypatch.setattr(fork_template, "_serve", serve) + args = ["--socket", str(tmp_path / "s"), "--learned", str(tmp_path / "m.json"), "--skip", "litellm"] + assert fork_template.main(args) == 0 + assert served == [str(tmp_path / "s")] + + monkeypatch.setattr(fork_template, "_serve", lambda path, *, lifeline: ({"argv": ["x"]}, [1, 2, 3])) + monkeypatch.setattr(fork_template, "_become", lambda *a: became.append(a)) + assert fork_template.main(args) == 1 + assert became[0][0] == {"argv": ["x"]} and became[0][4] == ["litellm"] + + +def test_the_template_entry_refuses_to_serve_with_a_thread_running(tmp_path: Path, monkeypatch, capsys) -> None: + import gc + + monkeypatch.setattr(gc, "freeze", lambda: None) + monkeypatch.setattr(fork_template, "_load", lambda seed, learned, skip: 0) + monkeypatch.setattr(fork_template.threading, "active_count", lambda: 2) + args = ["--socket", str(tmp_path / "s"), "--learned", str(tmp_path / "m.json")] + assert fork_template.main(args) == 3 + assert json.loads(capsys.readouterr().out.strip().splitlines()[-1])["ready"] is False + + +async def test_a_relay_that_cannot_start_leaves_the_agent_to_load_litellm(tmp_path: Path, monkeypatch) -> None: + from raven.acp_client.client import AcpClient + from raven.providers import relay_server + + async def no_relay(): + raise OSError("no room for a socket") + + monkeypatch.setattr(relay_server, "relay_env", no_relay) + stub = Path(__file__).with_name("acp_stub_server.py") + seen = tmp_path / "relay.txt" + script = _product( + tmp_path, + f""" + import os, runpy + open({str(seen)!r}, "w").write(os.environ.get("RAVEN_LLM_RELAY", "none")) + runpy.run_path({str(stub)!r}, run_name="__main__") + """, + ) + client = await AcpClient.launch(name="norelay", command=[sys.executable, str(script)], cwd=str(tmp_path)) + try: + await client.request("initialize", {"protocolVersion": 1, "clientCapabilities": {}}, timeout=60) + assert seen.read_text() == "none" + finally: + await client.close() + + +def test_a_new_event_loop_gets_a_new_template_and_the_old_one_ends(tmp_path: Path) -> None: + """A CLI turn is one asyncio.run; a template opened under the last one cannot serve the next.""" + script = _product(tmp_path, "print('ran', flush=True)\n") + + async def once() -> tuple[int, bytes]: + endpoint = await _launch(script) + out, _ = await _drain(endpoint) + endpoint.release() + template = fork_template._TEMPLATE + assert template is not None and template._proc is not None + return template._proc.pid, out + + first_pid, first_out = asyncio.run(once()) + second_pid, second_out = asyncio.run(once()) + try: + assert first_out == second_out == b"ran\n" + assert second_pid != first_pid + for _ in range(100): + try: + os.kill(first_pid, 0) + except ProcessLookupError: + break + time.sleep(0.05) + else: + pytest.fail("the template opened under the ended loop is still running") + finally: + asyncio.run(fork_template.close_template()) diff --git a/tests/test_acp_pool_idle_close.py b/tests/test_acp_pool_idle_close.py new file mode 100644 index 000000000..90bb861dd --- /dev/null +++ b/tests/test_acp_pool_idle_close.py @@ -0,0 +1,226 @@ +"""The ACP pool closes an agent's process once it has been idle for the configured budget.""" + +from __future__ import annotations + +import asyncio +import json +import sys +import time +from pathlib import Path +from types import SimpleNamespace +from typing import Any + +import pytest + +from raven.acp_client import pool as pool_mod +from raven.acp_client.pool import AcpConnectionPool, _Connection, _SessionElicitors, _SessionRouter, close_pool +from raven.utils.commands import command_quote + +_STUB = Path(__file__).with_name("acp_stub_server.py") +_LOADS = {"agentCapabilities": {"loadSession": True}} +_NO_LOAD = {"agentCapabilities": {}} + + +@pytest.fixture(autouse=True) +async def _no_pooled_connections(): + yield + await close_pool() + + +class _Client: + def __init__(self, *, idle: float = 0.0, busy: bool = False, alive: bool = True) -> None: + self.name = "agent" + self._idle_since = time.monotonic() - idle + self.busy = busy + self.alive = alive + self.closed = False + + @property + def idle_seconds(self) -> float: + return time.monotonic() - self._idle_since + + async def close(self) -> None: + self.closed = True + self.alive = False + + +def _held(client: _Client, initialize: Any, *, leased_ago: float = 3600.0, seen: tuple[str, ...] = ()) -> _Connection: + router = _SessionRouter("agent") + router.seen.update(seen) + connection = _Connection(client, router, _SessionElicitors("agent"), initialize=initialize) # type: ignore[arg-type] + connection.leased_at = time.monotonic() - leased_ago + return connection + + +def _pool_holding(**connections: _Connection) -> AcpConnectionPool: + pool = AcpConnectionPool() + for name, connection in connections.items(): + pool._connections[name] = {"": connection} + return pool + + +async def test_a_quiet_agent_that_can_reload_its_sessions_is_closed() -> None: + client = _Client(idle=120) + pool = _pool_holding(quiet=_held(client, _LOADS, seen=("s1",))) + + assert await pool.close_idle(60) == ["quiet"] + assert client.closed + assert pool.live("quiet") == [] + + +async def test_a_connection_whose_agent_already_ended_is_dropped() -> None: + gone = _Client(idle=0, alive=False) + pool = _pool_holding(gone=_held(gone, _NO_LOAD, leased_ago=0, seen=("s1",))) + + assert await pool.close_idle(60) == ["gone"] + assert pool._connections == {} + + +async def test_nothing_is_closed_while_a_request_is_in_flight_either_way() -> None: + """A held wake is a prompt still waiting for its answer, so it reads as busy too.""" + client = _Client(idle=120, busy=True) + pool = _pool_holding(working=_held(client, _LOADS)) + + assert await pool.close_idle(60) == [] + assert not client.closed + + +async def test_a_connection_just_handed_out_is_left_for_its_caller() -> None: + client = _Client(idle=120) + pool = _pool_holding(leased=_held(client, _LOADS, leased_ago=1.0)) + + assert await pool.close_idle(60) == [] + assert not client.closed + + +async def test_an_agent_heard_from_recently_is_kept() -> None: + client = _Client(idle=5) + pool = _pool_holding(chatty=_held(client, _LOADS)) + + assert await pool.close_idle(60) == [] + + +async def test_an_agent_that_cannot_reload_keeps_the_sessions_it_opened() -> None: + """Closing it would hand the next turn of a stateful instance a blank session.""" + keeps = _Client(idle=120) + unused = _Client(idle=120) + pool = _pool_holding( + holds=_held(keeps, _NO_LOAD, seen=("s1",)), + never_used=_held(unused, _NO_LOAD), + ) + + assert await pool.close_idle(60) == ["never_used"] + assert not keeps.closed + assert unused.closed + + +async def test_a_budget_of_zero_or_none_in_the_config_turns_closing_off( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + from raven.config import live as live_mod + + config = tmp_path / "config.json" + monkeypatch.setattr(live_mod, "default_live", lambda: live_mod.LiveConfig(config)) + client = _Client(idle=10_000) + pool = _pool_holding(quiet=_held(client, _LOADS)) + + for value in (0, None): + config.write_text(json.dumps({"subagents": {"runtime": {"idleCloseS": value}}}), encoding="utf-8") + assert await pool.close_idle() == [] + config.write_text(json.dumps({"subagents": {"runtime": {"idleCloseS": 30}}}), encoding="utf-8") + assert await pool.close_idle() == ["quiet"] + + +async def test_the_idle_check_starts_with_a_connection_and_ends_with_the_last(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(pool_mod, "_IDLE_CHECK_MIN_S", 0.01) + monkeypatch.setattr(AcpConnectionPool, "_idle_budget", staticmethod(lambda: 0.04)) + client = _Client(idle=0) + pool = _pool_holding(soon=_held(client, _LOADS, leased_ago=0)) + pool._watch_idle() + task = pool._idle_task + assert task is not None + + await asyncio.wait_for(task, timeout=5) + + assert client.closed + assert pool._connections == {} + + +async def test_an_idle_agent_process_ends_and_the_next_dispatch_starts_a_fresh_one(tmp_path: Path) -> None: + pool = AcpConnectionPool() + command = f"{command_quote(sys.executable)} {command_quote(str(_STUB))}" + first = await pool.acquire(name="stub", command=command, cwd=str(tmp_path), env={"ACP_STUB_MODE": "ok"}) + opened = await first.client.request("session/new", {"cwd": str(tmp_path), "mcpServers": []}) + first.router.reserve(opened["sessionId"]) + first_pid = first.client.endpoint.pid + first.leased_at -= 10 + + await asyncio.sleep(0.3) + assert await pool.close_idle(0.2) == ["stub"] + assert not first.alive + + second = await pool.acquire(name="stub", command=command, cwd=str(tmp_path), env={"ACP_STUB_MODE": "ok"}) + assert second.alive + assert second.client.endpoint.pid != first_pid + await pool.close_all() + + +def test_the_idle_budget_is_read_from_the_host_config(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None: + from raven.config import live as live_mod + + config = tmp_path / "config.json" + config.write_text(json.dumps({"subagents": {"runtime": {"idleCloseS": 42}}}), encoding="utf-8") + monkeypatch.setattr(live_mod, "default_live", lambda: live_mod.LiveConfig(config)) + assert AcpConnectionPool._idle_budget() == 42 + config.write_text("{}", encoding="utf-8") + assert AcpConnectionPool._idle_budget() == 600 + + +def test_a_busy_client_property_reads_requests_both_ways() -> None: + from raven.acp_client.client import AcpClient + + client = AcpClient(name="x", endpoint=SimpleNamespace(returncode=None)) # type: ignore[arg-type] + assert not client.busy + loop = asyncio.new_event_loop() + try: + client._pending[1] = loop.create_future() + assert client.busy + client._pending.clear() + client._answer_tasks.add(object()) # type: ignore[arg-type] + assert client.busy + finally: + loop.close() + + +def test_a_rejected_edit_keeps_the_last_valid_settings(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None: + from raven.config import live as live_mod + + config = tmp_path / "config.json" + reader = live_mod.LiveConfig(config) + config.write_text( + json.dumps({"subagents": {"runtime": {"idleCloseS": 30, "forkTemplate": False}}}), encoding="utf-8" + ) + assert live_mod.subagent_runtime(reader).idle_close_s == 30 + config.write_text(json.dumps({"subagents": {"runtime": {"idleCloseS": "soon"}}}), encoding="utf-8") + kept = live_mod.subagent_runtime(reader) + assert kept.idle_close_s == 30 and kept.fork_template is False + + +async def test_a_failed_idle_check_does_not_end_the_next_one(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(pool_mod, "_IDLE_CHECK_MIN_S", 0.01) + monkeypatch.setattr(AcpConnectionPool, "_idle_budget", staticmethod(lambda: 0.04)) + client = _Client(idle=0) + pool = _pool_holding(soon=_held(client, _LOADS, leased_ago=0)) + real = pool.close_idle + calls: list[int] = [] + + async def flaky(budget=None): + calls.append(1) + if len(calls) == 1: + raise RuntimeError("once") + return await real(budget) + + monkeypatch.setattr(pool, "close_idle", flaky) + pool._watch_idle() + await asyncio.wait_for(pool._idle_task, timeout=5) + assert len(calls) >= 2 and client.closed diff --git a/tests/test_provider_auth_method.py b/tests/test_provider_auth_method.py index 8c5a34781..fa6060ced 100644 --- a/tests/test_provider_auth_method.py +++ b/tests/test_provider_auth_method.py @@ -468,6 +468,9 @@ def test_only_the_auth_module_decides_configuredness_from_a_key() -> None: "raven/agent/tools/capabilities.py", "raven/config/update_providers.py", "raven/providers/litellm_provider.py", + # Forwards the key the agent's provider was built with to the host that + # sends the request with it (the LLM relay); rules on nothing. + "raven/providers/relay_provider.py", "raven/providers/factory.py", # The router reads the OpenRouter key in order to call OpenRouter with it; # the read moved here from the assembly root with the vendor knowledge, diff --git a/tests/test_provider_catalog.py b/tests/test_provider_catalog.py index 85a97d88a..d2fff982b 100644 --- a/tests/test_provider_catalog.py +++ b/tests/test_provider_catalog.py @@ -202,6 +202,7 @@ def _concrete_provider_subclasses() -> set[type]: import raven.providers.minimax_oauth_provider # noqa: F401 import raven.providers.openai_codex_provider # noqa: F401 import raven.providers.per_model_provider # noqa: F401 + import raven.providers.relay_provider # noqa: F401 seen: set[type] = set() stack = list(LLMProvider.__subclasses__()) @@ -220,7 +221,7 @@ def _concrete_provider_subclasses() -> set[type]: return seen -def test_exactly_seven_concrete_backend_classes() -> None: +def test_exactly_these_concrete_backend_classes() -> None: # This asserts class existence only, not the dispatch wiring. from raven.providers.anthropic_messages_provider import AnthropicMessagesProvider from raven.providers.azure_openai_provider import AzureOpenAIProvider @@ -230,6 +231,7 @@ def test_exactly_seven_concrete_backend_classes() -> None: from raven.providers.openai_codex_provider import OpenAICodexProvider from raven.providers.openai_responses_provider import OpenAIResponsesProvider from raven.providers.per_model_provider import PerModelProvider + from raven.providers.relay_provider import RelayProvider from raven.providers.resolving_provider import ResolvingProvider expected = { @@ -247,6 +249,9 @@ def test_exactly_seven_concrete_backend_classes() -> None: # the LiteLLM chat path (providers..protocol / modelProtocols). OpenAIResponsesProvider, AnthropicMessagesProvider, + # The LiteLLM chat path when the host that launched this agent makes its + # calls (the LLM relay): built in LiteLLMProvider's place by the factory. + RelayProvider, } assert _concrete_provider_subclasses() == expected for cls in expected: diff --git a/tests/test_provider_prompt_cache.py b/tests/test_provider_prompt_cache.py index 43494319e..ffdfb50e2 100644 --- a/tests/test_provider_prompt_cache.py +++ b/tests/test_provider_prompt_cache.py @@ -376,7 +376,7 @@ def test_every_module_that_writes_the_field_asks_whether_it_may(): writers = { "raven/providers/prompt_cache.py", # the answer itself - "raven/providers/litellm_provider.py", + "raven/providers/litellm_shape.py", # Marks its own system/tool blocks, and asks accepts_cache_control first. "raven/providers/anthropic_messages_provider.py", "raven/token_wise/cache_optimizer.py", diff --git a/tests/test_provider_relay.py b/tests/test_provider_relay.py new file mode 100644 index 000000000..7b2b8aca8 --- /dev/null +++ b/tests/test_provider_relay.py @@ -0,0 +1,776 @@ +"""The LLM relay: a Raven agent's LiteLLM calls made by the host that launched it.""" + +from __future__ import annotations + +import asyncio +import json +import os +import subprocess +import sys +import textwrap +from collections.abc import AsyncIterator +from typing import Any + +import httpx +import pytest + +from raven.contracts.llm_provider import ( + CallRecord, + ChatDelta, + ErrorClassification, + GenerationSettings, + LLMResponse, + RunMeta, + ToolCallRequest, + TruncationInfo, +) +from raven.providers import call_record, prompt_cache, relay, relay_provider, relay_server, relay_wire, usage_record +from raven.providers.base import LLMProvider +from raven.providers.first_byte import FirstByteTimeoutError, StreamIdleTimeoutError +from raven.providers.litellm_shape import LiteLLMShape + + +@pytest.fixture(autouse=True) +def _fresh_relay_state(monkeypatch: pytest.MonkeyPatch): + monkeypatch.delenv(relay.RELAY_ENV, raising=False) + monkeypatch.delenv(relay.RELAY_TOKEN_ENV, raising=False) + relay._reset_for_tests() + yield + relay._reset_for_tests() + + +# ---- the wire ---------------------------------------------------------------- + + +def test_the_provider_vocabulary_survives_the_trip() -> None: + record = CallRecord(http_status=200, served_by="vendor", headers={"x-request-id": "r1"}, body="ok") + response = LLMResponse( + content="hi", + tool_calls=[ + ToolCallRequest( + id="t1", + name="read_file", + arguments={"path": "a.py"}, + run_meta=RunMeta(truncation=TruncationInfo(at_tokens=12), arguments_repaired=True), + ) + ], + finish_reason="tool_calls", + usage={"prompt_tokens": 3}, + thinking_blocks=[{"type": "thinking", "thinking": "x", "signature": "s"}], + error_classification=ErrorClassification("network", retryable=True), + truncated=True, + max_tokens=100, + reasoning_ms=5, + call_record=record, + ) + delta = ChatDelta(content="a", tool_call_delta={"tool_calls": []}, finish_reason="stop", call_record=record) + generation = GenerationSettings(temperature=0.2, max_tokens=9, reasoning_effort="high") + + for value in (response, delta, generation): + assert relay_wire.decode(json.loads(json.dumps(relay_wire.encode(value)))) == value + + +def test_the_not_given_marker_becomes_the_far_sides_own() -> None: + here, there = object(), object() + encoded = json.loads(json.dumps(relay_wire.encode({"max_tokens": here, "model": "m"}, sentinel=here))) + assert relay_wire.decode(encoded, sentinel=there) == {"max_tokens": there, "model": "m"} + + +def test_a_payload_a_client_library_left_objects_in_still_travels() -> None: + from pydantic import BaseModel + + class Usage(BaseModel): + prompt_tokens: int = 4 + + assert relay_wire.decode(relay_wire.encode({"usage": Usage(), "n": (1, 2)})) == { + "usage": {"prompt_tokens": 4}, + "n": [1, 2], + } + + +def _litellm_errors() -> list[BaseException]: + import litellm + + request = httpx.Request("POST", "https://example.invalid/v1/chat/completions") + rate = httpx.Response(429, request=request, headers={"retry-after": "7"}, text='{"error": "slow down"}') + errors: list[BaseException] = [ + litellm.RateLimitError("rate limited", llm_provider="openai", model="gpt-4o", response=rate), + litellm.ContextWindowExceededError("context length exceeded", model="gpt-4o", llm_provider="openai"), + litellm.BadRequestError( + "cache_control is not permitted: invalid_request_error", model="claude", llm_provider="anthropic" + ), + litellm.APIConnectionError("connection reset", llm_provider="openai", model="gpt-4o"), + litellm.InternalServerError("upstream exploded", llm_provider="openai", model="gpt-4o"), + ] + return errors + + +def _plain_errors() -> list[BaseException]: + chained = ValueError("parse failed") + try: + try: + raise httpx.ReadTimeout("read timed out") + except httpx.ReadTimeout as inner: + raise chained from inner + except ValueError: + pass + return [ + FirstByteTimeoutError(phase="waiting for headers", budget=120.0, waited=120.4), + StreamIdleTimeoutError(idle=180.0), + TimeoutError(), + ConnectionResetError("peer reset"), + httpx.ReadTimeout("read timed out"), + chained, + ] + + +@pytest.mark.parametrize("make", [_litellm_errors, _plain_errors]) +def test_a_relayed_failure_is_judged_as_the_original_would_be(make) -> None: + for original in make(): + status = LLMProvider._extract_status_code(original) + rebuilt = relay_wire.decode_exception( + json.loads(json.dumps(relay_wire.encode_exception(original, status=status))) + ) + + assert LLMProvider.classify_error(rebuilt) == LLMProvider.classify_error(original), original + assert str(rebuilt) == str(original) + assert LLMProvider._extract_status_code(rebuilt) == status + assert isinstance(rebuilt, TimeoutError) == isinstance(original, TimeoutError) + assert prompt_cache.is_rejection(rebuilt) == prompt_cache.is_rejection(original) + assert call_record.from_exception(rebuilt, status=status) == call_record.from_exception(original, status=status) + + +def test_the_two_timeouts_come_back_as_their_own_classes() -> None: + first = FirstByteTimeoutError(phase="connecting", budget=30.0, waited=30.2) + idle = StreamIdleTimeoutError(idle=90.0) + for original in (first, idle): + rebuilt = relay_wire.decode_exception(relay_wire.encode_exception(original, status=None)) + assert type(rebuilt) is type(original) + assert str(rebuilt) == str(original) + + +# ---- the provider and the server ------------------------------------------------------------ + + +class _Wire(LLMProvider): + """What the host builds in these tests: answers from a script, remembers how it was called.""" + + def __init__(self, spec: dict[str, Any], script: dict[str, Any]) -> None: + super().__init__() + self.spec = spec + self.script = script + + def get_default_model(self) -> str: + return "m" + + async def chat(self, **kwargs: Any) -> LLMResponse: + self.script.setdefault("calls", []).append(("chat", kwargs, usage_record._owned.get())) + return LLMResponse(content=f"answered {kwargs['messages'][-1]['content']}", usage={"prompt_tokens": 1}) + + async def chat_stream(self, **kwargs: Any) -> AsyncIterator[ChatDelta]: + self.script.setdefault("calls", []).append(("stream", kwargs, usage_record._owned.get())) + try: + for piece in self.script.get("pieces", ["a", "b"]): + yield ChatDelta(content=piece) + if self.script.get("hang"): + await asyncio.sleep(3600) + if failure := self.script.get("raise"): + raise failure + yield ChatDelta(content=None, finish_reason="stop", usage={"completion_tokens": 2}) + except asyncio.CancelledError: + self.script["cancelled"] = True + raise + + +@pytest.fixture(params=["unix", "tcp"]) +async def served(request, monkeypatch: pytest.MonkeyPatch): + """A relay with a scripted provider, on each transport: TCP is the one Windows hosts use.""" + script: dict[str, Any] = {"specs": []} + + def build(spec: dict[str, Any]) -> LLMProvider: + script["specs"].append(spec) + return _Wire(spec, script) + + server = relay_server.RelayServer(build=build, tcp=request.param == "tcp") + await server.start() + for key, value in server.env().items(): + monkeypatch.setenv(key, value) + relay._reset_for_tests() + yield script, server + await server.close() + + +async def test_a_turn_streams_through_the_host(served) -> None: + script, _ = served + provider = relay_provider.litellm_provider(default_model="custom/m", provider_name="custom", api_base="http://x/v1") + assert isinstance(provider, relay_provider.RelayProvider) + + deltas = [d async for d in provider.chat_stream(messages=[{"role": "user", "content": "q"}], model="custom/m")] + + assert [d.content for d in deltas] == ["a", "b", None] + assert deltas[-1].usage == {"completion_tokens": 2} + kind, kwargs, owned = script["calls"][0] + assert kind == "stream" and kwargs["model"] == "custom/m" + assert kwargs["max_tokens"] is _Wire._SENTINEL + assert owned, "the host's call must be marked recorded by its caller, the agent" + + +async def test_a_single_call_answers_through_the_host(served) -> None: + provider = relay_provider.RelayProvider(default_model="custom/m", provider_name="custom") + response = await provider.chat(messages=[{"role": "user", "content": "q"}]) + assert response.content == "answered q" + + +async def test_the_host_builds_from_the_construction_as_it_stands_at_the_call(served) -> None: + """Callers set ``generation`` and ``model_overrides`` after building a provider.""" + script, _ = served + provider = relay_provider.RelayProvider(default_model="custom/m", provider_name="custom", api_key="k") + provider.generation = GenerationSettings(temperature=0.1, reasoning_effort="low") + provider.model_overrides = {"m": {"top_p": 0.5}} + + await provider.chat(messages=[{"role": "user", "content": "q"}]) + + spec = script["specs"][-1] + assert spec["api_key"] == "k" + assert spec["model_overrides"] == {"m": {"top_p": 0.5}} + assert relay_wire.decode(spec["generation"]) == GenerationSettings(temperature=0.1, reasoning_effort="low") + + +async def test_a_failure_on_the_host_reaches_the_agent_as_the_same_verdict(served) -> None: + script, _ = served + script["raise"] = StreamIdleTimeoutError(idle=5.0) + provider = relay_provider.RelayProvider(default_model="custom/m", provider_name="custom") + + with pytest.raises(StreamIdleTimeoutError) as caught: + async for _ in provider.chat_stream(messages=[{"role": "user", "content": "q"}]): + pass + assert provider.classify_error(caught.value).category == LLMProvider.classify_error(script["raise"]).category + + +async def test_a_stream_the_agent_abandons_is_called_off_on_the_host(served) -> None: + script, _ = served + script["hang"] = True + provider = relay_provider.RelayProvider(default_model="custom/m", provider_name="custom") + + stream = provider.chat_stream(messages=[{"role": "user", "content": "q"}]) + assert (await stream.__anext__()).content == "a" + await stream.aclose() + + for _ in range(100): + if script.get("cancelled"): + break + await asyncio.sleep(0.02) + assert script.get("cancelled") + + +async def test_a_request_without_the_token_is_refused(served, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv(relay.RELAY_TOKEN_ENV, "not-the-token") + relay._reset_for_tests() + provider = relay_provider.RelayProvider(default_model="custom/m", provider_name="custom") + with pytest.raises(relay.RelayFaultError, match="token"): + await provider.chat(messages=[{"role": "user", "content": "q"}]) + + +async def test_an_unreachable_host_leaves_the_call_to_this_process(monkeypatch: pytest.MonkeyPatch, tmp_path) -> None: + monkeypatch.setenv(relay.RELAY_ENV, f"unix:{tmp_path / 'gone.sock'}") + monkeypatch.setenv(relay.RELAY_TOKEN_ENV, "t") + relay._reset_for_tests() + provider = relay_provider.RelayProvider(default_model="custom/m", provider_name="custom") + local = _Wire({}, {}) + monkeypatch.setattr(provider, "_locally", lambda: local) + + response = await provider.chat(messages=[{"role": "user", "content": "q"}]) + + assert response.content == "answered q" + assert not relay.active() + + +async def test_a_question_that_times_out_leaves_the_next_one_to_the_host( + served, monkeypatch: pytest.MonkeyPatch +) -> None: + """Only a host that cannot be reached at all ends the relay for the process.""" + from raven.providers import rates + + def slow_answer(*_args: Any, **_kwargs: Any) -> Any: + raise relay.RelayUnavailableError("timed out") + + monkeypatch.setattr(relay, "ask", slow_answer) + local = relay._RELAYED["raven.providers.rates:_may_query_endpoint"]("gpt-4o") + + assert await asyncio.to_thread(rates._may_query_endpoint, "gpt-4o") == local + assert relay.active() + + +async def test_a_host_that_cannot_be_reached_ends_the_relay_for_this_process( + monkeypatch: pytest.MonkeyPatch, tmp_path +) -> None: + from raven.providers import rates + + monkeypatch.setenv(relay.RELAY_ENV, f"tcp:127.0.0.1:{_closed_port()}") + monkeypatch.setenv(relay.RELAY_TOKEN_ENV, "t") + relay._reset_for_tests() + + assert rates._may_query_endpoint("gpt-4o") == relay._RELAYED["raven.providers.rates:_may_query_endpoint"]("gpt-4o") + assert not relay.active() + + +def _closed_port() -> int: + import socket + + with socket.socket() as sock: + sock.bind(("127.0.0.1", 0)) + return sock.getsockname()[1] + + +def test_an_unreadable_relay_address_is_refused() -> None: + with pytest.raises(relay.RelayFaultError): + relay._split("pipe:/x") + + +def test_the_offer_is_read_once_and_taken_out_of_the_environment(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv(relay.RELAY_ENV, "unix:/x") + monkeypatch.setenv(relay.RELAY_TOKEN_ENV, "t") + relay._reset_for_tests() + + assert relay.target() == ("unix:/x", "t") + assert relay.RELAY_ENV not in os.environ and relay.RELAY_TOKEN_ENV not in os.environ + assert relay.active() + + +def test_without_an_offer_the_direct_provider_is_built() -> None: + from raven.providers.litellm_provider import LiteLLMProvider + + provider = relay_provider.litellm_provider(default_model="openai/gpt-4o") + assert type(provider) is LiteLLMProvider + assert relay_provider.is_litellm_transport(provider) + + +def test_the_relayed_provider_answers_the_direct_ones_questions_from_the_same_code() -> None: + from raven.providers.litellm_provider import LiteLLMProvider + + for kwargs in ( + {"default_model": "openrouter/anthropic/claude-sonnet-5", "provider_name": "openrouter", "api_key": "k"}, + {"default_model": "custom/qwen", "provider_name": "custom", "api_base": "http://h/v1/"}, + {"default_model": "anthropic/claude-opus-5", "provider_name": "anthropic"}, + ): + direct, relayed = LiteLLMProvider(**kwargs), relay_provider.RelayProvider(**kwargs) + model = kwargs["default_model"] + assert isinstance(relayed, LiteLLMShape) + assert relayed.wire_model_id(model) == direct.wire_model_id(model) + assert relayed.supports_prompt_caching(model) == direct.supports_prompt_caching(model) + assert relayed.supports_assistant_prefill(model) == direct.supports_assistant_prefill(model) + assert relayed.emits_unparsed_reasoning() == direct.emits_unparsed_reasoning() + assert relayed.can_serve(model) == direct.can_serve(model) + assert relayed.extra_headers == direct.extra_headers + assert relayed.api_base == direct.api_base + assert relayed.request_generation(model=model) == direct.request_generation(model=model) + + +# ---- catalogue questions --------------------------------------------------------------------- + + +async def test_a_catalogue_question_is_answered_by_the_host_and_remembered( + served, monkeypatch: pytest.MonkeyPatch +) -> None: + from raven.providers import rates + + asked: list[str] = [] + real = relay_server.run_relayed + + def counting(name: str, args: list[Any], kwargs: dict[str, Any]) -> Any: + asked.append(name) + return real(name, args, kwargs) + + monkeypatch.setattr(relay_server, "run_relayed", counting) + expected = relay._RELAYED["raven.providers.rates:_table_entry"]("gpt-4o") + + # From a worker thread: the answer comes from this test's own loop. + first = await asyncio.to_thread(rates._table_entry, "gpt-4o") + second = await asyncio.to_thread(rates._table_entry, "gpt-4o") + + assert first == second == expected + assert asked == ["raven.providers.rates:_table_entry"] + + +async def test_only_registered_functions_of_listed_modules_run_on_the_host() -> None: + with pytest.raises(relay.RelayFaultError): + relay.run_relayed("os:system", ["true"], {}) + with pytest.raises(relay.RelayFaultError): + relay.run_relayed("raven.providers.rates:resolve_context_window", ["gpt-4o"], {}) + + +def test_the_drivers_directory_is_found_without_importing_litellm() -> None: + script = textwrap.dedent( + """ + import sys + from raven.providers import rates + found = rates._drivers_dir() + print(found is not None and found.name == "llms" and found.is_dir(), "litellm" in sys.modules) + """ + ) + out = subprocess.run([sys.executable, "-c", script], capture_output=True, text=True, timeout=120, check=True) + assert out.stdout.split() == ["True", "False"] + + +def test_a_relayed_process_never_loads_litellm(tmp_path) -> None: + """The point of the relay, end to end: a turn's call and its catalogue lookups, no LiteLLM import.""" + host = textwrap.dedent( + """ + import asyncio, json, sys + from raven.contracts.llm_provider import ChatDelta, LLMResponse + from raven.providers.base import LLMProvider + from raven.providers.relay_server import RelayServer + + class Wire(LLMProvider): + def get_default_model(self): return "m" + async def chat(self, **kw): return LLMResponse(content="ok") + async def chat_stream(self, **kw): + yield ChatDelta(content="ok", finish_reason="stop") + + async def main(): + server = RelayServer(build=lambda spec: Wire()) + await server.start() + print(json.dumps(server.env()), flush=True) + await asyncio.get_running_loop().run_in_executor(None, sys.stdin.read) + await server.close() + + asyncio.run(main()) + """ + ) + agent = textwrap.dedent( + """ + import asyncio, sys + from raven.providers import rates + from raven.providers.relay_provider import litellm_provider + + async def main(): + provider = litellm_provider(default_model="custom/m", provider_name="custom", api_base="http://h/v1") + text = "".join([d.content or "" async for d in provider.chat_stream(messages=[{"role": "user", "content": "q"}])]) + window = await asyncio.to_thread(rates.resolve_context_window, "gpt-4o") + print(type(provider).__name__, text, window is not None, "litellm" in sys.modules, "openai" in sys.modules) + + asyncio.run(main()) + """ + ) + server = subprocess.Popen([sys.executable, "-c", host], stdin=subprocess.PIPE, stdout=subprocess.PIPE, text=True) + try: + assert server.stdout is not None + env = json.loads(server.stdout.readline()) + out = subprocess.run( + [sys.executable, "-c", agent], capture_output=True, text=True, timeout=180, env={**os.environ, **env} + ) + finally: + if server.stdin is not None: + server.stdin.close() + server.wait(timeout=30) + assert out.returncode == 0, out.stderr + assert out.stdout.split() == ["RelayProvider", "ok", "True", "False", "False"] + + +def test_the_estimators_table_drops_its_unread_copy_and_still_counts() -> None: + from raven.utils import tokens + + enc = tokens.encoding() + assert not hasattr(enc, "_mergeable_ranks") + assert tokens.estimate_message_tokens({"role": "user", "content": "hello world"}) > 0 + + +# ---- edges of the wire, the server and the provider ------------------------------------------- + + +def test_the_default_builder_makes_the_direct_provider_with_the_agents_settings() -> None: + from raven.providers.litellm_provider import LiteLLMProvider + + generation = GenerationSettings(temperature=0.3) + spec = relay_provider.RelayProvider(default_model="custom/m", provider_name="custom", api_base="http://h/v1") + spec.generation = generation + built = relay_server.build_litellm_provider(spec.spec()) + assert isinstance(built, LiteLLMProvider) + assert built.generation == generation and built.default_model == "custom/m" + + +async def test_frames_that_are_not_frames_are_refused() -> None: + with pytest.raises(relay_wire.FrameError): + relay_wire._unpack(b"\xff") + with pytest.raises(relay_wire.FrameError): + relay_wire._unpack(b"[1]") + with pytest.raises(relay_wire.FrameError): + relay_wire._length((relay_wire.MAX_FRAME_BYTES + 1).to_bytes(4, "big")) + + for data in (b"\x00\x00", b"\x00\x00\x00\x05ab"): + reader = asyncio.StreamReader() + reader.feed_data(data) + reader.feed_eof() + with pytest.raises(relay_wire.FrameError): + await relay_wire.read_frame(reader) + + import socket + + left, right = socket.socketpair() + with left, right: + left.sendall(b"\x00\x00\x00\x09{}") + left.shutdown(socket.SHUT_WR) + with pytest.raises(relay_wire.FrameError): + relay_wire.recv_frame(right) + + +def test_a_frame_over_the_bound_is_not_sent(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(relay_wire, "MAX_FRAME_BYTES", 8) + with pytest.raises(relay_wire.FrameError): + relay_wire._pack({"too": "long for eight"}) + + +def test_values_outside_the_vocabulary_are_reduced_not_dropped() -> None: + import dataclasses + + @dataclasses.dataclass + class Elsewhere: + n: int = 1 + + class Stubborn: + def model_dump(self) -> dict: + raise ValueError("no") + + def __str__(self) -> str: + return "stubborn" + + assert relay_wire.encode({"a": Elsewhere(), "b": Stubborn(), "c": {1, 1}}) == { + "a": {"n": 1}, + "b": "stubborn", + "c": [1], + } + record = relay_wire.decode( + {"__relay__": "CallRecord", "fields": {"http_status": 200, "headers": {"a": 1}, "gone": 1}} + ) + assert record == CallRecord(http_status=200, headers={"a": "1"}) + + +def test_a_timeout_whose_numbers_do_not_parse_still_comes_back_as_a_timeout() -> None: + rebuilt = relay_wire.decode_exception( + {"kind": "first_byte", "timeout": {"budget": "x"}, "message": "slow", "names": []} + ) + assert isinstance(rebuilt, TimeoutError) and str(rebuilt) == "slow" + again = relay_wire.decode_exception({"kind": "timeout", "message": "slow", "names": [], "body": "{}"}) + assert type(again) is type(rebuilt) and again.body == "{}" + + +async def test_the_host_builds_one_provider_per_construction_and_keeps_the_latest( + served, monkeypatch: pytest.MonkeyPatch +) -> None: + script, server = served + assert server.serving + provider = relay_provider.RelayProvider(default_model="custom/m", provider_name="custom") + await provider.chat(messages=[{"role": "user", "content": "q"}]) + await provider.chat(messages=[{"role": "user", "content": "q"}]) + assert len(script["specs"]) == 1 + + monkeypatch.setattr(relay_server, "_PROVIDER_CACHE_SIZE", 1) + other = relay_provider.RelayProvider(default_model="custom/other", provider_name="custom") + await other.chat(messages=[{"role": "user", "content": "q"}]) + assert len(server._providers) == 1 + + +async def _raw(server: relay_server.RelayServer, frame: dict | None) -> dict | None: + from raven.providers.relay import _split + + kind, where = _split(server.env()[relay.RELAY_ENV]) + if kind == "unix": + reader, writer = await asyncio.open_unix_connection(where) + else: + reader, writer = await asyncio.open_connection(*where) + try: + if frame is None: + return None + await relay_wire.write_frame(writer, {"token": server.env()[relay.RELAY_TOKEN_ENV], **frame}) + return await relay_wire.read_frame(reader) + finally: + writer.close() + + +async def test_requests_the_host_cannot_serve_are_answered_with_a_fault(served) -> None: + script, server = served + assert await _raw(server, None) is None + assert "unknown relay op" in (await _raw(server, {"op": "nope"}))["fault"] + assert "no provider construction" in (await _raw(server, {"op": "chat", "spec": "x"}))["fault"] + assert "names no function" in (await _raw(server, {"op": "call", "fn": 3}))["fault"] + assert "not a relayed" in (await _raw(server, {"op": "call", "fn": "os:system", "args": [], "kwargs": {}}))["fault"] + + +async def test_a_builder_that_fails_is_a_fault_not_a_dead_server(monkeypatch: pytest.MonkeyPatch) -> None: + def broken(spec): + raise ValueError("cannot build") + + server = relay_server.RelayServer(build=broken) + await server.start() + try: + reply = await _raw(server, {"op": "chat", "spec": {}, "args": {}}) + assert "cannot build" in reply["fault"] + assert server.serving + finally: + await server.close() + + +async def test_a_single_call_that_fails_on_the_host_raises_the_same_class_here(served) -> None: + script, _ = served + + async def refuse(**kwargs): + raise TimeoutError() + + provider = relay_provider.RelayProvider(default_model="custom/m", provider_name="custom") + await provider.chat(messages=[{"role": "user", "content": "q"}]) + built = next(iter(served[1]._providers.values())) + built.chat = refuse + with pytest.raises(TimeoutError): + await provider.chat(messages=[{"role": "user", "content": "q"}]) + + +async def test_a_stream_the_host_ends_oddly_or_drops_is_a_failure_here( + monkeypatch: pytest.MonkeyPatch, tmp_path +) -> None: + """A host speaking out of turn, played by a bare socket server.""" + replies = iter([{"what": 1}, None]) + + async def host(reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None: + await relay_wire.read_frame(reader) + reply = next(replies, None) + if reply is not None: + await relay_wire.write_frame(writer, reply) + writer.close() + + path = tmp_path / "odd.sock" + server = await asyncio.start_unix_server(host, path=str(path)) + monkeypatch.setenv(relay.RELAY_ENV, f"unix:{path}") + monkeypatch.setenv(relay.RELAY_TOKEN_ENV, "t") + relay._reset_for_tests() + provider = relay_provider.RelayProvider(default_model="custom/m", provider_name="custom") + try: + with pytest.raises(relay.RelayFaultError): + async for _ in provider.chat_stream(messages=[{"role": "user", "content": "q"}]): + pass + with pytest.raises(relay_wire.FrameError): + async for _ in provider.chat_stream(messages=[{"role": "user", "content": "q"}]): + pass + finally: + server.close() + + +async def test_a_stream_without_a_reachable_host_runs_here(monkeypatch: pytest.MonkeyPatch, tmp_path) -> None: + monkeypatch.setenv(relay.RELAY_ENV, f"unix:{tmp_path / 'gone.sock'}") + monkeypatch.setenv(relay.RELAY_TOKEN_ENV, "t") + relay._reset_for_tests() + provider = relay_provider.RelayProvider(default_model="custom/m", provider_name="custom") + local = _Wire({}, {}) + monkeypatch.setattr(provider, "_locally", lambda: local) + + assert [d.content async for d in provider.chat_stream(messages=[{"role": "user", "content": "q"}])] == [ + "a", + "b", + None, + ] + assert not relay.active() + + +def test_the_provider_built_here_follows_this_ones_settings_and_answers_for_the_wire() -> None: + provider = relay_provider.RelayProvider( + default_model="openrouter/openai/gpt-4o", provider_name="openrouter", api_key="k" + ) + provider.model_overrides = {"gpt": {"top_p": 0.1}} + local = provider._locally() + assert local.model_overrides == {"gpt": {"top_p": 0.1}} + assert provider._locally() is local + provider.disable_auto_cache_control = True + assert provider._locally().disable_auto_cache_control is True + model = "openrouter/openai/gpt-4o" + assert provider.wire_model_id(model) == local.wire_model_id(model) + assert provider.can_serve(model) == local.can_serve(model) + assert provider.emits_unparsed_reasoning() == local.emits_unparsed_reasoning() + + +def test_without_an_offer_nothing_is_asked(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv(relay.RELAY_ENV, "unix:/x") + relay._reset_for_tests() + assert relay.target() is None + with pytest.raises(relay.RelayUnavailableError): + relay.request_frame("call") + with pytest.raises(relay.RelayUnavailableError): + relay._connect_sync(1.0) + with pytest.raises(relay.RelayUnavailableError): + asyncio.run(relay.connect()) + with pytest.raises(relay_wire.FrameError): + relay.read_answer(None) + + +async def test_a_host_that_hangs_up_mid_question_leaves_the_relay_up(served, monkeypatch: pytest.MonkeyPatch) -> None: + def hang_up(sock): + raise relay_wire.FrameError("closed") + + monkeypatch.setattr(relay_wire, "recv_frame", hang_up) + with pytest.raises(relay.RelayUnavailableError) as caught: + await asyncio.to_thread(relay.ask, "call", fn="x") + assert not caught.value.connecting + + +async def test_unhashable_arguments_are_asked_every_time_and_shaped(served, monkeypatch: pytest.MonkeyPatch) -> None: + from raven.providers import rates + + asked: list[Any] = [] + + def answer(op: str, **fields: Any) -> Any: + asked.append(fields) + return [0.5, 1.5] + + monkeypatch.setattr(relay, "ask", answer) + monkeypatch.setattr(relay, "_ANSWER_CACHE_SIZE", 1) + assert await asyncio.to_thread(rates._try_litellm_rates, "m1", 1, 1) == (0.5, 1.5) + assert await asyncio.to_thread(rates._try_litellm_rates, "m2", 1, 1) == (0.5, 1.5) + assert await asyncio.to_thread(rates._try_litellm_rates, "m1", 1, 1) == (0.5, 1.5) + assert len(asked) == 3, "the answer cache holds one entry, so m1 was asked again" + + wrapped = relay.relayed()(lambda value: value) + assert await asyncio.to_thread(wrapped, ["a list"]) == [0.5, 1.5] + + +async def test_the_offer_starts_the_relay_once_and_shuts_it(monkeypatch: pytest.MonkeyPatch) -> None: + warmed: list[int] = [] + monkeypatch.setitem(sys.modules, "litellm", None) + monkeypatch.delitem(sys.modules, "litellm") + from raven.providers import litellm_setup + + monkeypatch.setattr(litellm_setup, "warm_up_in_background", lambda: warmed.append(1)) + try: + first = await relay_server.relay_env() + second = await relay_server.relay_env() + assert first == second and first[relay.RELAY_ENV] + assert warmed == [1] + finally: + await relay_server.close_relay() + assert relay_server._SERVER is None + + +def test_a_new_event_loop_gets_a_relay_that_answers(monkeypatch: pytest.MonkeyPatch) -> None: + """A relay left by an ended loop accepts connections nobody answers; the next loop must not hand it out.""" + from raven.providers import litellm_setup + + monkeypatch.setattr(litellm_setup, "warm_up_in_background", lambda: None) + + async def offer() -> dict[str, str]: + return await relay_server.relay_env() + + first = asyncio.run(offer()) + + async def ask_the_second() -> tuple[dict[str, str], Any]: + env = await relay_server.relay_env() + monkeypatch.setenv(relay.RELAY_ENV, env[relay.RELAY_ENV]) + monkeypatch.setenv(relay.RELAY_TOKEN_ENV, env[relay.RELAY_TOKEN_ENV]) + relay._reset_for_tests() + answer = await asyncio.to_thread( + relay.ask, "call", fn="raven.providers.rates:_table_entry", args=["gpt-4o"], kwargs={} + ) + await relay_server.close_relay() + return env, answer + + second, answer = asyncio.run(ask_the_second()) + assert second[relay.RELAY_ENV] != first[relay.RELAY_ENV] + assert answer == relay._RELAYED["raven.providers.rates:_table_entry"]("gpt-4o") + if first[relay.RELAY_ENV].startswith("unix:"): + assert not os.path.exists(first[relay.RELAY_ENV].removeprefix("unix:")) diff --git a/tests/test_provider_resolution_invariants.py b/tests/test_provider_resolution_invariants.py index 3954f2627..c3ff3819a 100644 --- a/tests/test_provider_resolution_invariants.py +++ b/tests/test_provider_resolution_invariants.py @@ -353,7 +353,7 @@ def test_the_wire_form_of_a_model_id_is_built_in_one_module() -> None: owner, # Decomposition, not construction: asks which upstream vendor a gateway # id names, to look up what that vendor's model can do. - "raven/providers/litellm_provider.py", + "raven/providers/litellm_shape.py", # Strips a known prefix off a model id to recover the vendor's own id for # a connectivity probe. Also decomposition. "raven/cli/onboard_commands.py", @@ -567,13 +567,21 @@ def test_every_provider_construction_site_passes_the_users_model_overrides() -> per-endpoint builder is exempt: it inherits the fallback's overrides rather than reading config itself. """ - exempt = {"raven/providers/per_model_provider.py"} + exempt = { + "raven/providers/per_model_provider.py", + # The LLM relay forwards a construction made at one of the sites checked + # here, overrides included, as a whole: to the host, or back into the + # direct class when the host cannot be reached. + "raven/providers/relay_provider.py", + "raven/providers/relay_server.py", + } + builders = {"LiteLLMProvider", "RelayProvider", "litellm_provider"} offenders = [] for path in _production_files(): if _rel(path) in exempt: continue source = path.read_text() - if "LiteLLMProvider(" not in source: + if not any(f"{name}(" in source for name in builders): continue tree = ast.parse(source) for node in ast.walk(tree): @@ -581,7 +589,7 @@ def test_every_provider_construction_site_passes_the_users_model_overrides() -> continue func = node.func name = func.attr if isinstance(func, ast.Attribute) else getattr(func, "id", None) - if name != "LiteLLMProvider": + if name not in builders: continue if not any(kw.arg == "model_overrides" for kw in node.keywords): offenders.append(f"{_rel(path)}:{node.lineno}") From 60c30e04f8d265d80241cc42037f3a4a632227a3 Mon Sep 17 00:00:00 2001 From: Dizhan Xue Date: Tue, 27 Oct 2026 14:04:32 +0000 Subject: [PATCH 02/19] fix(providers): return the relayed wrapper through cast so ty accepts it ty does not honour a mypy-style type: ignore on the relayed decorator's return, and the lint job failed on it with invalid-return-type. The wrapper stands in for the decorated function by construction, so it is cast to the decorated type instead of silenced. --- raven/providers/relay.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/raven/providers/relay.py b/raven/providers/relay.py index f42eecfe0..5afa6d181 100644 --- a/raven/providers/relay.py +++ b/raven/providers/relay.py @@ -34,7 +34,7 @@ import threading from collections import OrderedDict from collections.abc import Callable -from typing import Any, TypeVar +from typing import Any, TypeVar, cast from loguru import logger @@ -255,7 +255,7 @@ def call(*args: Any, **kwargs: Any) -> Any: answers.popitem(last=False) return value - return call # type: ignore[return-value] + return cast(F, call) return wrap From be8738b3484fa9c03d5d642e0a6e77c8355a96bf Mon Sep 17 00:00:00 2001 From: Dizhan Xue Date: Tue, 27 Oct 2026 14:04:39 +0000 Subject: [PATCH 03/19] fix(providers): treat a unix relay offer as unreachable without unix sockets On Windows asyncio has no open_unix_connection and socket has no AF_UNIX, so a unix: offer raised AttributeError past the relay's fallback instead of leaving the call to this process. Both connect paths now answer that case as an unreachable host, and the relay server picks loopback TCP by the same capability rather than by the platform name. The relay tests run the unix transport only where it exists. --- raven/providers/relay.py | 16 ++++++++++++ raven/providers/relay_server.py | 7 +++--- tests/test_provider_relay.py | 44 ++++++++++++++++++++++++++++++++- 3 files changed, 63 insertions(+), 4 deletions(-) diff --git a/raven/providers/relay.py b/raven/providers/relay.py index 5afa6d181..699f91438 100644 --- a/raven/providers/relay.py +++ b/raven/providers/relay.py @@ -145,11 +145,25 @@ def request_frame(op: str, **fields: Any) -> dict[str, Any]: return {"token": found[1], "op": op, **fields} +def _unreachable_kind(kind: str) -> RelayUnavailableError | None: + """Why an offer of this kind cannot be reached from this system, or ``None``. + + A system without Unix sockets (Windows) cannot connect to a ``unix:`` offer, + which only a host on another system could have made; that is a relay this + process does not have, not a failure of the call. + """ + if kind == "unix" and not (hasattr(socket, "AF_UNIX") and hasattr(asyncio, "open_unix_connection")): + return RelayUnavailableError("this system has no Unix sockets for a unix: relay", connecting=True) + return None + + def _connect_sync(timeout: float) -> socket.socket: found = target() if found is None: raise RelayUnavailableError("no relay was offered to this process") kind, where = _split(found[0]) + if (unreachable := _unreachable_kind(kind)) is not None: + raise unreachable try: if kind == "unix": sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) @@ -168,6 +182,8 @@ async def connect() -> tuple[asyncio.StreamReader, asyncio.StreamWriter]: if found is None: raise RelayUnavailableError("no relay was offered to this process") kind, where = _split(found[0]) + if (unreachable := _unreachable_kind(kind)) is not None: + raise unreachable limit = relay_wire.MAX_FRAME_BYTES try: if kind == "unix": diff --git a/raven/providers/relay_server.py b/raven/providers/relay_server.py index 014d67ce2..fe4440d7a 100644 --- a/raven/providers/relay_server.py +++ b/raven/providers/relay_server.py @@ -2,8 +2,9 @@ One server per host process, started the first time a Raven agent is launched with the relay on, listening where only this user can reach it -- a socket in a -fresh ``0700`` directory on POSIX, loopback TCP on Windows -- and answering -only requests carrying the token it handed out in the launch environment. +fresh ``0700`` directory, or loopback TCP where the system has no Unix sockets +(Windows) -- and answering only requests carrying the token it handed out in the +launch environment. Each request is one connection: open, one frame in, the answer out, close. A stream answers with a frame per delta and a closing frame, and a connection the @@ -68,7 +69,7 @@ class RelayServer: def __init__(self, *, build: ProviderBuilder = build_litellm_provider, tcp: bool | None = None) -> None: self._build = build # Loopback TCP where there is no Unix socket to put in a private directory. - self._tcp = sys.platform == "win32" if tcp is None else tcp + self._tcp = not hasattr(asyncio, "start_unix_server") if tcp is None else tcp self._token = secrets.token_urlsafe(32) self._server: asyncio.AbstractServer | None = None self._address = "" diff --git a/tests/test_provider_relay.py b/tests/test_provider_relay.py index 7b2b8aca8..4fc3df25c 100644 --- a/tests/test_provider_relay.py +++ b/tests/test_provider_relay.py @@ -182,7 +182,12 @@ async def chat_stream(self, **kwargs: Any) -> AsyncIterator[ChatDelta]: raise -@pytest.fixture(params=["unix", "tcp"]) +# TCP is the transport a Windows host uses, and the only one a system without +# Unix sockets can serve. +_TRANSPORTS = ["unix", "tcp"] if hasattr(asyncio, "start_unix_server") else ["tcp"] + + +@pytest.fixture(params=_TRANSPORTS) async def served(request, monkeypatch: pytest.MonkeyPatch): """A relay with a scripted provider, on each transport: TCP is the one Windows hosts use.""" script: dict[str, Any] = {"specs": []} @@ -624,6 +629,7 @@ async def refuse(**kwargs): await provider.chat(messages=[{"role": "user", "content": "q"}]) +@pytest.mark.skipif(not hasattr(asyncio, "start_unix_server"), reason="the bare host listens on a Unix socket") async def test_a_stream_the_host_ends_oddly_or_drops_is_a_failure_here( monkeypatch: pytest.MonkeyPatch, tmp_path ) -> None: @@ -774,3 +780,39 @@ async def ask_the_second() -> tuple[dict[str, str], Any]: assert answer == relay._RELAYED["raven.providers.rates:_table_entry"]("gpt-4o") if first[relay.RELAY_ENV].startswith("unix:"): assert not os.path.exists(first[relay.RELAY_ENV].removeprefix("unix:")) + + +def test_a_host_without_unix_sockets_serves_on_loopback_tcp(monkeypatch: pytest.MonkeyPatch) -> None: + assert relay_server.RelayServer()._tcp is not hasattr(asyncio, "start_unix_server") + monkeypatch.delattr(asyncio, "start_unix_server", raising=False) + assert relay_server.RelayServer()._tcp is True + + +async def test_a_unix_address_on_a_system_without_unix_sockets_runs_here( + monkeypatch: pytest.MonkeyPatch, tmp_path +) -> None: + """Windows has no Unix-socket API: a ``unix:`` offer there is a relay this process cannot reach.""" + import socket as socket_mod + + from raven.providers import rates + + monkeypatch.setenv(relay.RELAY_ENV, f"unix:{tmp_path / 'relay.sock'}") + monkeypatch.setenv(relay.RELAY_TOKEN_ENV, "t") + monkeypatch.delattr(asyncio, "open_unix_connection", raising=False) + monkeypatch.delattr(socket_mod, "AF_UNIX", raising=False) + relay._reset_for_tests() + local = relay._RELAYED["raven.providers.rates:_may_query_endpoint"]("gpt-4o") + + assert rates._may_query_endpoint("gpt-4o") == local + assert not relay.active() + + # The first read took the offer out of the environment. + monkeypatch.setenv(relay.RELAY_ENV, f"unix:{tmp_path / 'relay.sock'}") + monkeypatch.setenv(relay.RELAY_TOKEN_ENV, "t") + relay._reset_for_tests() + assert relay.active() + provider = relay_provider.RelayProvider(default_model="custom/m", provider_name="custom") + here = _Wire({}, {}) + monkeypatch.setattr(provider, "_locally", lambda: here) + assert (await provider.chat(messages=[{"role": "user", "content": "q"}])).content == "answered q" + assert not relay.active() From e91a788758f0015b6a54dccfba00436366e9a03b Mon Sep 17 00:00:00 2001 From: Dizhan Xue Date: Tue, 27 Oct 2026 14:04:54 +0000 Subject: [PATCH 04/19] fix(*): start the fork template through the entry, beside the cli The template process loads, by name, the modules agents import -- the CLI, the ACP server, the RPC stack, the agent loop -- and runs python -m raven, which the acp_client shelf may not do. import-linter could not see it because the edges were strings. That side moves to raven/cli/_fork_template.py, a launcher, and the host starts it through Raven's entry point (python -m raven --fork-template), which hands it the arguments. acp_client names no module of it and imports nothing of it; a test holds acp_client to no importlib or runpy loading. Going through the entry also settles the certificate store before the template imports anything. The template used to preload aiohttp, which builds its default TLS context on import, before any switch, so a forked agent kept a context that did not trust the system store. The in-process test that runs python -m raven now opts out of that switch, which had leaked truststore into the test process. The template-process tests move to tests/test_cli_fork_template.py. --- CONTEXT.md | 21 +- raven/acp_client/fork_template.py | 404 ++--------------------------- raven/cli/_fork_template.py | 405 ++++++++++++++++++++++++++++++ raven/cli/entry.py | 12 +- tests/test_acp_fork_template.py | 285 +++------------------ tests/test_cli_fork_template.py | 290 +++++++++++++++++++++ 6 files changed, 774 insertions(+), 643 deletions(-) create mode 100644 raven/cli/_fork_template.py create mode 100644 tests/test_cli_fork_template.py diff --git a/CONTEXT.md b/CONTEXT.md index ff453caa6..efa28cfc1 100644 --- a/CONTEXT.md +++ b/CONTEXT.md @@ -3409,15 +3409,20 @@ attempt and the `relayed` catalogue lookups crossing to the host. Retry, classif records and traces stay in the agent; the host's calls run under `recorded_by_caller` so a call is billed once. An agent that cannot reach the host loads LiteLLM and works as before. -**Fork template** (`raven/acp_client/fork_template.py`, `subagents.runtime.forkTemplate`): +**Fork template** (`raven/acp_client/fork_template.py`, `raven/cli/_fork_template.py`, `subagents.runtime.forkTemplate`): On Linux, one process per host that has imported what Raven agents import, frozen out of the -collector's reach, and forks each Raven agent the host launches, so the agents share those -pages. It starts through the parent-death wrapper like any agent; each forked agent leads a -session of its own and arms the parent-death signal against the template. A product launcher's -`os.execv` of `python -m raven` runs Raven in the same process. Agents write the modules they -imported beside the template, per product, and the next template loads the modules at least two -products import (one product's own while only one has been seen): a module one product alone imports -is shared by no other process, so it is left for that product to import after the fork. Only the host's own interpreter starting +collector's reach, and forks each Raven agent the host launches, so the agents share those pages. +Two halves: the host's side in `acp_client` decides when to ask, starts and ends the template and +holds the endpoint a forked agent runs behind; the template's own side is a launcher beside the +CLI, because it loads the surfaces by name and runs `python -m raven`, which the client shelf may +not. The client starts it through Raven's entry point, which settles the certificate store first, +and never imports it. It starts through the parent-death wrapper like any agent and its stdin is a +lifeline only the host holds. Each forked agent leads a session of its own and arms the +parent-death signal against the template. A product launcher's `os.execv` of `python -m raven` +runs Raven in the same process. Agents write the modules they imported beside the template, per +product, and the next template loads the modules at least two products import (one product's own +while only one has been seen): a module one product alone imports is shared by no other process, +so it is left for that product to import after the fork. Only the host's own interpreter starting `-m raven` or a product launcher is forked; anything else launches the ordinary way. **Leftover agent** (`raven/acp_client/leftovers.py`): diff --git a/raven/acp_client/fork_template.py b/raven/acp_client/fork_template.py index 72bab171a..121b0bdb4 100644 --- a/raven/acp_client/fork_template.py +++ b/raven/acp_client/fork_template.py @@ -8,25 +8,20 @@ launches; measured on the product agents, what each one costs on its own falls by about two thirds. -**The template.** ``python -m raven.acp_client.fork_template``, started by the -host through the same parent-death wrapper as any agent (``leftovers``), so it -ends with the host. It imports a seed of the modules every agent needs plus -the modules agents were seen to import before -- each agent writes its module -list beside the template after a while, under its product, and the next -template loads what at least two products share (one product's own list while -only one has been seen; see :func:`preload_names`) -- then collects, freezes the heap out of the collector's reach -- a collection writes -to every object it visits, which would copy the shared pages back apart -- and -then answers fork requests on a socket in a private directory. It runs no -thread, and refuses to fork if an import started one. +This module is the host's side: when to ask, starting and ending the template, +and the endpoint a forked agent runs behind. The template's own side -- what it +loads, the fork and the agent it becomes -- is a launcher and lives beside the +CLI (``raven.cli._fork_template``). This shelf starts it through Raven's entry +point, as ``python -m raven --fork-template``, and never imports it. + +**The template** is started through the same parent-death wrapper as any agent +(``leftovers``), so it ends with the host, and its stdin is a lifeline only the +host holds: letting go of it asks the template to stop. **A fork.** The host hands over the agent's argv, environment and working directory, and the three pipe ends that become its stdin, stdout and stderr. -The child starts a session of its own (so the host ends it with ``killpg`` as -it would any agent), arms the parent-death signal against the template, and -runs the agent's program in place of the ``exec`` a fresh interpreter would -have done: a product launcher's own ``os.execv`` of ``python -m raven`` is -answered the same way, by running Raven in the same process. The template -reaps the child and tells the host its exit status over the request's +The child leads a session of its own, so the host ends it with ``killpg`` as it +would any agent, and the template reports its exit status over the request's connection, which stays open for that. **What is eligible.** Only an agent this interpreter would run as Raven: the @@ -37,40 +32,29 @@ from __future__ import annotations -import argparse import asyncio -import collections import contextlib -import gc -import importlib import json import os import signal import socket import sys import tempfile -import threading -import time from collections.abc import Sequence from pathlib import Path from typing import Any -#: Modules every Raven agent imports before it serves, imported by every -#: template whatever it has learned. -SEED_MODULES = ( - "raven.config.product_render", - "raven.cli.commands", - "raven.acp.server", - "raven.acp.methods", - "raven.rpc.bootstrap", - "raven.agent.loop", -) - #: How long after starting an agent writes the modules it has imported. Long #: enough to cover its first session and turn, where nearly all of them load. LEARN_AFTER_S = 30.0 _LEARNED_FILE = "modules.json" -_LEARNED_CAP = 20_000 + +#: How the template process is started: through Raven's entry point, which settles +#: the certificate store before anything is imported and hands this argument to the +#: launcher beside the CLI. That launcher imports the surfaces agents import, by +#: name, which this shelf may not. +TEMPLATE_ARG = "--fork-template" +_TEMPLATE_PROGRAM: tuple[str, ...] = ("-m", "raven", TEMPLATE_ARG) _READY_TIMEOUT_S = 60.0 _REQUEST_TIMEOUT_S = 30.0 @@ -245,7 +229,7 @@ async def _start(self, env: dict[str, str]) -> None: self._dir = Path(tempfile.mkdtemp(prefix="raven-fork-template-")) self._socket = str(self._dir / "template.sock") learned = get_runtime_subdir("acp-template") / _LEARNED_FILE - argv = [sys.executable, "-P", "-m", __name__, "--socket", self._socket, "--learned", str(learned)] + argv = [sys.executable, "-P", *_TEMPLATE_PROGRAM, "--socket", self._socket, "--learned", str(learned)] argv += ["--learn-after", str(LEARN_AFTER_S)] for prefix in self._skip: argv += ["--skip", prefix] @@ -491,353 +475,3 @@ async def close_template() -> None: template.abandon() elif template is not None: await template.close() - - -# ---- the template's side ------------------------------------------------------- - - -#: Not loaded ahead whatever was learned: the entry point itself, and this -#: module under its import name (it already runs as ``__main__``). -_NEVER_LOAD = frozenset({"__main__", "raven.acp_client.fork_template"}) - - -def product_of(argv: Sequence[str]) -> str: - """The name an agent's modules are learned under: its product folder, or ``raven`` for ``-m raven``.""" - if len(argv) > 1 and argv[1] != "-m": - return Path(argv[1]).parent.name or "raven" - return "raven" - - -def preload_names(learned: Any) -> list[str]: - """What a template loads ahead of the agents: the modules they share, or one product's own. - - A module only one product imports is shared by nobody -- a host normally - runs one process per product -- so loading it ahead would hold it in the - template, unused, while any other product runs. Measured on the five - shipped products: they share about two thousand modules, each adds a few - dozen of its own, and the deck engine adds about a thousand. With only one - product learned yet, its own set is all there is to share. - """ - if not isinstance(learned, dict) or learned.get("python") != list(sys.version_info[:2]): - return [] - products = learned.get("products") - if not isinstance(products, dict): - return [] - sets = [{n for n in names if isinstance(n, str)} for names in products.values() if isinstance(names, list)] - if len(sets) == 1: - return sorted(sets[0]) - counts = collections.Counter(name for names in sets for name in names) - return sorted(name for name, seen in counts.items() if seen >= 2) - - -def _load(seed: Sequence[str], learned: Path, skip: Sequence[str]) -> int: - """Import what the agents will import; the module count afterwards.""" - names = list(seed) - try: - names += preload_names(json.loads(learned.read_text(encoding="utf-8"))) - except (OSError, ValueError): - pass - for name in names: - if name in sys.modules or name in _NEVER_LOAD or name.endswith(".__main__") or name.startswith(tuple(skip)): - continue - try: - importlib.import_module(name) - except BaseException: # noqa: BLE001 - an import that fails here fails the same way in the agent - continue - # Not an import: every agent loads the token estimator's table on its first - # count, and the table's core is read-only once built, so loaded here it is - # shared by every agent for as long as they run. - with contextlib.suppress(Exception): - from raven.utils.tokens import encoding - - encoding() - return len(sys.modules) - - -def _learn(path: Path, skip: Sequence[str], product: str) -> None: - """Add the modules this agent imported to what its product is known to import.""" - try: - import fcntl - - mine = { - name - for name, module in list(sys.modules.items()) - if getattr(module, "__spec__", None) is not None - and name != "__main__" - and not name.endswith(".__main__") - and not name.startswith(tuple(skip)) - } - path.parent.mkdir(parents=True, exist_ok=True) - with open(path.with_suffix(".lock"), "a") as lock: - fcntl.flock(lock, fcntl.LOCK_EX) - try: - data = json.loads(path.read_text(encoding="utf-8")) - except (OSError, ValueError): - data = {} - products = data.get("products") if isinstance(data, dict) else None - if not isinstance(products, dict) or data.get("python") != list(sys.version_info[:2]): - products = {} - known = products.get(product) - known = {n for n in known if isinstance(n, str)} if isinstance(known, list) else set() - merged = sorted(known | mine)[:_LEARNED_CAP] - if merged != sorted(known): - products[product] = merged - tmp = path.with_suffix(f".{os.getpid()}.tmp") - tmp.write_text( - json.dumps({"python": list(sys.version_info[:2]), "products": products}), encoding="utf-8" - ) - os.replace(tmp, path) - except Exception: # noqa: BLE001 - learning is an optimisation for the next template - pass - - -def _arm_parent_death( - parent: int, -) -> None: # pragma: no cover - runs only in a forked agent; see tests/test_acp_fork_template.py - """End this process when the template does; it ends with the host the same way.""" - try: - import ctypes - - armed = ctypes.CDLL(None).prctl(1, signal.SIGKILL, 0, 0, 0) == 0 - except Exception: # noqa: BLE001 - without the signal the agent still ends on its closed stdin - armed = False - if armed and os.getppid() != parent: - os._exit(1) - - -def _finish(code: int) -> None: # pragma: no cover - runs only in a forked agent; see tests/test_acp_fork_template.py - """End the agent the way an interpreter ends: threads joined, exit handlers run, streams flushed. - - Never by unwinding: the frames below the agent's program are the template's, - and a launcher's own exit handling must not run after Raven, which it would - never have reached past a real ``exec``. - """ - import atexit - - with contextlib.suppress(Exception): - threading._shutdown() # type: ignore[attr-defined] - with contextlib.suppress(Exception): - atexit._run_exitfuncs() - for stream in (sys.stdout, sys.stderr): - with contextlib.suppress(Exception): - stream.flush() - os._exit(code) - - -def _exit_code(exc: BaseException) -> int: - """What an interpreter exits with when ``exc`` ends its program, printing what it would print.""" - if isinstance(exc, SystemExit): - if exc.code is None: - return 0 - if isinstance(exc.code, int): - return exc.code - with contextlib.suppress(Exception): - print(exc.code, file=sys.stderr) - return 1 - with contextlib.suppress(Exception): - sys.excepthook(type(exc), exc, exc.__traceback__) - return 1 - - -def _run_as_raven(args: Sequence[str]) -> None: - """``python -m raven `` in this process, then the end of it.""" - import runpy - - sys.argv = ["raven", *args] - sys.path.insert(0, os.getcwd()) - try: - runpy.run_module("raven", run_name="__main__", alter_sys=True) - code = 0 - except BaseException as exc: # noqa: BLE001 - the program's end, however it came - code = _exit_code(exc) - _finish(code) - - -def _run_agent(argv: list[str]) -> None: - """Run the agent's program here, as a fresh interpreter would have run it, and end the process.""" - import runpy - - real_execv = os.execv - - def execv(path: Any, args: Any) -> None: - args = [str(a) for a in args] - if os.path.abspath(str(path)) == os.path.abspath(sys.executable) and args[1:3] == ["-m", "raven"]: - os.execv = real_execv - _run_as_raven(args[3:]) - real_execv(path, args) - - os.execv = execv - if argv[1] == "-m": - _run_as_raven(argv[3:]) - script = os.path.abspath(argv[1]) - sys.argv = [argv[1], *argv[2:]] - sys.path.insert(0, os.path.dirname(script)) - try: - runpy.run_path(script, run_name="__main__") - code = 0 - except BaseException as exc: # noqa: BLE001 - the program's end, however it came - code = _exit_code(exc) - _finish(code) - - -def _become( # pragma: no cover - runs only in a forked agent; see tests/test_acp_fork_template.py - request: dict[str, Any], fds: list[int], template: int, learned: Path, skip: Sequence[str], learn_after: float -) -> None: - """Turn this freshly forked child into the agent the request names. Does not return.""" - os.setsid() - _arm_parent_death(template) - for signum in (signal.SIGCHLD, signal.SIGTERM): - signal.signal(signum, signal.SIG_DFL) - with contextlib.suppress(ValueError): - signal.set_wakeup_fd(-1) - for target, fd in zip((0, 1, 2), fds, strict=True): - os.dup2(fd, target) - os.close(fd) - os.chdir(request["cwd"]) - os.environ.clear() - os.environ.update({str(k): str(v) for k, v in (request.get("env") or {}).items()}) - timer = threading.Timer(learn_after, _learn, args=(learned, skip, product_of(request["argv"]))) - timer.daemon = True - timer.start() - _run_agent([str(a) for a in request["argv"]]) - - -def _serve(path: str, *, lifeline: int) -> tuple[dict[str, Any], list[int]] | None: - """Fork agents on request until the host lets go of ``lifeline``. - - In a forked child this returns the request it is to become; in the template - it returns ``None`` once ``lifeline`` -- the template's stdin, whose other end - only the host holds -- reads as closed. - """ - import selectors - - from raven.providers import relay_wire - - template = os.getpid() - listener = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) - listener.bind(path) - listener.listen(64) - wake_r, wake_w = os.pipe() - os.set_blocking(wake_w, False) - signal.set_wakeup_fd(wake_w) - signal.signal(signal.SIGCHLD, lambda *_: None) - selector = selectors.DefaultSelector() - selector.register(listener, selectors.EVENT_READ) - selector.register(wake_r, selectors.EVENT_READ) - selector.register(lifeline, selectors.EVENT_READ) - reporting: dict[int, socket.socket] = {} - - def let_go(*, child: bool) -> None: - selector.close() - listener.close() - for conn in reporting.values(): - conn.close() - if not child: - signal.set_wakeup_fd(-1) - signal.signal(signal.SIGCHLD, signal.SIG_DFL) - os.close(wake_r) - os.close(wake_w) - - print( - json.dumps( - {"ready": True, "pid": template, "modules": len(sys.modules), "seconds": time.monotonic() - _STARTED} - ), - flush=True, - ) - while True: - for key, _ in selector.select(): - if key.fileobj == lifeline: - if not os.read(lifeline, 4096): - let_go(child=False) - return None - continue - if key.fileobj == wake_r: - with contextlib.suppress(BlockingIOError): - os.read(wake_r, 4096) - _report_exits(reporting) - continue - conn, _ = listener.accept() - conn.settimeout(_REQUEST_TIMEOUT_S) - fds: list[int] = [] - try: - _, fds, _, _ = socket.recv_fds(conn, 1, 3) - request = relay_wire.recv_frame(conn) - if len(fds) != 3 or not isinstance(request.get("argv"), list) or not request.get("cwd"): - raise ValueError("a fork request carries an argv, a cwd and three descriptors") - if threading.active_count() != 1: - raise RuntimeError("the template is running a thread and cannot fork safely") - pid = os.fork() - except Exception as exc: # noqa: BLE001 - one bad request must not end the template - for fd in fds: - os.close(fd) - with contextlib.suppress(OSError): - relay_wire.send_frame(conn, {"error": f"{type(exc).__name__}: {exc}"}) - conn.close() - continue - if pid == 0: # pragma: no cover - the forked agent's side of the fork - conn.close() - let_go(child=True) - return request, fds - for fd in fds: - os.close(fd) - try: - relay_wire.send_frame(conn, {"pid": pid}) - reporting[pid] = conn - except OSError: - conn.close() - - -def _report_exits(reporting: dict[int, socket.socket]) -> None: - """Reap every child that has ended and tell its requester how.""" - from raven.providers import relay_wire - - while True: - try: - pid, status = os.waitpid(-1, os.WNOHANG) - except ChildProcessError: - return - if pid == 0: - return - conn = reporting.pop(pid, None) - if conn is not None: - with contextlib.suppress(OSError): - relay_wire.send_frame(conn, {"exit": os.waitstatus_to_exitcode(status)}) - conn.close() - - -_STARTED = time.monotonic() - - -def main(argv: Sequence[str] | None = None) -> int: - parser = argparse.ArgumentParser(prog="raven-fork-template") - parser.add_argument("--socket", required=True) - parser.add_argument("--learned", required=True) - parser.add_argument("--skip", action="append", default=[]) - parser.add_argument("--learn-after", type=float, default=LEARN_AFTER_S) - args = parser.parse_args(argv) - learned = Path(args.learned) - # stdout is the channel the host reads the ready line from; whatever an - # import prints goes to stderr instead, where the host only logs it. - sys.stdout.flush() - ready_channel = os.dup(1) - os.dup2(2, 1) - _load(SEED_MODULES, learned, args.skip) - sys.stdout.flush() - os.dup2(ready_channel, 1) - os.close(ready_channel) - gc.collect() - gc.freeze() - if threading.active_count() != 1: - names = ", ".join(t.name for t in threading.enumerate() if t is not threading.current_thread()) - print(json.dumps({"ready": False, "reason": f"an import started a thread ({names})"}), flush=True) - return 3 - forked = _serve(args.socket, lifeline=0) - if forked is None: - return 0 - request, fds = forked - _become(request, fds, os.getppid(), learned, args.skip, args.learn_after) - return 1 - - -if __name__ == "__main__": - sys.exit(main()) diff --git a/raven/cli/_fork_template.py b/raven/cli/_fork_template.py new file mode 100644 index 000000000..9496d4724 --- /dev/null +++ b/raven/cli/_fork_template.py @@ -0,0 +1,405 @@ +"""The fork template's own process: load what Raven agents load, then fork them on request. + +Started by the host as ``python -P -m raven --fork-template`` (see +``raven.acp_client.fork_template``, the host's side): through Raven's entry +point like every Raven process, so the certificate store is settled before it +loads anything, and the TLS contexts its imports build are the ones its agents +would have built. It is a launcher: it +imports, by name, the modules a Raven agent imports -- the CLI, the ACP server, +the RPC stack, the agent loop -- and runs ``python -m raven`` in the agents it +forks, which is why it sits beside the CLI rather than with the client that +starts it. + +**Loading.** A seed of the modules every agent needs plus the modules agents +were seen to import before: each agent writes its module list beside the +template after a while, under its product, and the next template loads what at +least two products share (one product's own list while only one has been seen; +see :func:`preload_names`). Then it collects and freezes the heap out of the +collector's reach -- a collection writes to every object it visits, which would +copy the shared pages back apart -- and answers fork requests on a socket in a +private directory. It runs no thread, and refuses to fork if an import started +one. + +**A fork.** The child starts a session of its own, arms the parent-death signal +against the template, and runs the agent's program in place of the ``exec`` a +fresh interpreter would have done: a product launcher's own ``os.execv`` of +``python -m raven`` is answered the same way, by running Raven in the same +process. The template reaps the child and tells the host its exit status. +""" + +from __future__ import annotations + +import argparse +import collections +import contextlib +import gc +import importlib +import json +import os +import signal +import socket +import sys +import threading +import time +from collections.abc import Sequence +from pathlib import Path +from typing import Any + +from raven.acp_client.fork_template import LEARN_AFTER_S + +#: Modules every Raven agent imports before it serves, imported by every +#: template whatever it has learned. +SEED_MODULES = ( + "raven.config.product_render", + "raven.cli.commands", + "raven.acp.server", + "raven.acp.methods", + "raven.rpc.bootstrap", + "raven.agent.loop", +) + +_LEARNED_CAP = 20_000 +_REQUEST_TIMEOUT_S = 30.0 + + +#: Not loaded ahead whatever was learned: the program an agent ran as, which is +#: its own and not the template's, and this module, which is already running. +_NEVER_LOAD = frozenset({"__main__", "raven.cli._fork_template"}) + + +def product_of(argv: Sequence[str]) -> str: + """The name an agent's modules are learned under: its product folder, or ``raven`` for ``-m raven``.""" + if len(argv) > 1 and argv[1] != "-m": + return Path(argv[1]).parent.name or "raven" + return "raven" + + +def preload_names(learned: Any) -> list[str]: + """What a template loads ahead of the agents: the modules they share, or one product's own. + + A module only one product imports is shared by nobody -- a host normally + runs one process per product -- so loading it ahead would hold it in the + template, unused, while any other product runs. Measured on the five + shipped products: they share about two thousand modules, each adds a few + dozen of its own, and the deck engine adds about a thousand. With only one + product learned yet, its own set is all there is to share. + """ + if not isinstance(learned, dict) or learned.get("python") != list(sys.version_info[:2]): + return [] + products = learned.get("products") + if not isinstance(products, dict): + return [] + sets = [{n for n in names if isinstance(n, str)} for names in products.values() if isinstance(names, list)] + if len(sets) == 1: + return sorted(sets[0]) + counts = collections.Counter(name for names in sets for name in names) + return sorted(name for name, seen in counts.items() if seen >= 2) + + +def _load(seed: Sequence[str], learned: Path, skip: Sequence[str]) -> int: + """Import what the agents will import; the module count afterwards.""" + names = list(seed) + try: + names += preload_names(json.loads(learned.read_text(encoding="utf-8"))) + except (OSError, ValueError): + pass + for name in names: + if name in sys.modules or name in _NEVER_LOAD or name.endswith(".__main__") or name.startswith(tuple(skip)): + continue + try: + importlib.import_module(name) + except BaseException: # noqa: BLE001 - an import that fails here fails the same way in the agent + continue + # Not an import: every agent loads the token estimator's table on its first + # count, and the table's core is read-only once built, so loaded here it is + # shared by every agent for as long as they run. + with contextlib.suppress(Exception): + from raven.utils.tokens import encoding + + encoding() + return len(sys.modules) + + +def _learn(path: Path, skip: Sequence[str], product: str) -> None: + """Add the modules this agent imported to what its product is known to import.""" + try: + import fcntl + + mine = { + name + for name, module in list(sys.modules.items()) + if getattr(module, "__spec__", None) is not None + and name != "__main__" + and not name.endswith(".__main__") + and not name.startswith(tuple(skip)) + } + path.parent.mkdir(parents=True, exist_ok=True) + with open(path.with_suffix(".lock"), "a") as lock: + fcntl.flock(lock, fcntl.LOCK_EX) + try: + data = json.loads(path.read_text(encoding="utf-8")) + except (OSError, ValueError): + data = {} + products = data.get("products") if isinstance(data, dict) else None + if not isinstance(products, dict) or data.get("python") != list(sys.version_info[:2]): + products = {} + known = products.get(product) + known = {n for n in known if isinstance(n, str)} if isinstance(known, list) else set() + merged = sorted(known | mine)[:_LEARNED_CAP] + if merged != sorted(known): + products[product] = merged + tmp = path.with_suffix(f".{os.getpid()}.tmp") + tmp.write_text( + json.dumps({"python": list(sys.version_info[:2]), "products": products}), encoding="utf-8" + ) + os.replace(tmp, path) + except Exception: # noqa: BLE001 - learning is an optimisation for the next template + pass + + +def _arm_parent_death( + parent: int, +) -> None: # pragma: no cover - runs only in a forked agent; see tests/test_acp_fork_template.py + """End this process when the template does; it ends with the host the same way.""" + try: + import ctypes + + armed = ctypes.CDLL(None).prctl(1, signal.SIGKILL, 0, 0, 0) == 0 + except Exception: # noqa: BLE001 - without the signal the agent still ends on its closed stdin + armed = False + if armed and os.getppid() != parent: + os._exit(1) + + +def _finish(code: int) -> None: # pragma: no cover - runs only in a forked agent; see tests/test_acp_fork_template.py + """End the agent the way an interpreter ends: threads joined, exit handlers run, streams flushed. + + Never by unwinding: the frames below the agent's program are the template's, + and a launcher's own exit handling must not run after Raven, which it would + never have reached past a real ``exec``. + """ + import atexit + + with contextlib.suppress(Exception): + threading._shutdown() # type: ignore[attr-defined] + with contextlib.suppress(Exception): + atexit._run_exitfuncs() + for stream in (sys.stdout, sys.stderr): + with contextlib.suppress(Exception): + stream.flush() + os._exit(code) + + +def _exit_code(exc: BaseException) -> int: + """What an interpreter exits with when ``exc`` ends its program, printing what it would print.""" + if isinstance(exc, SystemExit): + if exc.code is None: + return 0 + if isinstance(exc.code, int): + return exc.code + with contextlib.suppress(Exception): + print(exc.code, file=sys.stderr) + return 1 + with contextlib.suppress(Exception): + sys.excepthook(type(exc), exc, exc.__traceback__) + return 1 + + +def _run_as_raven(args: Sequence[str]) -> None: + """``python -m raven `` in this process, then the end of it.""" + import runpy + + sys.argv = ["raven", *args] + sys.path.insert(0, os.getcwd()) + try: + runpy.run_module("raven", run_name="__main__", alter_sys=True) + code = 0 + except BaseException as exc: # noqa: BLE001 - the program's end, however it came + code = _exit_code(exc) + _finish(code) + + +def _run_agent(argv: list[str]) -> None: + """Run the agent's program here, as a fresh interpreter would have run it, and end the process.""" + import runpy + + real_execv = os.execv + + def execv(path: Any, args: Any) -> None: + args = [str(a) for a in args] + if os.path.abspath(str(path)) == os.path.abspath(sys.executable) and args[1:3] == ["-m", "raven"]: + os.execv = real_execv + _run_as_raven(args[3:]) + real_execv(path, args) + + os.execv = execv + if argv[1] == "-m": + _run_as_raven(argv[3:]) + script = os.path.abspath(argv[1]) + sys.argv = [argv[1], *argv[2:]] + sys.path.insert(0, os.path.dirname(script)) + try: + runpy.run_path(script, run_name="__main__") + code = 0 + except BaseException as exc: # noqa: BLE001 - the program's end, however it came + code = _exit_code(exc) + _finish(code) + + +def _become( # pragma: no cover - runs only in a forked agent; see tests/test_acp_fork_template.py + request: dict[str, Any], fds: list[int], template: int, learned: Path, skip: Sequence[str], learn_after: float +) -> None: + """Turn this freshly forked child into the agent the request names. Does not return.""" + os.setsid() + _arm_parent_death(template) + for signum in (signal.SIGCHLD, signal.SIGTERM): + signal.signal(signum, signal.SIG_DFL) + with contextlib.suppress(ValueError): + signal.set_wakeup_fd(-1) + for target, fd in zip((0, 1, 2), fds, strict=True): + os.dup2(fd, target) + os.close(fd) + os.chdir(request["cwd"]) + os.environ.clear() + os.environ.update({str(k): str(v) for k, v in (request.get("env") or {}).items()}) + timer = threading.Timer(learn_after, _learn, args=(learned, skip, product_of(request["argv"]))) + timer.daemon = True + timer.start() + _run_agent([str(a) for a in request["argv"]]) + + +def _serve(path: str, *, lifeline: int) -> tuple[dict[str, Any], list[int]] | None: + """Fork agents on request until the host lets go of ``lifeline``. + + In a forked child this returns the request it is to become; in the template + it returns ``None`` once ``lifeline`` -- the template's stdin, whose other end + only the host holds -- reads as closed. + """ + import selectors + + from raven.providers import relay_wire + + template = os.getpid() + listener = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) + listener.bind(path) + listener.listen(64) + wake_r, wake_w = os.pipe() + os.set_blocking(wake_w, False) + signal.set_wakeup_fd(wake_w) + signal.signal(signal.SIGCHLD, lambda *_: None) + selector = selectors.DefaultSelector() + selector.register(listener, selectors.EVENT_READ) + selector.register(wake_r, selectors.EVENT_READ) + selector.register(lifeline, selectors.EVENT_READ) + reporting: dict[int, socket.socket] = {} + + def let_go(*, child: bool) -> None: + selector.close() + listener.close() + for conn in reporting.values(): + conn.close() + if not child: + signal.set_wakeup_fd(-1) + signal.signal(signal.SIGCHLD, signal.SIG_DFL) + os.close(wake_r) + os.close(wake_w) + + print( + json.dumps( + {"ready": True, "pid": template, "modules": len(sys.modules), "seconds": time.monotonic() - _STARTED} + ), + flush=True, + ) + while True: + for key, _ in selector.select(): + if key.fileobj == lifeline: + if not os.read(lifeline, 4096): + let_go(child=False) + return None + continue + if key.fileobj == wake_r: + with contextlib.suppress(BlockingIOError): + os.read(wake_r, 4096) + _report_exits(reporting) + continue + conn, _ = listener.accept() + conn.settimeout(_REQUEST_TIMEOUT_S) + fds: list[int] = [] + try: + _, fds, _, _ = socket.recv_fds(conn, 1, 3) + request = relay_wire.recv_frame(conn) + if len(fds) != 3 or not isinstance(request.get("argv"), list) or not request.get("cwd"): + raise ValueError("a fork request carries an argv, a cwd and three descriptors") + if threading.active_count() != 1: + raise RuntimeError("the template is running a thread and cannot fork safely") + pid = os.fork() + except Exception as exc: # noqa: BLE001 - one bad request must not end the template + for fd in fds: + os.close(fd) + with contextlib.suppress(OSError): + relay_wire.send_frame(conn, {"error": f"{type(exc).__name__}: {exc}"}) + conn.close() + continue + if pid == 0: # pragma: no cover - the forked agent's side of the fork + conn.close() + let_go(child=True) + return request, fds + for fd in fds: + os.close(fd) + try: + relay_wire.send_frame(conn, {"pid": pid}) + reporting[pid] = conn + except OSError: + conn.close() + + +def _report_exits(reporting: dict[int, socket.socket]) -> None: + """Reap every child that has ended and tell its requester how.""" + from raven.providers import relay_wire + + while True: + try: + pid, status = os.waitpid(-1, os.WNOHANG) + except ChildProcessError: + return + if pid == 0: + return + conn = reporting.pop(pid, None) + if conn is not None: + with contextlib.suppress(OSError): + relay_wire.send_frame(conn, {"exit": os.waitstatus_to_exitcode(status)}) + conn.close() + + +_STARTED = time.monotonic() + + +def main(argv: Sequence[str] | None = None) -> int: + parser = argparse.ArgumentParser(prog="raven-fork-template") + parser.add_argument("--socket", required=True) + parser.add_argument("--learned", required=True) + parser.add_argument("--skip", action="append", default=[]) + parser.add_argument("--learn-after", type=float, default=LEARN_AFTER_S) + args = parser.parse_args(argv) + learned = Path(args.learned) + # stdout is the channel the host reads the ready line from; whatever an + # import prints goes to stderr instead, where the host only logs it. + sys.stdout.flush() + ready_channel = os.dup(1) + os.dup2(2, 1) + _load(SEED_MODULES, learned, args.skip) + sys.stdout.flush() + os.dup2(ready_channel, 1) + os.close(ready_channel) + gc.collect() + gc.freeze() + if threading.active_count() != 1: + names = ", ".join(t.name for t in threading.enumerate() if t is not threading.current_thread()) + print(json.dumps({"ready": False, "reason": f"an import started a thread ({names})"}), flush=True) + return 3 + forked = _serve(args.socket, lifeline=0) + if forked is None: + return 0 + request, fds = forked + _become(request, fds, os.getppid(), learned, args.skip, args.learn_after) + return 1 diff --git a/raven/cli/entry.py b/raven/cli/entry.py index db552dae0..286555a8e 100644 --- a/raven/cli/entry.py +++ b/raven/cli/entry.py @@ -16,9 +16,10 @@ def run() -> None: """Console-script entry point. - The supervisor `raven web` leaves running is the one launch taken without the - command table: it only restarts the gateway, and whatever it imports stays - resident for as long as the page is up. + Two launches are taken without the command table. The supervisor `raven web` + leaves running only restarts the gateway, and whatever it imports stays + resident for as long as the page is up. The fork template loads what Raven + agents load and forks them, so what it imports is exactly what it chooses. """ from raven.security.tls import use_system_ca @@ -31,6 +32,11 @@ def run() -> None: if supervise_from_argv(sys.argv[1:]): return + if sys.argv[1:2] == ["--fork-template"]: + from raven.cli._fork_template import main as serve_forks + + raise SystemExit(serve_forks(sys.argv[2:])) + from raven.cli.commands import run as run_commands run_commands() diff --git a/tests/test_acp_fork_template.py b/tests/test_acp_fork_template.py index 4b3e2fd63..22ecfb72e 100644 --- a/tests/test_acp_fork_template.py +++ b/tests/test_acp_fork_template.py @@ -271,197 +271,6 @@ async def test_switching_both_off_launches_a_raven_agent_as_before(tmp_path: Pat await client.close() -# ---- the template's side, driven in this process --------------------------------- - - -class _Ended(Exception): - def __init__(self, code: int) -> None: - super().__init__(code) - self.code = code - - -@pytest.fixture -def ended(monkeypatch: pytest.MonkeyPatch) -> list[int]: - """Run the agent-side functions here: the exit becomes an exception, and what they change is put back. - - The codes are recorded because the real exit never returns, while this one - unwinds through the frames above it, which may read it as a second ending. - """ - codes: list[int] = [] - - def finish(code: int) -> None: - codes.append(code) - raise _Ended(code) - - monkeypatch.setattr(fork_template, "_finish", finish) - monkeypatch.setattr(os, "execv", os.execv) - monkeypatch.setattr(sys, "argv", list(sys.argv)) - monkeypatch.setattr(sys, "path", list(sys.path)) - return codes - - -def _learned(**products: list[str]) -> dict: - return {"python": list(sys.version_info[:2]), "products": products} - - -def test_the_template_loads_the_seed_and_what_agents_imported_but_not_what_it_skips( - tmp_path: Path, monkeypatch -) -> None: - for name in ("colorsys", "netrc", "this_module_does_not_exist"): - monkeypatch.delitem(sys.modules, name, raising=False) - learned = tmp_path / "modules.json" - learned.write_text(json.dumps(_learned(only=["netrc", "this_module_does_not_exist"]))) - - fork_template._load(("colorsys",), learned, ()) - assert "colorsys" in sys.modules and "netrc" in sys.modules - - monkeypatch.delitem(sys.modules, "netrc") - fork_template._load((), learned, ("net",)) - assert "netrc" not in sys.modules - - learned.write_text(json.dumps({"python": [2, 7], "products": {"only": ["netrc"]}})) - fork_template._load((), learned, ()) - assert "netrc" not in sys.modules - - -def test_one_products_own_modules_are_loaded_ahead_only_while_it_is_the_only_one_seen() -> None: - """What only one product imports is shared by no other process, so it is not loaded ahead for the rest.""" - assert fork_template.preload_names(_learned(code=["core", "code.only"])) == ["code.only", "core"] - two = _learned( - code=["core", "both", "code.only"], ppt=["core", "both", "deck.engine"], design=["core", "design.only"] - ) - assert fork_template.preload_names(two) == ["both", "core"] - assert fork_template.preload_names(_learned()) == [] - assert fork_template.preload_names({"python": [2, 7], "products": {"code": ["core"]}}) == [] - assert fork_template.preload_names({"python": list(sys.version_info[:2]), "modules": ["core"]}) == [] - assert fork_template.preload_names(["not", "a", "record"]) == [] - assert fork_template.preload_names(_learned(code=["core", 7], odd="not a list")) == ["core"] - - -def test_an_agent_is_learned_under_its_product_folder() -> None: - assert fork_template.product_of([sys.executable, "/x/agents/raven-code/run.py", "--acp"]) == "raven-code" - assert fork_template.product_of([sys.executable, "-m", "raven", "acp"]) == "raven" - assert fork_template.product_of([sys.executable, "run.py"]) == "raven" - - -def test_an_agent_adds_to_its_products_record_and_starts_over_on_another_python(tmp_path: Path) -> None: - learned = tmp_path / "t" / "modules.json" - learned.parent.mkdir() - learned.write_text(json.dumps(_learned(code=["kept.before"], ppt=["deck.engine"]))) - - fork_template._learn(learned, ("json",), "code") - products = json.loads(learned.read_text())["products"] - assert "kept.before" in products["code"] and "os" in products["code"] - assert not any(n.startswith("json") for n in products["code"]) and "__main__" not in products["code"] - assert products["ppt"] == ["deck.engine"], "another product's record is left as it was" - - learned.write_text(json.dumps({"python": [2, 7], "products": {"code": ["stale.one"]}})) - fork_template._learn(learned, (), "code") - assert "stale.one" not in json.loads(learned.read_text())["products"]["code"] - - -def test_an_agents_end_is_read_the_way_an_interpreter_reads_it(capsys) -> None: - assert fork_template._exit_code(SystemExit(None)) == 0 - assert fork_template._exit_code(SystemExit(4)) == 4 - assert fork_template._exit_code(SystemExit("bad config")) == 1 - assert "bad config" in capsys.readouterr().err - assert fork_template._exit_code(ValueError("boom")) == 1 - assert "ValueError: boom" in capsys.readouterr().err - - -def test_a_script_runs_with_its_own_directory_first_on_the_path(tmp_path: Path, ended) -> None: - script = _product( - tmp_path, "import sys\nsys.exit(len(sys.argv) * 10 + (sys.path[0] == __file__.rsplit('/', 1)[0]))\n" - ) - with pytest.raises(_Ended) as caught: - fork_template._run_agent([sys.executable, str(script), "a", "b"]) - assert caught.value.code == 31 - - -def test_a_launchers_exec_of_raven_becomes_raven_here_and_any_other_exec_is_real( - tmp_path: Path, ended, monkeypatch, capsys -) -> None: - script = _product( - tmp_path, "import os, sys\nos.execv(sys.executable, [sys.executable, '-m', 'raven', '--version'])\n" - ) - with pytest.raises(_Ended): - fork_template._run_agent([sys.executable, str(script)]) - assert ended[0] == 0 - assert "Raven v" in capsys.readouterr().out - - execs: list[list[str]] = [] - - def real(path, args): - execs.append([str(path), *args]) - raise _Ended(99) - - monkeypatch.setattr(os, "execv", real) - other = _product(tmp_path, "import os\nos.execv('/bin/true', ['/bin/true'])\n") - with pytest.raises(_Ended): - fork_template._run_agent([sys.executable, str(other)]) - assert execs == [["/bin/true", "/bin/true"]] - - -def test_a_module_launch_runs_raven_here(ended, capsys) -> None: - with pytest.raises(_Ended) as caught: - fork_template._run_agent([sys.executable, "-m", "raven", "--version"]) - assert caught.value.code == 0 - assert "Raven v" in capsys.readouterr().out - - -# The test process runs threads (the host side below, pytest's own); the forked -# child does nothing but exit, which is what the warning's deadlock cannot reach. -@pytest.mark.filterwarnings("ignore:This process .* is multi-threaded:DeprecationWarning") -def test_the_template_loop_forks_reports_the_exit_and_stops_when_let_go(tmp_path: Path, monkeypatch) -> None: - import socket - import threading - - from raven.providers import relay_wire - - counts = iter([2]) - monkeypatch.setattr(fork_template.threading, "active_count", lambda: next(counts, 1)) - path = str(tmp_path / "t.sock") - life_r, life_w = os.pipe() - os.write(life_w, b"x") - seen: dict[str, object] = {} - - def host() -> None: - while not os.path.exists(path): - time.sleep(0.01) - pipes = [os.pipe() for _ in range(3)] - with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as bad: - bad.connect(path) - socket.send_fds(bad, [b"F"], [pipes[0][0]]) - relay_wire.send_frame(bad, {"argv": ["x"], "cwd": "/"}) - seen["refused"] = relay_wire.recv_frame(bad) - with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as busy: - busy.connect(path) - socket.send_fds(busy, [b"F"], [pipes[0][0], pipes[1][1], pipes[2][1]]) - relay_wire.send_frame(busy, {"argv": [sys.executable, "x.py"], "env": {}, "cwd": "/"}) - seen["threaded"] = relay_wire.recv_frame(busy) - with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as sock: - sock.connect(path) - socket.send_fds(sock, [b"F"], [pipes[0][0], pipes[1][1], pipes[2][1]]) - relay_wire.send_frame(sock, {"argv": [sys.executable, "x.py"], "env": {}, "cwd": "/"}) - seen["forked"] = relay_wire.recv_frame(sock) - seen["exit"] = relay_wire.recv_frame(sock) - os.close(life_w) - - thread = threading.Thread(target=host) - thread.start() - forked = fork_template._serve(path, lifeline=life_r) - if forked is not None: - os._exit(9 if forked[0]["argv"] == [sys.executable, "x.py"] and len(forked[1]) == 3 else 1) - thread.join(timeout=30) - os.close(life_r) - - assert "error" in seen["refused"] - assert "thread" in seen["threaded"]["error"] - assert isinstance(seen["forked"].get("pid"), int) - assert seen["exit"] == {"exit": 9} - assert signal.getsignal(signal.SIGCHLD) == signal.SIG_DFL - - def test_an_argv_the_check_cannot_read_is_not_forked(monkeypatch, tmp_path: Path) -> None: script = _product(tmp_path, "pass\n") assert not raven_agent_argv([None, str(script)]) # type: ignore[list-item] @@ -592,62 +401,6 @@ async def test_a_template_built_for_other_skips_is_replaced_while_it_has_no_agen assert fork_template._TEMPLATE is not old and old is not None and not old.running -def test_learning_survives_an_unreadable_record_and_an_unwritable_place(tmp_path: Path) -> None: - fork_template._load((), tmp_path / "missing.json", ()) - corrupt = tmp_path / "modules.json" - corrupt.write_text("{not json") - fork_template._learn(corrupt, (), "code") - assert "os" in json.loads(corrupt.read_text())["products"]["code"] - corrupt.write_text(json.dumps(["not", "a", "record"])) - fork_template._learn(corrupt, (), "code") - assert "os" in json.loads(corrupt.read_text())["products"]["code"] - blocked = tmp_path / "file-not-dir" - blocked.write_text("") - fork_template._learn(blocked / "modules.json", (), "code") - - -def test_a_script_that_just_ends_exits_zero(tmp_path: Path, ended) -> None: - script = _product(tmp_path, "x = 1\n") - with pytest.raises(_Ended): - fork_template._run_agent([sys.executable, str(script)]) - assert ended == [0] - - -def test_the_template_entry_loads_freezes_and_serves(tmp_path: Path, monkeypatch) -> None: - import gc - - served: list[str] = [] - became: list[tuple] = [] - monkeypatch.setattr(gc, "freeze", lambda: None) - monkeypatch.setattr(fork_template, "_load", lambda seed, learned, skip: 0) - monkeypatch.setattr(fork_template.threading, "active_count", lambda: 1) - - def serve(path, *, lifeline): - served.append(path) - return None - - monkeypatch.setattr(fork_template, "_serve", serve) - args = ["--socket", str(tmp_path / "s"), "--learned", str(tmp_path / "m.json"), "--skip", "litellm"] - assert fork_template.main(args) == 0 - assert served == [str(tmp_path / "s")] - - monkeypatch.setattr(fork_template, "_serve", lambda path, *, lifeline: ({"argv": ["x"]}, [1, 2, 3])) - monkeypatch.setattr(fork_template, "_become", lambda *a: became.append(a)) - assert fork_template.main(args) == 1 - assert became[0][0] == {"argv": ["x"]} and became[0][4] == ["litellm"] - - -def test_the_template_entry_refuses_to_serve_with_a_thread_running(tmp_path: Path, monkeypatch, capsys) -> None: - import gc - - monkeypatch.setattr(gc, "freeze", lambda: None) - monkeypatch.setattr(fork_template, "_load", lambda seed, learned, skip: 0) - monkeypatch.setattr(fork_template.threading, "active_count", lambda: 2) - args = ["--socket", str(tmp_path / "s"), "--learned", str(tmp_path / "m.json")] - assert fork_template.main(args) == 3 - assert json.loads(capsys.readouterr().out.strip().splitlines()[-1])["ready"] is False - - async def test_a_relay_that_cannot_start_leaves_the_agent_to_load_litellm(tmp_path: Path, monkeypatch) -> None: from raven.acp_client.client import AcpClient from raven.providers import relay_server @@ -674,6 +427,21 @@ async def no_relay(): await client.close() +async def test_a_forked_agent_trusts_the_store_a_launched_one_would(tmp_path: Path, monkeypatch) -> None: + """The template imports aiohttp, which builds its default TLS context on import: after the switch to the + system's certificate store, as in an agent launched fresh.""" + monkeypatch.delenv("RAVEN_NO_SYSTEM_CA", raising=False) + script = _product( + tmp_path, + "import aiohttp.connector, truststore\n" + "print(isinstance(aiohttp.connector._SSL_CONTEXT_VERIFIED, truststore.SSLContext), flush=True)\n", + ) + endpoint = await _launch(script) + out, _ = await _drain(endpoint) + endpoint.release() + assert out == b"True\n" + + def test_a_new_event_loop_gets_a_new_template_and_the_old_one_ends(tmp_path: Path) -> None: """A CLI turn is one asyncio.run; a template opened under the last one cannot serve the next.""" script = _product(tmp_path, "print('ran', flush=True)\n") @@ -701,3 +469,26 @@ async def once() -> tuple[int, bytes]: pytest.fail("the template opened under the ended loop is still running") finally: asyncio.run(fork_template.close_template()) + + +def test_the_client_shelf_loads_no_module_by_name() -> None: + """``raven.acp_client`` sits under the surfaces, and starts the template through the entry point. + + import-linter reads import statements, so a module loaded through importlib + or runpy is an edge no contract sees. The template's own side, which does + load the surfaces by name, is a launcher beside the CLI. + """ + import ast + + import raven.acp_client + + root = Path(raven.acp_client.__file__).parent + found = [] + for path in sorted(root.rglob("*.py")): + for node in ast.walk(ast.parse(path.read_text(encoding="utf-8"))): + if isinstance(node, ast.Call): + fn = node.func + name = fn.attr if isinstance(fn, ast.Attribute) else getattr(fn, "id", "") + if name in {"import_module", "__import__", "run_module", "run_path"}: + found.append(f"{path.name}:{node.lineno} {name}") + assert found == [] diff --git a/tests/test_cli_fork_template.py b/tests/test_cli_fork_template.py new file mode 100644 index 000000000..2dcd87199 --- /dev/null +++ b/tests/test_cli_fork_template.py @@ -0,0 +1,290 @@ +"""The fork template's own process, driven in this one: what it loads, learns, forks and becomes (Linux only).""" + +from __future__ import annotations + +import json +import os +import signal +import sys +import textwrap +import time +from pathlib import Path + +import pytest + +from raven.acp_client import fork_template +from raven.acp_client.fork_template import template_available +from raven.cli import _fork_template as template +from raven.cli import entry +from raven.security import tls + +pytestmark = pytest.mark.skipif(not template_available(), reason="forking agents is Linux-only") + + +def _product(tmp_path: Path, body: str) -> Path: + folder = tmp_path / "product" + folder.mkdir(exist_ok=True) + (folder / "subagent.json").write_text('{"name": "P"}', encoding="utf-8") + script = folder / "run.py" + script.write_text(textwrap.dedent(body), encoding="utf-8") + return script + + +class _Ended(Exception): + def __init__(self, code: int) -> None: + super().__init__(code) + self.code = code + + +@pytest.fixture +def ended(monkeypatch: pytest.MonkeyPatch) -> list[int]: + """Run the agent-side functions here: the exit becomes an exception, and what they change is put back. + + The codes are recorded because the real exit never returns, while this one + unwinds through the frames above it, which may read it as a second ending. + """ + codes: list[int] = [] + + def finish(code: int) -> None: + codes.append(code) + raise _Ended(code) + + monkeypatch.setattr(template, "_finish", finish) + monkeypatch.setattr(os, "execv", os.execv) + # Raven's entry switches the whole process to the system's certificate store, + # and this process is pytest's. + monkeypatch.setenv(tls.OPT_OUT_ENV, "1") + monkeypatch.setattr(sys, "argv", list(sys.argv)) + monkeypatch.setattr(sys, "path", list(sys.path)) + return codes + + +def _learned(**products: list[str]) -> dict: + return {"python": list(sys.version_info[:2]), "products": products} + + +def test_the_template_loads_the_seed_and_what_agents_imported_but_not_what_it_skips( + tmp_path: Path, monkeypatch +) -> None: + for name in ("colorsys", "netrc", "this_module_does_not_exist"): + monkeypatch.delitem(sys.modules, name, raising=False) + learned = tmp_path / "modules.json" + learned.write_text(json.dumps(_learned(only=["netrc", "this_module_does_not_exist"]))) + + template._load(("colorsys",), learned, ()) + assert "colorsys" in sys.modules and "netrc" in sys.modules + + monkeypatch.delitem(sys.modules, "netrc") + template._load((), learned, ("net",)) + assert "netrc" not in sys.modules + + learned.write_text(json.dumps({"python": [2, 7], "products": {"only": ["netrc"]}})) + template._load((), learned, ()) + assert "netrc" not in sys.modules + + +def test_one_products_own_modules_are_loaded_ahead_only_while_it_is_the_only_one_seen() -> None: + """What only one product imports is shared by no other process, so it is not loaded ahead for the rest.""" + assert template.preload_names(_learned(code=["core", "code.only"])) == ["code.only", "core"] + two = _learned( + code=["core", "both", "code.only"], ppt=["core", "both", "deck.engine"], design=["core", "design.only"] + ) + assert template.preload_names(two) == ["both", "core"] + assert template.preload_names(_learned()) == [] + assert template.preload_names({"python": [2, 7], "products": {"code": ["core"]}}) == [] + assert template.preload_names({"python": list(sys.version_info[:2]), "modules": ["core"]}) == [] + assert template.preload_names(["not", "a", "record"]) == [] + assert template.preload_names(_learned(code=["core", 7], odd="not a list")) == ["core"] + + +def test_an_agent_is_learned_under_its_product_folder() -> None: + assert template.product_of([sys.executable, "/x/agents/raven-code/run.py", "--acp"]) == "raven-code" + assert template.product_of([sys.executable, "-m", "raven", "acp"]) == "raven" + assert template.product_of([sys.executable, "run.py"]) == "raven" + + +def test_an_agent_adds_to_its_products_record_and_starts_over_on_another_python(tmp_path: Path) -> None: + learned = tmp_path / "t" / "modules.json" + learned.parent.mkdir() + learned.write_text(json.dumps(_learned(code=["kept.before"], ppt=["deck.engine"]))) + + template._learn(learned, ("json",), "code") + products = json.loads(learned.read_text())["products"] + assert "kept.before" in products["code"] and "os" in products["code"] + assert not any(n.startswith("json") for n in products["code"]) and "__main__" not in products["code"] + assert products["ppt"] == ["deck.engine"], "another product's record is left as it was" + + learned.write_text(json.dumps({"python": [2, 7], "products": {"code": ["stale.one"]}})) + template._learn(learned, (), "code") + assert "stale.one" not in json.loads(learned.read_text())["products"]["code"] + + +def test_an_agents_end_is_read_the_way_an_interpreter_reads_it(capsys) -> None: + assert template._exit_code(SystemExit(None)) == 0 + assert template._exit_code(SystemExit(4)) == 4 + assert template._exit_code(SystemExit("bad config")) == 1 + assert "bad config" in capsys.readouterr().err + assert template._exit_code(ValueError("boom")) == 1 + assert "ValueError: boom" in capsys.readouterr().err + + +def test_a_script_runs_with_its_own_directory_first_on_the_path(tmp_path: Path, ended) -> None: + script = _product( + tmp_path, "import sys\nsys.exit(len(sys.argv) * 10 + (sys.path[0] == __file__.rsplit('/', 1)[0]))\n" + ) + with pytest.raises(_Ended) as caught: + template._run_agent([sys.executable, str(script), "a", "b"]) + assert caught.value.code == 31 + + +def test_a_launchers_exec_of_raven_becomes_raven_here_and_any_other_exec_is_real( + tmp_path: Path, ended, monkeypatch, capsys +) -> None: + script = _product( + tmp_path, "import os, sys\nos.execv(sys.executable, [sys.executable, '-m', 'raven', '--version'])\n" + ) + with pytest.raises(_Ended): + template._run_agent([sys.executable, str(script)]) + assert ended[0] == 0 + assert "Raven v" in capsys.readouterr().out + + execs: list[list[str]] = [] + + def real(path, args): + execs.append([str(path), *args]) + raise _Ended(99) + + monkeypatch.setattr(os, "execv", real) + other = _product(tmp_path, "import os\nos.execv('/bin/true', ['/bin/true'])\n") + with pytest.raises(_Ended): + template._run_agent([sys.executable, str(other)]) + assert execs == [["/bin/true", "/bin/true"]] + + +def test_a_module_launch_runs_raven_here(ended, capsys) -> None: + with pytest.raises(_Ended) as caught: + template._run_agent([sys.executable, "-m", "raven", "--version"]) + assert caught.value.code == 0 + assert "Raven v" in capsys.readouterr().out + + +# The test process runs threads (the host side below, pytest's own); the forked +# child does nothing but exit, which is what the warning's deadlock cannot reach. +@pytest.mark.filterwarnings("ignore:This process .* is multi-threaded:DeprecationWarning") +def test_the_template_loop_forks_reports_the_exit_and_stops_when_let_go(tmp_path: Path, monkeypatch) -> None: + import socket + import threading + + from raven.providers import relay_wire + + counts = iter([2]) + monkeypatch.setattr(template.threading, "active_count", lambda: next(counts, 1)) + path = str(tmp_path / "t.sock") + life_r, life_w = os.pipe() + os.write(life_w, b"x") + seen: dict[str, object] = {} + + def host() -> None: + while not os.path.exists(path): + time.sleep(0.01) + pipes = [os.pipe() for _ in range(3)] + with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as bad: + bad.connect(path) + socket.send_fds(bad, [b"F"], [pipes[0][0]]) + relay_wire.send_frame(bad, {"argv": ["x"], "cwd": "/"}) + seen["refused"] = relay_wire.recv_frame(bad) + with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as busy: + busy.connect(path) + socket.send_fds(busy, [b"F"], [pipes[0][0], pipes[1][1], pipes[2][1]]) + relay_wire.send_frame(busy, {"argv": [sys.executable, "x.py"], "env": {}, "cwd": "/"}) + seen["threaded"] = relay_wire.recv_frame(busy) + with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as sock: + sock.connect(path) + socket.send_fds(sock, [b"F"], [pipes[0][0], pipes[1][1], pipes[2][1]]) + relay_wire.send_frame(sock, {"argv": [sys.executable, "x.py"], "env": {}, "cwd": "/"}) + seen["forked"] = relay_wire.recv_frame(sock) + seen["exit"] = relay_wire.recv_frame(sock) + os.close(life_w) + + thread = threading.Thread(target=host) + thread.start() + forked = template._serve(path, lifeline=life_r) + if forked is not None: + os._exit(9 if forked[0]["argv"] == [sys.executable, "x.py"] and len(forked[1]) == 3 else 1) + thread.join(timeout=30) + os.close(life_r) + + assert "error" in seen["refused"] + assert "thread" in seen["threaded"]["error"] + assert isinstance(seen["forked"].get("pid"), int) + assert seen["exit"] == {"exit": 9} + assert signal.getsignal(signal.SIGCHLD) == signal.SIG_DFL + + +def test_learning_survives_an_unreadable_record_and_an_unwritable_place(tmp_path: Path) -> None: + template._load((), tmp_path / "missing.json", ()) + corrupt = tmp_path / "modules.json" + corrupt.write_text("{not json") + template._learn(corrupt, (), "code") + assert "os" in json.loads(corrupt.read_text())["products"]["code"] + corrupt.write_text(json.dumps(["not", "a", "record"])) + template._learn(corrupt, (), "code") + assert "os" in json.loads(corrupt.read_text())["products"]["code"] + blocked = tmp_path / "file-not-dir" + blocked.write_text("") + template._learn(blocked / "modules.json", (), "code") + + +def test_a_script_that_just_ends_exits_zero(tmp_path: Path, ended) -> None: + script = _product(tmp_path, "x = 1\n") + with pytest.raises(_Ended): + template._run_agent([sys.executable, str(script)]) + assert ended == [0] + + +def test_the_template_entry_loads_freezes_and_serves(tmp_path: Path, monkeypatch) -> None: + import gc + + served: list[str] = [] + became: list[tuple] = [] + monkeypatch.setattr(gc, "freeze", lambda: None) + monkeypatch.setattr(template, "_load", lambda seed, learned, skip: 0) + monkeypatch.setattr(template.threading, "active_count", lambda: 1) + + def serve(path, *, lifeline): + served.append(path) + return None + + monkeypatch.setattr(template, "_serve", serve) + args = ["--socket", str(tmp_path / "s"), "--learned", str(tmp_path / "m.json"), "--skip", "litellm"] + assert template.main(args) == 0 + assert served == [str(tmp_path / "s")] + + monkeypatch.setattr(template, "_serve", lambda path, *, lifeline: ({"argv": ["x"]}, [1, 2, 3])) + monkeypatch.setattr(template, "_become", lambda *a: became.append(a)) + assert template.main(args) == 1 + assert became[0][0] == {"argv": ["x"]} and became[0][4] == ["litellm"] + + +def test_the_template_entry_refuses_to_serve_with_a_thread_running(tmp_path: Path, monkeypatch, capsys) -> None: + import gc + + monkeypatch.setattr(gc, "freeze", lambda: None) + monkeypatch.setattr(template, "_load", lambda seed, learned, skip: 0) + monkeypatch.setattr(template.threading, "active_count", lambda: 2) + args = ["--socket", str(tmp_path / "s"), "--learned", str(tmp_path / "m.json")] + assert template.main(args) == 3 + assert json.loads(capsys.readouterr().out.strip().splitlines()[-1])["ready"] is False + + +def test_the_entry_settles_trust_then_hands_the_template_its_arguments(monkeypatch) -> None: + """The host starts the template through Raven's entry point, so the switch comes before every import.""" + calls: list = [] + monkeypatch.setattr(tls, "use_system_ca", lambda: calls.append("trust")) + monkeypatch.setattr(template, "main", lambda argv: calls.append(list(argv)) or 0) + monkeypatch.setattr(sys, "argv", ["raven", fork_template.TEMPLATE_ARG, "--socket", "s"]) + + with pytest.raises(SystemExit) as ended: + entry.run() + assert ended.value.code == 0 + assert calls == ["trust", ["--socket", "s"]] From 3f2abd1f4a3b2b3350ce2ce669533c7dfcbe7f42 Mon Sep 17 00:00:00 2001 From: Dizhan Xue Date: Tue, 27 Oct 2026 14:05:05 +0000 Subject: [PATCH 05/19] fix(acp_client): end the fork template with the event loop that started it A template was ended by close() or, once its loop had gone, by abandon(), which killed it but neither reaped it nor closed its transport. On a pidfd host the closed loop's watcher could no longer reap it, so it stayed a zombie, and the transport's finalizer later raised RuntimeError: Event loop is closed. A task on the template's own loop now holds it, and when asyncio.run cancels that task at the end of the run it closes the template while the loop can still reap it and close its pipes. close() waits out the group kill as well, and a close cut short by the loop's end still kills the template and closes its pipes. abandon() stays as the fallback for a loop closed some other way, and marks the transport closed so its finalizer does not run again. --- CONTEXT.md | 18 +++--- raven/acp_client/fork_template.py | 95 ++++++++++++++++++++++--------- tests/test_acp_fork_template.py | 59 +++++++++++++++---- 3 files changed, 126 insertions(+), 46 deletions(-) diff --git a/CONTEXT.md b/CONTEXT.md index efa28cfc1..8a746ed17 100644 --- a/CONTEXT.md +++ b/CONTEXT.md @@ -3416,14 +3416,16 @@ Two halves: the host's side in `acp_client` decides when to ask, starts and ends holds the endpoint a forked agent runs behind; the template's own side is a launcher beside the CLI, because it loads the surfaces by name and runs `python -m raven`, which the client shelf may not. The client starts it through Raven's entry point, which settles the certificate store first, -and never imports it. It starts through the parent-death wrapper like any agent and its stdin is a -lifeline only the host holds. Each forked agent leads a session of its own and arms the -parent-death signal against the template. A product launcher's `os.execv` of `python -m raven` -runs Raven in the same process. Agents write the modules they imported beside the template, per -product, and the next template loads the modules at least two products import (one product's own -while only one has been seen): a module one product alone imports is shared by no other process, -so it is left for that product to import after the fork. Only the host's own interpreter starting -`-m raven` or a product launcher is forked; anything else launches the ordinary way. +and never imports it. It starts through the parent-death wrapper like any agent, its stdin is a +lifeline only the host holds, and it belongs to the event loop that started it: its pipes and the +watcher that reaps it are that loop's, so it is ended with that loop and a later loop starts +another. Each forked agent leads a session of its own and arms the parent-death signal against the +template. A product launcher's `os.execv` of `python -m raven` runs Raven in the same process. +Agents write the modules they imported beside the template, per product, and the next template +loads the modules at least two products import (one product's own while only one has been seen): a +module one product alone imports is shared by no other process, so it is left for that product to +import after the fork. Only the host's own interpreter starting `-m raven` or a product launcher +is forked; anything else launches the ordinary way. **Leftover agent** (`raven/acp_client/leftovers.py`): An ACP agent process group whose launching host is no longer running. An ACP agent runs in a diff --git a/raven/acp_client/fork_template.py b/raven/acp_client/fork_template.py index 121b0bdb4..bbe77689d 100644 --- a/raven/acp_client/fork_template.py +++ b/raven/acp_client/fork_template.py @@ -16,7 +16,10 @@ **The template** is started through the same parent-death wrapper as any agent (``leftovers``), so it ends with the host, and its stdin is a lifeline only the -host holds: letting go of it asks the template to stop. +host holds: letting go of it asks the template to stop. It belongs to the event +loop it was started on -- its pipes and the watcher that reaps it are that +loop's -- so it is ended with that loop (``ForkTemplate._hold``), and a host on a +later loop starts another. **A fork.** The host hands over the agent's argv, environment and working directory, and the three pipe ends that become its stdin, stdout and stderr. @@ -200,6 +203,7 @@ def __init__(self, *, skip: Sequence[str] = ()) -> None: self._forking = 0 self._idle_task: asyncio.Task | None = None self._log_task: asyncio.Task | None = None + self._hold_task: asyncio.Task | None = None self._lock = asyncio.Lock() self.loop = asyncio.get_running_loop() """The loop its process and pipes were opened under; the handle is no use to another.""" @@ -245,6 +249,7 @@ async def _start(self, env: dict[str, str]) -> None: ) leftovers.record(self._proc.pid, "fork-template") self._log_task = asyncio.get_running_loop().create_task(self._log_stderr(self._proc), name="fork-template-log") + self._hold_task = asyncio.get_running_loop().create_task(self._hold(self._proc), name="fork-template-hold") ready = await asyncio.wait_for(self._ready_report(self._proc), timeout=_READY_TIMEOUT_S) if not ready.get("ready"): raise RuntimeError(str(ready.get("reason") or "the template exited before it was ready")) @@ -373,11 +378,30 @@ async def wait_then_close() -> None: self._idle_task = asyncio.get_running_loop().create_task(wait_then_close(), name="fork-template-idle") + async def _hold(self, proc: asyncio.subprocess.Process) -> None: + """Keep the template for as long as its loop runs, and end it with the loop. + + The template's pipes and the watcher that reaps it belong to the loop it + was started on. ``asyncio.run`` cancels the tasks still pending when its + coroutine returns and lets them finish before it closes the loop, so, + cancelled here, the template is ended, reaped and its pipes closed while + that loop still can. A loop closed any other way leaves it to :meth:`abandon`. + """ + try: + await proc.wait() + except asyncio.CancelledError: + if self._proc is proc: + await self.close() + raise + def abandon(self) -> None: """End the template without the loop it was started on, which may be closed by now. + The fallback for a loop that closed without letting :meth:`_hold` finish. The agents it forked end with it (they armed its death signal); a loop - that has ended has no use for them either, their pipes being its. + that has ended has no use for them either, their pipes being its. Closing + the transport here may fail on that loop, but it marks the transport closed + first, so its finalizer does not try again on a loop that is gone. """ from raven.acp_client import leftovers @@ -386,6 +410,10 @@ def abandon(self) -> None: with contextlib.suppress(ProcessLookupError, PermissionError): os.killpg(proc.pid, signal.SIGKILL) leftovers.forget(proc.pid) + transport = getattr(proc, "_transport", None) + if transport is not None: + with contextlib.suppress(Exception): + transport.close() if self._dir is not None: import shutil @@ -396,30 +424,45 @@ async def close(self) -> None: from raven.acp_client import leftovers proc, self._proc = self._proc, None - if proc is not None and proc.returncode is None: - # Letting go of its stdin is how the template is asked to stop; the - # group kill after the grace is for one that does not. - if proc.stdin is not None: - with contextlib.suppress(Exception): - proc.stdin.close() - with contextlib.suppress(asyncio.TimeoutError, ProcessLookupError): - await asyncio.wait_for(proc.wait(), timeout=_TEARDOWN_GRACE_S) - with contextlib.suppress(ProcessLookupError, PermissionError): - os.killpg(proc.pid, signal.SIGKILL) - if proc is not None: - leftovers.forget(proc.pid) - transport = getattr(proc, "_transport", None) - if transport is not None: - with contextlib.suppress(Exception): - transport.close() - if self._log_task is not None: - self._log_task.cancel() - self._log_task = None - if self._dir is not None: - import shutil - - shutil.rmtree(self._dir, ignore_errors=True) - self._dir = None + hold, self._hold_task = self._hold_task, None + try: + if proc is not None and proc.returncode is None: + # Letting go of its stdin is how the template is asked to stop; the + # group kill after the grace is for one that does not, and is waited + # out too, so the template is reaped before this returns. + if proc.stdin is not None: + with contextlib.suppress(Exception): + proc.stdin.close() + with contextlib.suppress(asyncio.TimeoutError, ProcessLookupError): + await asyncio.wait_for(proc.wait(), timeout=_TEARDOWN_GRACE_S) + if proc.returncode is None: + with contextlib.suppress(ProcessLookupError, PermissionError): + os.killpg(proc.pid, signal.SIGKILL) + with contextlib.suppress(asyncio.TimeoutError, ProcessLookupError): + await asyncio.wait_for(proc.wait(), timeout=_TEARDOWN_GRACE_S) + finally: + # Also when a wait above is cancelled, as when its loop ends: the + # template is not left running, and its pipes are closed while that + # loop can still close them. + if proc is not None: + if proc.returncode is None: + with contextlib.suppress(ProcessLookupError, PermissionError): + os.killpg(proc.pid, signal.SIGKILL) + leftovers.forget(proc.pid) + transport = getattr(proc, "_transport", None) + if transport is not None: + with contextlib.suppress(Exception): + transport.close() + if self._log_task is not None: + self._log_task.cancel() + self._log_task = None + if hold is not None and hold is not asyncio.current_task(): + hold.cancel() + if self._dir is not None: + import shutil + + shutil.rmtree(self._dir, ignore_errors=True) + self._dir = None _TEMPLATE: ForkTemplate | None = None diff --git a/tests/test_acp_fork_template.py b/tests/test_acp_fork_template.py index 22ecfb72e..1b8317e06 100644 --- a/tests/test_acp_fork_template.py +++ b/tests/test_acp_fork_template.py @@ -442,8 +442,12 @@ async def test_a_forked_agent_trusts_the_store_a_launched_one_would(tmp_path: Pa assert out == b"True\n" -def test_a_new_event_loop_gets_a_new_template_and_the_old_one_ends(tmp_path: Path) -> None: - """A CLI turn is one asyncio.run; a template opened under the last one cannot serve the next.""" +# The pipes and the child watcher belong to the loop: a template still held when +# that loop closes leaves an unreaped child on a pidfd host and a transport whose +# finalizer fails on the closed loop, which pytest reports as unraisable. +@pytest.mark.filterwarnings("error::pytest.PytestUnraisableExceptionWarning") +def test_a_template_ends_with_the_event_loop_that_started_it(tmp_path: Path) -> None: + """A CLI turn is one asyncio.run: the template it opened is gone, reaped, when that run returns.""" script = _product(tmp_path, "print('ran', flush=True)\n") async def once() -> tuple[int, bytes]: @@ -454,23 +458,54 @@ async def once() -> tuple[int, bytes]: assert template is not None and template._proc is not None return template._proc.pid, out - first_pid, first_out = asyncio.run(once()) - second_pid, second_out = asyncio.run(once()) try: + first_pid, first_out = asyncio.run(once()) + assert not Path(f"/proc/{first_pid}").exists(), "the template outlived its loop, or was left unreaped" + second_pid, second_out = asyncio.run(once()) assert first_out == second_out == b"ran\n" assert second_pid != first_pid - for _ in range(100): - try: - os.kill(first_pid, 0) - except ProcessLookupError: - break - time.sleep(0.05) - else: - pytest.fail("the template opened under the ended loop is still running") + assert not Path(f"/proc/{second_pid}").exists() finally: asyncio.run(fork_template.close_template()) +async def _deaf_template(tmp_path: Path, monkeypatch) -> fork_template.ForkTemplate: + """A template that reports ready and never reads its stdin, so letting go of it does not stop it.""" + (tmp_path / "deaf_template.py").write_text( + "import json, time\nprint(json.dumps({'ready': True}), flush=True)\ntime.sleep(600)\n", encoding="utf-8" + ) + monkeypatch.setattr(fork_template, "_TEMPLATE_PROGRAM", ("-m", "deaf_template")) + template = fork_template.ForkTemplate() + await template._start({**os.environ, "PYTHONPATH": str(tmp_path)}) + assert template.running + return template + + +async def test_a_template_that_ignores_its_lifeline_is_killed_and_reaped(tmp_path: Path, monkeypatch) -> None: + template = await _deaf_template(tmp_path, monkeypatch) + proc = template._proc + assert proc is not None + + await template.close() + assert proc.returncode is not None + + +async def test_a_close_cut_short_still_ends_the_template_and_closes_its_pipes(tmp_path: Path, monkeypatch) -> None: + """The loop ending cancels a close still waiting out the grace: an idle close, or the hold task's own.""" + template = await _deaf_template(tmp_path, monkeypatch) + proc = template._proc + assert proc is not None + + closing = asyncio.get_running_loop().create_task(template.close()) + await asyncio.sleep(0.2) + assert not closing.done() + closing.cancel() + with pytest.raises(asyncio.CancelledError): + await closing + assert proc._transport.is_closing() + await asyncio.wait_for(proc.wait(), timeout=10) + + def test_the_client_shelf_loads_no_module_by_name() -> None: """``raven.acp_client`` sits under the surfaces, and starts the template through the entry point. From 1e2a46e1ef7aa3256ba71962cccab3f8ec783646 Mon Sep 17 00:00:00 2001 From: Dizhan Xue Date: Tue, 27 Oct 2026 14:05:20 +0000 Subject: [PATCH 06/19] fix(*): leave a fork template and relay alone while their loop still runs A host can run a second event loop beside its own: the RPC server runs a CLI command such as raven agent -m on a worker thread, under its own asyncio.run. When that loop launched a Raven agent, the launch took the first loop's template and relay for ones an ended loop had left behind and abandoned them, and every agent forked from that template died with it through the parent-death signal. Only a template or relay whose loop has stopped is abandoned now. A loop running beside them launches its agents the ordinary way and is offered the relay's socket, which the loop that owns it goes on answering. --- CONTEXT.md | 26 +++++++++++++++----------- raven/acp_client/fork_template.py | 8 +++++++- raven/providers/relay_server.py | 8 +++++++- tests/test_acp_fork_template.py | 20 ++++++++++++++++++++ tests/test_provider_relay.py | 15 +++++++++++++++ 5 files changed, 64 insertions(+), 13 deletions(-) diff --git a/CONTEXT.md b/CONTEXT.md index 8a746ed17..8212b1e07 100644 --- a/CONTEXT.md +++ b/CONTEXT.md @@ -3407,7 +3407,9 @@ token in the launch environment (`RAVEN_LLM_RELAY`), and an agent that reads the same answers (both stand on `LiteLLMShape`), with only the single `chat` / `chat_stream` attempt and the `relayed` catalogue lookups crossing to the host. Retry, classification, usage records and traces stay in the agent; the host's calls run under `recorded_by_caller` so a call -is billed once. An agent that cannot reach the host loads LiteLLM and works as before. +is billed once. The relay is served by the event loop that started it: a later loop starts +another, and a loop running beside it in another thread is offered the same socket. An agent +that cannot reach the host loads LiteLLM and works as before. **Fork template** (`raven/acp_client/fork_template.py`, `raven/cli/_fork_template.py`, `subagents.runtime.forkTemplate`): On Linux, one process per host that has imported what Raven agents import, frozen out of the @@ -3416,16 +3418,18 @@ Two halves: the host's side in `acp_client` decides when to ask, starts and ends holds the endpoint a forked agent runs behind; the template's own side is a launcher beside the CLI, because it loads the surfaces by name and runs `python -m raven`, which the client shelf may not. The client starts it through Raven's entry point, which settles the certificate store first, -and never imports it. It starts through the parent-death wrapper like any agent, its stdin is a -lifeline only the host holds, and it belongs to the event loop that started it: its pipes and the -watcher that reaps it are that loop's, so it is ended with that loop and a later loop starts -another. Each forked agent leads a session of its own and arms the parent-death signal against the -template. A product launcher's `os.execv` of `python -m raven` runs Raven in the same process. -Agents write the modules they imported beside the template, per product, and the next template -loads the modules at least two products import (one product's own while only one has been seen): a -module one product alone imports is shared by no other process, so it is left for that product to -import after the fork. Only the host's own interpreter starting `-m raven` or a product launcher -is forked; anything else launches the ordinary way. +and never imports it. It starts through the parent-death +wrapper like any agent, its stdin is a lifeline only the host holds, and it belongs to the event +loop that started it: its pipes and the watcher that reaps it are that loop's, so it is ended with +that loop and a later loop starts another; a loop running beside it in another thread launches the +ordinary way rather than end the agents it still has. Each forked agent leads a session of its own +and arms the parent-death signal against the template. A product launcher's `os.execv` of +`python -m raven` runs Raven in the same process. Agents write the modules they imported beside +the template, per product, and the next template loads the modules at least two products import +(one product's own while only one has been seen): a module one product alone imports is shared by +no other process, so it is left for that product to import after the fork. Only the host's own +interpreter starting `-m raven` or a product launcher is forked; anything else launches the +ordinary way. **Leftover agent** (`raven/acp_client/leftovers.py`): An ACP agent process group whose launching host is no longer running. An ACP agent runs in a diff --git a/raven/acp_client/fork_template.py b/raven/acp_client/fork_template.py index bbe77689d..c1a16f343 100644 --- a/raven/acp_client/fork_template.py +++ b/raven/acp_client/fork_template.py @@ -19,7 +19,8 @@ host holds: letting go of it asks the template to stop. It belongs to the event loop it was started on -- its pipes and the watcher that reaps it are that loop's -- so it is ended with that loop (``ForkTemplate._hold``), and a host on a -later loop starts another. +later loop starts another. A loop running beside it in another thread launches +the ordinary way rather than end the agents that loop still has. **A fork.** The host hands over the agent's argv, environment and working directory, and the three pipe ends that become its stdin, stdout and stderr. @@ -494,6 +495,11 @@ async def fork_launch( # caller hears why; a fork would only see the agent die. return None if _TEMPLATE is not None and _TEMPLATE.loop is not asyncio.get_running_loop(): + if _TEMPLATE.loop.is_running(): + # A loop in another thread (a CLI command the RPC server runs on a + # worker) still holds it, and its agents would die with it: this + # one launches the ordinary way. + return None # Started under a loop that has ended (a CLI turn is one asyncio.run). _TEMPLATE.abandon() _TEMPLATE = None diff --git a/raven/providers/relay_server.py b/raven/providers/relay_server.py index fe4440d7a..ffedb4e46 100644 --- a/raven/providers/relay_server.py +++ b/raven/providers/relay_server.py @@ -77,7 +77,7 @@ def __init__(self, *, build: ProviderBuilder = build_litellm_provider, tcp: bool self._providers: OrderedDict[str, LLMProvider] = OrderedDict() self._building: dict[str, asyncio.Future[LLMProvider]] = {} self.loop: asyncio.AbstractEventLoop | None = None - """The loop that answers this server's connections; a server is no use to another.""" + """The loop that answers this server's connections, and the only one that may close it.""" @property def serving(self) -> bool: @@ -260,6 +260,12 @@ async def relay_env() -> dict[str, str]: global _SERVER, _STARTING loop = asyncio.get_running_loop() if _SERVER is not None and _SERVER.loop is not loop: + if _SERVER.loop is not None and _SERVER.loop.is_running(): + # A loop in another thread serves it, and its socket answers an + # agent launched from this loop as well as its own. + if _SERVER.serving: + return _SERVER.env() + raise RuntimeError("the relay belongs to an event loop running in another thread") # Started under a loop that has ended (a CLI turn is one asyncio.run). _SERVER.abandon() _SERVER, _STARTING = None, None diff --git a/tests/test_acp_fork_template.py b/tests/test_acp_fork_template.py index 1b8317e06..ec80c2785 100644 --- a/tests/test_acp_fork_template.py +++ b/tests/test_acp_fork_template.py @@ -469,6 +469,26 @@ async def once() -> tuple[int, bytes]: asyncio.run(fork_template.close_template()) +async def test_a_loop_in_another_thread_leaves_this_loops_agents_running(tmp_path: Path) -> None: + """The RPC server runs a CLI command's own asyncio.run on a worker thread, beside its own loop.""" + script = _product(tmp_path, "import time\nprint('up', flush=True)\ntime.sleep(600)\n") + endpoint = await _launch(script) + try: + assert await asyncio.wait_for(endpoint.readline(), timeout=60) == b"up\n" + template = fork_template._TEMPLATE + + async def launch_beside() -> object: + base = dict(os.environ) + return await fork_launch([sys.executable, str(script)], base, str(tmp_path), base_env=base) + + assert await asyncio.to_thread(asyncio.run, launch_beside()) is None + assert fork_template._TEMPLATE is template and template is not None and template.running + os.kill(endpoint.pid, 0) + assert endpoint.returncode is None + finally: + endpoint.release() + + async def _deaf_template(tmp_path: Path, monkeypatch) -> fork_template.ForkTemplate: """A template that reports ready and never reads its stdin, so letting go of it does not stop it.""" (tmp_path / "deaf_template.py").write_text( diff --git a/tests/test_provider_relay.py b/tests/test_provider_relay.py index 4fc3df25c..2d4638f02 100644 --- a/tests/test_provider_relay.py +++ b/tests/test_provider_relay.py @@ -782,6 +782,21 @@ async def ask_the_second() -> tuple[dict[str, str], Any]: assert not os.path.exists(first[relay.RELAY_ENV].removeprefix("unix:")) +async def test_a_loop_in_another_thread_is_offered_the_relay_this_loop_serves(monkeypatch: pytest.MonkeyPatch) -> None: + """The RPC server runs a CLI command's own asyncio.run on a worker thread, beside its own loop.""" + from raven.providers import litellm_setup + + monkeypatch.setattr(litellm_setup, "warm_up_in_background", lambda: None) + try: + here = await relay_server.relay_env() + server = relay_server._SERVER + beside = await asyncio.to_thread(asyncio.run, relay_server.relay_env()) + assert beside == here + assert relay_server._SERVER is server and server is not None and server.serving + finally: + await relay_server.close_relay() + + def test_a_host_without_unix_sockets_serves_on_loopback_tcp(monkeypatch: pytest.MonkeyPatch) -> None: assert relay_server.RelayServer()._tcp is not hasattr(asyncio, "start_unix_server") monkeypatch.delattr(asyncio, "start_unix_server", raising=False) From 032a2bef4706f7e9760ce5b5c849c1621e63518a Mon Sep 17 00:00:00 2001 From: Dizhan Xue Date: Wed, 28 Oct 2026 06:09:34 +0000 Subject: [PATCH 07/19] test(providers): keep the bare relay host's socket path short on macOS The test that plays a misbehaving relay host bound its Unix socket under tmp_path, which on macOS is already past the 104-byte cap on a socket path, so it failed with AF_UNIX path too long. It now binds in a short directory of its own, as the relay server itself does. --- tests/test_provider_relay.py | 19 +++++++++++++------ 1 file changed, 13 insertions(+), 6 deletions(-) diff --git a/tests/test_provider_relay.py b/tests/test_provider_relay.py index 2d4638f02..d73426fe1 100644 --- a/tests/test_provider_relay.py +++ b/tests/test_provider_relay.py @@ -5,8 +5,10 @@ import asyncio import json import os +import shutil import subprocess import sys +import tempfile import textwrap from collections.abc import AsyncIterator from typing import Any @@ -630,10 +632,13 @@ async def refuse(**kwargs): @pytest.mark.skipif(not hasattr(asyncio, "start_unix_server"), reason="the bare host listens on a Unix socket") -async def test_a_stream_the_host_ends_oddly_or_drops_is_a_failure_here( - monkeypatch: pytest.MonkeyPatch, tmp_path -) -> None: - """A host speaking out of turn, played by a bare socket server.""" +async def test_a_stream_the_host_ends_oddly_or_drops_is_a_failure_here(monkeypatch: pytest.MonkeyPatch) -> None: + """A host speaking out of turn, played by a bare socket server. + + The socket goes in a short directory of its own, as the relay server puts + its own: a socket path is capped (104 bytes on macOS), and a test's + ``tmp_path`` there is already longer. + """ replies = iter([{"what": 1}, None]) async def host(reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None: @@ -643,8 +648,9 @@ async def host(reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> No await relay_wire.write_frame(writer, reply) writer.close() - path = tmp_path / "odd.sock" - server = await asyncio.start_unix_server(host, path=str(path)) + folder = tempfile.mkdtemp(prefix="raven-relay-") + path = os.path.join(folder, "odd.sock") + server = await asyncio.start_unix_server(host, path=path) monkeypatch.setenv(relay.RELAY_ENV, f"unix:{path}") monkeypatch.setenv(relay.RELAY_TOKEN_ENV, "t") relay._reset_for_tests() @@ -658,6 +664,7 @@ async def host(reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> No pass finally: server.close() + shutil.rmtree(folder, ignore_errors=True) async def test_a_stream_without_a_reachable_host_runs_here(monkeypatch: pytest.MonkeyPatch, tmp_path) -> None: From cf3c94b2ab5d8f291fbcb5112abfdccaff8f65c4 Mon Sep 17 00:00:00 2001 From: Dizhan Xue Date: Wed, 28 Oct 2026 06:09:39 +0000 Subject: [PATCH 08/19] fix(config): keep the last runtime settings when the section is not a mapping subagent_runtime read anything but a mapping at subagents.runtime as the section having left the file, so an edit such as "runtime": false forgot the last valid settings and brought the defaults back: relay on, fork template on, idle close at 600 s. Only a missing section or a null brings the defaults back now. Any other value is a rejected edit, and the last admitted answer stays, as the reader's docstring says. --- raven/config/live.py | 10 +++++++--- tests/test_acp_pool_idle_close.py | 24 ++++++++++++++++++++++++ 2 files changed, 31 insertions(+), 3 deletions(-) diff --git a/raven/config/live.py b/raven/config/live.py index adb0562bf..1c9293dc5 100644 --- a/raven/config/live.py +++ b/raven/config/live.py @@ -596,10 +596,14 @@ def subagent_runtime(live: LiveConfig) -> SubagentRuntimeConfig: """ from raven.config.schema import SubagentRuntimeConfig - section = live.get("subagents.runtime") - present = isinstance(section, dict) + parent = live.raw().get("subagents") + # Only a section that left the file, or a null (the schema's own "unset"), + # brings the defaults back. Anything else in its place -- ``false``, a + # string, a non-mapping ``subagents`` -- is an edit that failed. + section = parent.get("runtime") if isinstance(parent, dict) else parent + present = section is not None value = None - if present: + if isinstance(section, dict): try: value = SubagentRuntimeConfig.model_validate(section) except ValueError: diff --git a/tests/test_acp_pool_idle_close.py b/tests/test_acp_pool_idle_close.py index 90bb861dd..dcc4d3b55 100644 --- a/tests/test_acp_pool_idle_close.py +++ b/tests/test_acp_pool_idle_close.py @@ -206,6 +206,30 @@ def test_a_rejected_edit_keeps_the_last_valid_settings(monkeypatch: pytest.Monke assert kept.idle_close_s == 30 and kept.fork_template is False +@pytest.mark.parametrize("edit", [{"runtime": False}, {"runtime": "off"}, {"runtime": []}, False]) +def test_a_runtime_edit_that_is_not_a_mapping_is_rejected_not_removed( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path, edit: Any +) -> None: + """Only a section that left the file, or a null, brings the defaults back; any other value is a bad edit.""" + from raven.config import live as live_mod + + config = tmp_path / "config.json" + reader = live_mod.LiveConfig(config) + config.write_text( + json.dumps({"subagents": {"runtime": {"idleCloseS": 30, "llmRelay": False, "forkTemplate": False}}}), + encoding="utf-8", + ) + assert live_mod.subagent_runtime(reader).idle_close_s == 30 + config.write_text(json.dumps({"subagents": edit}), encoding="utf-8") + kept = live_mod.subagent_runtime(reader) + assert (kept.idle_close_s, kept.llm_relay, kept.fork_template) == (30, False, False) + + config.write_text(json.dumps({"subagents": {"runtime": None}}), encoding="utf-8") + assert live_mod.subagent_runtime(reader).idle_close_s == 600 + config.write_text(json.dumps({}), encoding="utf-8") + assert live_mod.subagent_runtime(reader).idle_close_s == 600 + + async def test_a_failed_idle_check_does_not_end_the_next_one(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setattr(pool_mod, "_IDLE_CHECK_MIN_S", 0.01) monkeypatch.setattr(AcpConnectionPool, "_idle_budget", staticmethod(lambda: 0.04)) From ce693a3ae3347ce44e1b23f5587ab43d8f211c36 Mon Sep 17 00:00:00 2001 From: Dizhan Xue Date: Wed, 28 Oct 2026 06:09:46 +0000 Subject: [PATCH 09/19] fix(providers): build an agent's provider on the host without exporting its key The relay host built a LiteLLMProvider from the agent's construction, and the constructor wrote the agent's key into os.environ: over the host's own value for a gateway, beside it otherwise. The host's later providers could pick the agent's key up, and every process the host started after inherited it. LiteLLMShape takes export_env, and the host passes False. Every call already passes the key explicitly, and LiteLLM's chat path was checked to send it with the variables unset (zai, fireworks_ai, huggingface, openrouter). --- raven/providers/litellm_shape.py | 7 ++++++- raven/providers/relay_server.py | 4 ++-- tests/test_provider_relay.py | 34 ++++++++++++++++++++++++++++++++ 3 files changed, 42 insertions(+), 3 deletions(-) diff --git a/raven/providers/litellm_shape.py b/raven/providers/litellm_shape.py index 1a025bdd4..4dcb8b790 100644 --- a/raven/providers/litellm_shape.py +++ b/raven/providers/litellm_shape.py @@ -73,6 +73,7 @@ def __init__( tls_verify: bool = True, *, unparsed_reasoning: bool | None = None, + export_env: bool = True, ): addressed_spec = find_by_name(provider_name) or find_by_model(default_model) if api_base and addressed_spec and addressed_spec.strip_api_base_trailing_slash: @@ -113,7 +114,11 @@ def __init__( if self._gateway and self._gateway.name == "openrouter": self.extra_headers = {**_OPENROUTER_ATTRIBUTION, **self.extra_headers} - if api_key: + # A provider built in one process for another (the LLM relay's host) + # leaves the environment alone: every call passes the key explicitly, + # and the variables would carry one agent's key into the host's own + # providers and into every process the host starts after. + if api_key and export_env: self._setup_env(api_key, api_base, default_model) def _setup_env(self, api_key: str, api_base: str | None, model: str) -> None: diff --git a/raven/providers/relay_server.py b/raven/providers/relay_server.py index ffedb4e46..b3db1ea8e 100644 --- a/raven/providers/relay_server.py +++ b/raven/providers/relay_server.py @@ -52,12 +52,12 @@ def build_litellm_provider(spec: dict[str, Any]) -> LLMProvider: - """The provider an agent would have built for itself from ``spec``.""" + """The provider an agent would have built for itself from ``spec``, its key kept out of the host's environment.""" from raven.providers.litellm_provider import LiteLLMProvider spec = dict(spec) generation = relay_wire.decode(spec.pop("generation", None)) - provider = LiteLLMProvider(**spec) + provider = LiteLLMProvider(**spec, export_env=False) if generation is not None: provider.generation = generation return provider diff --git a/tests/test_provider_relay.py b/tests/test_provider_relay.py index d73426fe1..35a611b09 100644 --- a/tests/test_provider_relay.py +++ b/tests/test_provider_relay.py @@ -496,6 +496,40 @@ def test_the_default_builder_makes_the_direct_provider_with_the_agents_settings( assert built.generation == generation and built.default_model == "custom/m" +@pytest.mark.parametrize( + ("provider_name", "model", "names"), + [ + ("openrouter", "openrouter/openai/gpt-4o-mini", ("OPENROUTER_API_KEY",)), + ("zai", "zai/glm-4.6", ("ZAI_API_KEY", "ZHIPUAI_API_KEY")), + ], +) +def test_building_an_agents_provider_leaves_the_hosts_environment_alone( + monkeypatch: pytest.MonkeyPatch, provider_name: str, model: str, names: tuple[str, ...] +) -> None: + """The host's own key stays its own, and nothing it starts inherits the agent's. + + Every call passes the key explicitly, so the provider built for an agent + needs nothing in the environment; the agent's own process still exports it. + """ + from raven.providers.litellm_provider import LiteLLMProvider + + # The agent's side, which exports the key in its own process as before. + spec = relay_provider.RelayProvider(default_model=model, provider_name=provider_name, api_key="agent-key").spec() + monkeypatch.setenv(names[0], "host-key") + for name in names[1:]: + monkeypatch.delenv(name, raising=False) + built = relay_server.build_litellm_provider(spec) + + assert built.api_key == "agent-key" + assert os.environ[names[0]] == "host-key" + assert all(name not in os.environ for name in names[1:]) + assert relay_server.RelayServer()._build is relay_server.build_litellm_provider + + monkeypatch.delenv(names[0]) + LiteLLMProvider(default_model=model, provider_name=provider_name, api_key="agent-key") + assert os.environ[names[0]] == "agent-key" + + async def test_frames_that_are_not_frames_are_refused() -> None: with pytest.raises(relay_wire.FrameError): relay_wire._unpack(b"\xff") From 4d29d4a536dd82a2196e46ce6f9667499d9dad3f Mon Sep 17 00:00:00 2001 From: Dizhan Xue Date: Wed, 28 Oct 2026 06:09:58 +0000 Subject: [PATCH 10/19] fix(providers): read the relay token before any frame The relay server read a whole frame, up to 512 MiB and with no deadline, before it compared the token. On loopback TCP, where any local account can connect, a peer without the token could hold connections open and make the host buffer large frames. A client now sends the token first, alone on a line, and the server reads exactly that many bytes under a five-second deadline before it reads a frame. A silent peer is dropped, and a wrong token is refused without anything more being read. --- raven/providers/relay.py | 8 ++++-- raven/providers/relay_server.py | 27 ++++++++++++++++--- raven/providers/relay_wire.py | 9 +++++++ tests/test_provider_relay.py | 48 +++++++++++++++++++++++++++------ 4 files changed, 78 insertions(+), 14 deletions(-) diff --git a/raven/providers/relay.py b/raven/providers/relay.py index 699f91438..990a795b9 100644 --- a/raven/providers/relay.py +++ b/raven/providers/relay.py @@ -142,7 +142,7 @@ def request_frame(op: str, **fields: Any) -> dict[str, Any]: found = target() if found is None: raise RelayUnavailableError("no relay was offered to this process") - return {"token": found[1], "op": op, **fields} + return {"op": op, **fields} def _unreachable_kind(kind: str) -> RelayUnavailableError | None: @@ -171,6 +171,7 @@ def _connect_sync(timeout: float) -> socket.socket: sock.connect(where) else: sock = socket.create_connection(where, timeout=_CONNECT_TIMEOUT_S) + sock.sendall(relay_wire.greeting(found[1])) except OSError as exc: raise RelayUnavailableError(str(exc), connecting=True) from exc sock.settimeout(timeout) @@ -191,9 +192,12 @@ async def connect() -> tuple[asyncio.StreamReader, asyncio.StreamWriter]: else: host, port = where # type: ignore[misc] opening = asyncio.open_connection(host, port, limit=limit) - return await asyncio.wait_for(opening, timeout=_CONNECT_TIMEOUT_S) + reader, writer = await asyncio.wait_for(opening, timeout=_CONNECT_TIMEOUT_S) except (OSError, asyncio.TimeoutError) as exc: raise RelayUnavailableError(str(exc), connecting=True) from exc + # The token goes first, alone: the host reads no frame from a peer without it. + writer.write(relay_wire.greeting(found[1])) + return reader, writer def read_answer(reply: dict[str, Any] | None, *, sentinel: object | None = None) -> Any: diff --git a/raven/providers/relay_server.py b/raven/providers/relay_server.py index b3db1ea8e..a9eeefc25 100644 --- a/raven/providers/relay_server.py +++ b/raven/providers/relay_server.py @@ -6,7 +6,10 @@ (Windows) -- and answering only requests carrying the token it handed out in the launch environment. -Each request is one connection: open, one frame in, the answer out, close. A +Each request is one connection: open, the token, one frame in, the answer out, +close. The token comes first, on its own and under a deadline, so a peer +without it -- any local account, on loopback TCP -- is dropped before a frame +is read or buffered. A stream answers with a frame per delta and a closing frame, and a connection the agent closes early is how a cancelled turn reaches the call: the host stops it. @@ -48,6 +51,9 @@ #: model, a fallback); the bound is for a host that serves many agents. _PROVIDER_CACHE_SIZE = 32 +#: How long a new connection has to present the token before it is dropped. +_ADMIT_TIMEOUT_S = 5.0 + ProviderBuilder = Callable[[dict[str, Any]], LLMProvider] @@ -153,14 +159,27 @@ async def _provider(self, spec: Any) -> LLMProvider: self._providers.popitem(last=False) return provider + async def _admit(self, reader: asyncio.StreamReader) -> bool | None: + """Whether the peer presented this relay's token: ``None`` when it said nothing in time.""" + expected = relay_wire.greeting(self._token) + try: + async with asyncio.timeout(_ADMIT_TIMEOUT_S): + presented = await reader.readexactly(len(expected)) + except (TimeoutError, asyncio.IncompleteReadError, ConnectionError): + return None + return hmac.compare_digest(presented, expected) + async def _serve(self, reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None: try: - request = await relay_wire.read_frame(reader) - if request is None: + admitted = await self._admit(reader) + if admitted is None: return - if not hmac.compare_digest(str(request.get("token") or ""), self._token): + if not admitted: await relay_wire.write_frame(writer, {"fault": "this relay does not serve that token"}) return + request = await relay_wire.read_frame(reader) + if request is None: + return op = request.get("op") if op == "stream": await self._stream(request, reader, writer) diff --git a/raven/providers/relay_wire.py b/raven/providers/relay_wire.py index 76c6bbeb3..63ccf72de 100644 --- a/raven/providers/relay_wire.py +++ b/raven/providers/relay_wire.py @@ -93,6 +93,15 @@ def _length(header: bytes) -> int: return size +def greeting(token: str) -> bytes: + """What a client sends before its frame: the token alone, on a line. + + A host reads exactly this many bytes, under a deadline, before it reads a + frame, so a peer without the token never gets a frame read or buffered. + """ + return token.encode("utf-8") + b"\n" + + async def write_frame(writer: asyncio.StreamWriter, message: dict[str, Any]) -> None: writer.write(_pack(message)) await writer.drain() diff --git a/tests/test_provider_relay.py b/tests/test_provider_relay.py index 35a611b09..87387f2b5 100644 --- a/tests/test_provider_relay.py +++ b/tests/test_provider_relay.py @@ -278,6 +278,42 @@ async def test_a_request_without_the_token_is_refused(served, monkeypatch: pytes await provider.chat(messages=[{"role": "user", "content": "q"}]) +async def _open(server: relay_server.RelayServer) -> tuple[asyncio.StreamReader, asyncio.StreamWriter]: + from raven.providers.relay import _split + + kind, where = _split(server.env()[relay.RELAY_ENV]) + if kind == "unix": + return await asyncio.open_unix_connection(where) + return await asyncio.open_connection(*where) + + +async def test_a_peer_without_the_token_is_dropped_before_the_host_reads_a_frame( + served, monkeypatch: pytest.MonkeyPatch +) -> None: + """On loopback TCP any local account can connect: silence is dropped, and nothing past the token's length is read. + + The second peer announces a frame far larger than it sends; a host that read + frames before tokens would sit waiting to buffer it. + """ + _, server = served + monkeypatch.setattr(relay_server, "_ADMIT_TIMEOUT_S", 0.2) + + reader, writer = await _open(server) + try: + assert await asyncio.wait_for(reader.read(), timeout=5) == b"" + finally: + writer.close() + + reader, writer = await _open(server) + try: + writer.write((256 * 1024 * 1024).to_bytes(4, "big") + b"x" * 4096) + await writer.drain() + reply = await asyncio.wait_for(relay_wire.read_frame(reader), timeout=5) + assert reply is not None and "token" in reply["fault"] + finally: + writer.close() + + async def test_an_unreachable_host_leaves_the_call_to_this_process(monkeypatch: pytest.MonkeyPatch, tmp_path) -> None: monkeypatch.setenv(relay.RELAY_ENV, f"unix:{tmp_path / 'gone.sock'}") monkeypatch.setenv(relay.RELAY_TOKEN_ENV, "t") @@ -612,17 +648,12 @@ async def test_the_host_builds_one_provider_per_construction_and_keeps_the_lates async def _raw(server: relay_server.RelayServer, frame: dict | None) -> dict | None: - from raven.providers.relay import _split - - kind, where = _split(server.env()[relay.RELAY_ENV]) - if kind == "unix": - reader, writer = await asyncio.open_unix_connection(where) - else: - reader, writer = await asyncio.open_connection(*where) + reader, writer = await _open(server) try: if frame is None: return None - await relay_wire.write_frame(writer, {"token": server.env()[relay.RELAY_TOKEN_ENV], **frame}) + writer.write(relay_wire.greeting(server.env()[relay.RELAY_TOKEN_ENV])) + await relay_wire.write_frame(writer, frame) return await relay_wire.read_frame(reader) finally: writer.close() @@ -676,6 +707,7 @@ async def test_a_stream_the_host_ends_oddly_or_drops_is_a_failure_here(monkeypat replies = iter([{"what": 1}, None]) async def host(reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None: + assert await reader.readline() == relay_wire.greeting("t") await relay_wire.read_frame(reader) reply = next(replies, None) if reply is not None: From a5f6146b7ce9eee1659a48477c6b40e6fd99a85a Mon Sep 17 00:00:00 2001 From: Dizhan Xue Date: Wed, 28 Oct 2026 06:10:13 +0000 Subject: [PATCH 11/19] fix(acp_client): close idle agents only from the loop that launched them The pool kept one idle-check task and restarted it on whichever loop acquired next, and close_idle closed every connection held. When the RPC server ran a CLI command such as raven agent -m on a worker thread, that loop's check could close the main loop's agents, killing them after a stalled teardown and binding the pool's per-name locks to the worker loop. Each connection now records the loop it was made on, each loop runs its own idle check, and a check takes the locks of and closes only its own loop's connections. --- raven/acp_client/pool.py | 51 ++++++++++++++++++++++++------- tests/test_acp_pool_idle_close.py | 28 +++++++++++++++-- 2 files changed, 66 insertions(+), 13 deletions(-) diff --git a/raven/acp_client/pool.py b/raven/acp_client/pool.py index e4765808f..3cf134a43 100644 --- a/raven/acp_client/pool.py +++ b/raven/acp_client/pool.py @@ -22,6 +22,7 @@ from __future__ import annotations import asyncio +import contextlib import hashlib import json import time @@ -305,6 +306,18 @@ def dispatch(self, params: dict[str, Any]) -> bool: return True +def _running_loop() -> asyncio.AbstractEventLoop | None: + try: + return asyncio.get_running_loop() + except RuntimeError: + return None + + +def _on(connections: Any, loop: asyncio.AbstractEventLoop) -> bool: + """Whether any of ``connections`` belongs to ``loop`` (one made outside any loop belongs to every one).""" + return any(connection.loop is None or connection.loop is loop for connection in connections) + + class _Connection: """A client, the router whose lifetime is tied to it, and one lock per session.""" @@ -342,6 +355,9 @@ def __init__( actually serving requests, where a stored snapshot only describes the one that was verified.""" self._session_locks: dict[str, asyncio.Lock] = {} + self.loop = _running_loop() + """The loop this connection was made on. Its client's pipes are that loop's, + so only that loop's idle check may close it.""" self.leased_at = time.monotonic() """When the pool last handed this connection out. A caller sends its request right after taking it, so the idle check leaves a connection @@ -383,7 +399,7 @@ class AcpConnectionPool: def __init__(self) -> None: self._connections: dict[str, dict[str, _Connection]] = {} self._locks: dict[str, asyncio.Lock] = {} - self._idle_task: asyncio.Task | None = None + self._idle_tasks: dict[asyncio.AbstractEventLoop, asyncio.Task] = {} def connections(self, name: str) -> list[_Connection]: """Every connection held under ``name``, live or not.""" @@ -669,17 +685,21 @@ async def close_idle(self, budget: float | None = None) -> list[str]: ``budget`` defaults to the host's live setting. Each candidate is judged again under its agent's lock, the one ``acquire`` holds while it hands a connection out, so a connection being reused is never closed - underneath its caller. + underneath its caller. Only the running loop's connections are + considered: another loop's are that loop's to close. """ budget = self._idle_budget() if budget is None else budget if budget is None: return [] + loop = asyncio.get_running_loop() closing: list[tuple[str, _Connection]] = [] - for name in list(self._connections): + for name in [name for name, held in self._connections.items() if _on(held.values(), loop)]: async with self._lock_for(name): held = self._connections.get(name, {}) now = time.monotonic() for bkey, connection in list(held.items()): + if not _on((connection,), loop): + continue # A connection whose agent already ended goes too: nothing # would hand it out again, and holding it keeps this check running. if not connection.alive or self._closable(connection, budget, now): @@ -694,15 +714,18 @@ async def close_idle(self, budget: float | None = None) -> list[str]: return [name for name, _ in closing] def _watch_idle(self) -> None: - """Start the idle check if it is not running; it ends once nothing is held.""" + """Start this loop's idle check if it is not running; it ends once the loop holds nothing.""" loop = asyncio.get_running_loop() - task = self._idle_task - if task is not None and not task.done() and task.get_loop() is loop: + for other in [other for other in self._idle_tasks if other is not loop and other.is_closed()]: + del self._idle_tasks[other] + task = self._idle_tasks.get(loop) + if task is not None and not task.done(): return - self._idle_task = loop.create_task(self._idle_loop(), name="acp-idle-close") + self._idle_tasks[loop] = loop.create_task(self._idle_loop(), name="acp-idle-close") async def _idle_loop(self) -> None: - while self._connections: + loop = asyncio.get_running_loop() + while any(_on(held.values(), loop) for held in list(self._connections.values())): budget = self._idle_budget() share = budget / _IDLE_CHECK_SHARE if budget is not None else _IDLE_CHECK_MAX_S await asyncio.sleep(min(max(share, _IDLE_CHECK_MIN_S), _IDLE_CHECK_MAX_S)) @@ -720,9 +743,15 @@ async def close_all(self) -> None: request runs. Serially and unbounded, one such connection held every other one -- and the exit waiting behind them -- for that whole time. """ - task, self._idle_task = self._idle_task, None - if task is not None and not task.get_loop().is_closed(): - task.cancel() + tasks, self._idle_tasks = self._idle_tasks, {} + running = asyncio.get_running_loop() + for loop, task in tasks.items(): + if loop is running: + task.cancel() + elif not loop.is_closed(): + # Another thread's loop: its task is cancelled from that loop. + with contextlib.suppress(RuntimeError): + loop.call_soon_threadsafe(task.cancel) names = list(self._connections) if not names: return diff --git a/tests/test_acp_pool_idle_close.py b/tests/test_acp_pool_idle_close.py index dcc4d3b55..63cf19710 100644 --- a/tests/test_acp_pool_idle_close.py +++ b/tests/test_acp_pool_idle_close.py @@ -137,7 +137,7 @@ async def test_the_idle_check_starts_with_a_connection_and_ends_with_the_last(mo client = _Client(idle=0) pool = _pool_holding(soon=_held(client, _LOADS, leased_ago=0)) pool._watch_idle() - task = pool._idle_task + task = pool._idle_tasks.get(asyncio.get_running_loop()) assert task is not None await asyncio.wait_for(task, timeout=5) @@ -230,6 +230,30 @@ def test_a_runtime_edit_that_is_not_a_mapping_is_rejected_not_removed( assert live_mod.subagent_runtime(reader).idle_close_s == 600 +async def test_an_idle_check_on_another_loop_leaves_this_loops_connections(monkeypatch: pytest.MonkeyPatch) -> None: + """The RPC server runs a CLI command's own ``asyncio.run`` on a worker thread, beside its own loop. + + A connection's pipes belong to the loop that launched it, so only that + loop's check may close it. + """ + monkeypatch.setattr(pool_mod, "_IDLE_CHECK_MIN_S", 0.01) + monkeypatch.setattr(AcpConnectionPool, "_idle_budget", staticmethod(lambda: 0.04)) + here = _Client(idle=3600) + pool = _pool_holding(here=_held(here, _LOADS)) + + async def beside() -> list[str]: + closed = await pool.close_idle(0.04) + pool._watch_idle() + await asyncio.sleep(0.2) + return closed + + assert await asyncio.to_thread(asyncio.run, beside()) == [] + assert not here.closed and "here" in pool._connections + + assert await pool.close_idle(0.04) == ["here"] + assert here.closed + + async def test_a_failed_idle_check_does_not_end_the_next_one(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setattr(pool_mod, "_IDLE_CHECK_MIN_S", 0.01) monkeypatch.setattr(AcpConnectionPool, "_idle_budget", staticmethod(lambda: 0.04)) @@ -246,5 +270,5 @@ async def flaky(budget=None): monkeypatch.setattr(pool, "close_idle", flaky) pool._watch_idle() - await asyncio.wait_for(pool._idle_task, timeout=5) + await asyncio.wait_for(pool._idle_tasks[asyncio.get_running_loop()], timeout=5) assert len(calls) >= 2 and client.closed From a84e2355acbffa19b86f00beb5320cd1599e7664 Mon Sep 17 00:00:00 2001 From: Dizhan Xue Date: Wed, 28 Oct 2026 06:10:21 +0000 Subject: [PATCH 12/19] fix(providers): size a relayed call under the agent's turn A relayed chat_stream resolved max_tokens on the host, through send_max_tokens and the model binding active there. The relay server's handlers ran in a copy of the context the server was started in, so that binding was none or the host turn that first delegated, never the agent's: an agent whose loop reserved 16384 sent 32768. The agent now sends its turn's context window with each call, and the server serves each connection in a fresh contextvars.Context and makes the call under a binding with that window, so both sides size the ceiling with the same function and the same inputs. Measured across two processes with the real provider: 16384 relayed and 16384 direct. CONTEXT.md's LLM relay entry also states the token-first and environment guarantees. --- CONTEXT.md | 20 ++++++++------ raven/providers/relay_provider.py | 10 ++++++- raven/providers/relay_server.py | 38 +++++++++++++++++++++++--- tests/test_provider_relay.py | 44 +++++++++++++++++++++++++++++++ 4 files changed, 99 insertions(+), 13 deletions(-) diff --git a/CONTEXT.md b/CONTEXT.md index 8212b1e07..ad6891be7 100644 --- a/CONTEXT.md +++ b/CONTEXT.md @@ -3401,15 +3401,19 @@ stateful instance's next turn reopens its session in a fresh process with `sessi _Avoid_: calling it a timeout - nothing in flight is cut short. **LLM relay** (`raven/providers/relay.py`, `relay_provider.py`, `relay_server.py`, `subagents.runtime.llmRelay`): -The host making a Raven agent's LiteLLM calls for it. The host offers a local socket and a -token in the launch environment (`RAVEN_LLM_RELAY`), and an agent that reads them builds a +The host making a Raven agent's LiteLLM calls for it. The host offers a local socket and a token +in the launch environment (`RAVEN_LLM_RELAY`), and an agent that reads them builds a `RelayProvider` where it would have built a `LiteLLMProvider`: the same construction and the -same answers (both stand on `LiteLLMShape`), with only the single `chat` / `chat_stream` -attempt and the `relayed` catalogue lookups crossing to the host. Retry, classification, usage -records and traces stay in the agent; the host's calls run under `recorded_by_caller` so a call -is billed once. The relay is served by the event loop that started it: a later loop starts -another, and a loop running beside it in another thread is offered the same socket. An agent -that cannot reach the host loads LiteLLM and works as before. +same answers (both stand on `LiteLLMShape`), with only the single `chat` / `chat_stream` attempt +and the `relayed` catalogue lookups crossing to the host. Retry, classification, usage records +and traces stay in the agent; the host's calls run under `recorded_by_caller` so a call is +billed once. The relay is served by the event loop that started it: a later loop starts another, +and a loop running beside it in another thread is offered the same socket. A connection presents +the token first, alone and under a deadline, so a peer without it is dropped before a frame is +read. The host builds an agent's provider without writing the agent's key into its own +environment, and makes each call in a context of its own under the window of the agent's turn, +so the output ceiling on the wire is the one the agent's loop reserved. An agent that cannot +reach the host loads LiteLLM and works as before. **Fork template** (`raven/acp_client/fork_template.py`, `raven/cli/_fork_template.py`, `subagents.runtime.forkTemplate`): On Linux, one process per host that has imported what Raven agents import, frozen out of the diff --git a/raven/providers/relay_provider.py b/raven/providers/relay_provider.py index 1c17d1b4f..ae644b58d 100644 --- a/raven/providers/relay_provider.py +++ b/raven/providers/relay_provider.py @@ -12,6 +12,7 @@ from raven.contracts.llm_provider import ChatDelta, GenerationSettings, LLMResponse from raven.providers import relay_wire from raven.providers.base import LLMProvider +from raven.providers.binding import active_binding from raven.providers.litellm_shape import LiteLLMShape from raven.providers.relay import ( RelayFaultError, @@ -84,7 +85,14 @@ def _locally(self) -> LLMProvider: return local def _call(self, **args: Any) -> dict[str, Any]: - return {"spec": self.spec(), "args": relay_wire.encode(args, sentinel=self._SENTINEL)} + # The running turn's window travels with the call: the host sizes the + # output ceiling under it, as this process would have. + bound = active_binding() + return { + "spec": self.spec(), + "args": relay_wire.encode(args, sentinel=self._SENTINEL), + "window": bound.context_window if bound is not None else None, + } async def chat( self, diff --git a/raven/providers/relay_server.py b/raven/providers/relay_server.py index a9eeefc25..1f3cd8f95 100644 --- a/raven/providers/relay_server.py +++ b/raven/providers/relay_server.py @@ -29,6 +29,7 @@ import asyncio import atexit import contextlib +import contextvars import hashlib import hmac import json @@ -45,6 +46,7 @@ from raven.providers import relay_wire, usage_record from raven.providers.base import LLMProvider +from raven.providers.binding import ModelBinding, use_binding from raven.providers.relay import RELAY_ENV, RELAY_TOKEN_ENV, RelayFaultError, run_relayed #: Distinct constructions kept built. An agent normally has one or two (its @@ -82,6 +84,7 @@ def __init__(self, *, build: ProviderBuilder = build_litellm_provider, tcp: bool self._dir: Path | None = None self._providers: OrderedDict[str, LLMProvider] = OrderedDict() self._building: dict[str, asyncio.Future[LLMProvider]] = {} + self._handlers: set[asyncio.Task[None]] = set() self.loop: asyncio.AbstractEventLoop | None = None """The loop that answers this server's connections, and the only one that may close it.""" @@ -96,7 +99,7 @@ def env(self) -> dict[str, str]: async def start(self) -> None: self.loop = asyncio.get_running_loop() if self._tcp: - self._server = await asyncio.start_server(self._serve, host="127.0.0.1", port=0) + self._server = await asyncio.start_server(self._accept, host="127.0.0.1", port=0) host, port = self._server.sockets[0].getsockname()[:2] self._address = f"tcp:{host}:{port}" else: @@ -104,7 +107,7 @@ async def start(self) -> None: # from reaching the socket at all; the token is the second lock. self._dir = Path(tempfile.mkdtemp(prefix="raven-llm-relay-")) path = self._dir / "relay.sock" - self._server = await asyncio.start_unix_server(self._serve, path=str(path)) + self._server = await asyncio.start_unix_server(self._accept, path=str(path)) self._address = f"unix:{path}" atexit.register(shutil.rmtree, self._dir, True) logger.debug("llm relay: serving at {}", self._address) @@ -159,6 +162,33 @@ async def _provider(self, spec: Any) -> LLMProvider: self._providers.popitem(last=False) return provider + def _accept(self, reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None: + """Serve one connection in a context of its own. + + A handler would otherwise run in a copy of the context :meth:`start` ran + in, which may be inside a host turn: that turn's model binding would size + an agent's call, and its other state would follow the call too. + """ + task = asyncio.get_running_loop().create_task(self._serve(reader, writer), context=contextvars.Context()) + self._handlers.add(task) + task.add_done_callback(self._handlers.discard) + + @staticmethod + def _turn( + request: dict[str, Any], provider: LLMProvider, args: dict[str, Any] + ) -> contextlib.AbstractContextManager: + """The agent's turn as far as sizing a call goes: its window, under the model the call names. + + ``send_max_tokens`` reads the window from the active binding on both sides, + so the request body carries the ceiling the agent's loop reserved. No + window means the agent called outside a turn, and so does the host. + """ + window = request.get("window") + if not isinstance(window, int) or isinstance(window, bool) or window <= 0: + return contextlib.nullcontext() + model = str(args.get("model") or provider.get_default_model()) + return use_binding(ModelBinding(provider=provider, model=model, configured_window=window)) + async def _admit(self, reader: asyncio.StreamReader) -> bool | None: """Whether the peer presented this relay's token: ``None`` when it said nothing in time.""" expected = relay_wire.greeting(self._token) @@ -234,7 +264,7 @@ async def _stream( args = relay_wire.decode(request.get("args") or {}, sentinel=provider._SENTINEL) async def pump() -> None: - with usage_record.recorded_by_caller(): + with usage_record.recorded_by_caller(), self._turn(request, provider, args): async for delta in provider.chat_stream(**args): await relay_wire.write_frame(writer, {"delta": relay_wire.encode(delta)}) await relay_wire.write_frame(writer, {"done": True}) @@ -248,7 +278,7 @@ async def _chat(self, request: dict[str, Any], reader: asyncio.StreamReader, wri args = relay_wire.decode(request.get("args") or {}, sentinel=provider._SENTINEL) async def answer() -> None: - with usage_record.recorded_by_caller(): + with usage_record.recorded_by_caller(), self._turn(request, provider, args): response = await provider.chat(**args) await relay_wire.write_frame(writer, {"result": relay_wire.encode(response)}) diff --git a/tests/test_provider_relay.py b/tests/test_provider_relay.py index 87387f2b5..db412f331 100644 --- a/tests/test_provider_relay.py +++ b/tests/test_provider_relay.py @@ -28,6 +28,7 @@ ) from raven.providers import call_record, prompt_cache, relay, relay_provider, relay_server, relay_wire, usage_record from raven.providers.base import LLMProvider +from raven.providers.binding import ModelBinding, active_binding, use_binding from raven.providers.first_byte import FirstByteTimeoutError, StreamIdleTimeoutError from raven.providers.litellm_shape import LiteLLMShape @@ -165,11 +166,17 @@ def __init__(self, spec: dict[str, Any], script: dict[str, Any]) -> None: def get_default_model(self) -> str: return "m" + def _note_binding(self) -> None: + bound = active_binding() + self.script.setdefault("bindings", []).append(None if bound is None else (bound.model, bound.context_window)) + async def chat(self, **kwargs: Any) -> LLMResponse: + self._note_binding() self.script.setdefault("calls", []).append(("chat", kwargs, usage_record._owned.get())) return LLMResponse(content=f"answered {kwargs['messages'][-1]['content']}", usage={"prompt_tokens": 1}) async def chat_stream(self, **kwargs: Any) -> AsyncIterator[ChatDelta]: + self._note_binding() self.script.setdefault("calls", []).append(("stream", kwargs, usage_record._owned.get())) try: for piece in self.script.get("pieces", ["a", "b"]): @@ -222,6 +229,43 @@ async def test_a_turn_streams_through_the_host(served) -> None: assert owned, "the host's call must be marked recorded by its caller, the agent" +async def test_a_relayed_call_is_sized_under_the_agents_turn(served) -> None: + """The output ceiling is the running turn's: the host sees the agent's window, as the agent's own call would. + + ``send_max_tokens`` reads the window from the active binding, so a relayed + stream sized under the host's context would carry a different number than + the one the agent's loop reserved. + """ + script, _ = served + provider = relay_provider.RelayProvider(default_model="custom/m", provider_name="custom", api_base="http://x/v1") + turn = ModelBinding(provider=provider, model="custom/m", configured_window=32768) + + with use_binding(turn): + [d async for d in provider.chat_stream(messages=[{"role": "user", "content": "q"}], model="custom/m")] + await provider.chat(messages=[{"role": "user", "content": "q"}], model="custom/m") + [d async for d in provider.chat_stream(messages=[{"role": "user", "content": "q"}], model="custom/m")] + + assert script["bindings"] == [("custom/m", 32768), ("custom/m", 32768), None] + + +async def test_the_host_answers_in_a_context_of_its_own(monkeypatch: pytest.MonkeyPatch) -> None: + """A relay started inside a host turn: none of that turn reaches an agent's call.""" + script: dict[str, Any] = {"specs": []} + host_turn = ModelBinding(provider=_Wire({}, script), model="host-model", configured_window=1_000_000) + with use_binding(host_turn): + server = relay_server.RelayServer(build=lambda spec: _Wire(spec, script)) + await server.start() + try: + for key, value in server.env().items(): + monkeypatch.setenv(key, value) + relay._reset_for_tests() + provider = relay_provider.RelayProvider(default_model="custom/m", provider_name="custom") + [d async for d in provider.chat_stream(messages=[{"role": "user", "content": "q"}], model="custom/m")] + finally: + await server.close() + assert script["bindings"] == [None] + + async def test_a_single_call_answers_through_the_host(served) -> None: provider = relay_provider.RelayProvider(default_model="custom/m", provider_name="custom") response = await provider.chat(messages=[{"role": "user", "content": "q"}]) From 3f49649f460cca3cc873f699ea5a1a471bc52af6 Mon Sep 17 00:00:00 2001 From: Dizhan Xue Date: Wed, 28 Oct 2026 06:10:32 +0000 Subject: [PATCH 13/19] fix(*): ask a raven agent for resident work before closing it idle Idle close judged an agent from the wire alone, so a Raven agent that had armed a cron wake or was running a plugin service looked idle and was closed after 600 s: the wake never fired, and the service died with the process. Raven agents now declare raven.residentWork in agentCapabilities._meta and answer _raven/residentWork from their own state: the enabled cron jobs of their partition while the scheduler runs, and the plugin services started. The pool asks only an agent that declared it, only once closing it is otherwise settled, and keeps it while it says yes. An agent that does not answer in time, or never declared it, is closed as before. The default stays 600 s. Resident work is a new term in CONTEXT.md, and the config field and the agent integrations page say the same. --- CONTEXT.md | 23 +++++-- docs-site/docs/agent-integrations.md | 2 +- docs-site/docs/agent-integrations.zh.md | 2 +- raven/acp/capabilities.py | 9 ++- raven/acp/methods.py | 16 +++++ raven/acp/protocol.py | 8 +++ raven/acp_client/capabilities.py | 4 ++ raven/acp_client/pool.py | 33 +++++++++- raven/acp_client/protocol.py | 7 ++ raven/agent/loop/wiring.py | 5 ++ raven/config/schema.py | 5 +- .../schedulers/cron/service.py | 10 +++ tests/acp_stub_server.py | 29 ++++++++ tests/test_acp_methods.py | 41 ++++++++++++ tests/test_acp_pool_idle_close.py | 66 ++++++++++++++++++- 15 files changed, 246 insertions(+), 14 deletions(-) diff --git a/CONTEXT.md b/CONTEXT.md index ad6891be7..cd1b97dce 100644 --- a/CONTEXT.md +++ b/CONTEXT.md @@ -3394,11 +3394,24 @@ or `transcript.jsonl`, not this. **Idle close** (`AcpConnectionPool.close_idle`, `subagents.runtime.idleCloseS`): The pool ending a pooled ACP agent's process once the connection has been quiet for the -configured budget: no request of the host's unanswered (a held wake is one), none of the -agent's being answered, no frame from the agent and no lease for the whole budget. Only an -agent that advertises `loadSession`, or one that never opened a session, is closed, so a -stateful instance's next turn reopens its session in a fresh process with `session/load`. -_Avoid_: calling it a timeout - nothing in flight is cut short. +configured budget: no request of the host's unanswered (a held wake is one), none of the agent's +being answered, no frame from the agent and no lease for the whole budget. Only an agent that +advertises `loadSession`, or one that never opened a session, is closed, so a stateful +instance's next turn reopens its session in a fresh process with `session/load`. An agent that +declares `raven.residentWork` is asked first and kept while it holds Resident work. Each event +loop closes only the connections made on it, whose pipes are its own. +_Avoid_: calling it a timeout - nothing in flight is cut short, nor anything held for later. + +**Resident work** (`_raven/residentWork`, `CronService.armed_jobs`, `running_plugin_services`): +What a Raven agent holds that runs with nobody asking: an enabled cron job of its own partition +while its scheduler runs, and a running plugin Service. Neither puts anything on the wire, and +closing the process loses both: no other partition claims the job, and the Service lives only as +long as the process. The agent declares `raven.residentWork` in `agentCapabilities._meta` and +answers `_raven/residentWork` as things stand at the moment of asking, so a one-shot wake that +has fired no longer counts. The host decides on it, so it is a protocol extension the host asks, +not Response Meta. +_Avoid_: a held wake (`raven.holdTurn`) for it - that keeps a prompt open and the agent busy; +resident work is what an agent holds between prompts. **LLM relay** (`raven/providers/relay.py`, `relay_provider.py`, `relay_server.py`, `subagents.runtime.llmRelay`): The host making a Raven agent's LiteLLM calls for it. The host offers a local socket and a token diff --git a/docs-site/docs/agent-integrations.md b/docs-site/docs/agent-integrations.md index 35cc3ce1e..85278c2de 100644 --- a/docs-site/docs/agent-integrations.md +++ b/docs-site/docs/agent-integrations.md @@ -82,7 +82,7 @@ applies to the next launch. | Key | Default | What it does | | --- | --- | --- | -| `idleCloseS` | `600` | Closes an agent's process once it has been idle this many seconds: no request in flight either way and nothing said. Only an agent that can reload its sessions is closed, so the next turn resumes the conversation in a fresh process. `0` or `null` keeps agents until the host exits. | +| `idleCloseS` | `600` | Closes an agent's process once it has been idle this many seconds: no request in flight either way and nothing said. An agent that still holds an armed reminder or a running plugin service is kept. Only an agent that can reload its sessions is closed, so the next turn resumes the conversation in a fresh process. `0` or `null` keeps agents until the host exits. | | `llmRelay` | `true` | The host makes a native agent's LiteLLM calls, so the agent never loads LiteLLM. The agent still builds each request, retries and records usage itself. | | `forkTemplate` | `true` | On Linux, native agents are forked from one process that has already loaded Raven, and share those pages instead of each loading its own copy. Other systems launch agents as before. | diff --git a/docs-site/docs/agent-integrations.zh.md b/docs-site/docs/agent-integrations.zh.md index ca2ce5382..f43872e61 100644 --- a/docs-site/docs/agent-integrations.zh.md +++ b/docs-site/docs/agent-integrations.zh.md @@ -70,7 +70,7 @@ launcher。readiness probe 成功或收到委派回执，都不代表任务已 | 键 | 默认值 | 作用 | | --- | --- | --- | -| `idleCloseS` | `600` | Agent 空闲达到该秒数（双方都没有在途请求，也没有任何消息）后关闭其进程。只关闭能重新加载会话的 Agent，因此下一轮会在新进程中续上原对话。`0` 或 `null` 表示一直保留到宿主退出。 | +| `idleCloseS` | `600` | Agent 空闲达到该秒数（双方都没有在途请求，也没有任何消息）后关闭其进程。仍持有已设定的提醒或正在运行的插件服务的 Agent 会被保留。只关闭能重新加载会话的 Agent，因此下一轮会在新进程中续上原对话。`0` 或 `null` 表示一直保留到宿主退出。 | | `llmRelay` | `true` | 由宿主代为发出原生 Agent 的 LiteLLM 调用，Agent 自身不再加载 LiteLLM。请求组装、重试和用量记录仍由 Agent 自己完成。 | | `forkTemplate` | `true` | 在 Linux 上，原生 Agent 从一个已加载 Raven 的进程 fork 出来，共享这些内存页，而不是各自加载一份。其他系统仍按原方式启动。 | diff --git a/raven/acp/capabilities.py b/raven/acp/capabilities.py index b13089e26..2682bd2a3 100644 --- a/raven/acp/capabilities.py +++ b/raven/acp/capabilities.py @@ -81,7 +81,14 @@ def agent_capabilities() -> dict[str, Any]: # turn already in flight -- a person speaking to an agent whose prompt # is held open while it waits on its own scheduled work. Advertised by # presence, which is how raven's own client reads it. - "_meta": {protocol.SESSION_MCP_CAPABILITY: {}, protocol.STEER_CAPABILITY: True}, + # And ``_raven/residentWork``: whether this process holds work nobody is + # asking for yet (an armed cron job, a running plugin service), asked + # by raven's client before it closes the agent for idleness. + "_meta": { + protocol.SESSION_MCP_CAPABILITY: {}, + protocol.STEER_CAPABILITY: True, + protocol.RESIDENT_WORK_CAPABILITY: True, + }, # ``list``, ``resume``, ``close`` and ``delete``, and an empty object is # how the schema spells "supported". Each one is declared because the # method behind it exists and is honoured: ``_session_resume`` reopens a diff --git a/raven/acp/methods.py b/raven/acp/methods.py index b96fadd59..5861dec73 100644 --- a/raven/acp/methods.py +++ b/raven/acp/methods.py @@ -284,12 +284,28 @@ async def _route(self, method: str, params: dict[str, Any], *, is_request: bool) return await self._session_cancel(params) if method == protocol.STEER_METHOD: return await self._session_steer(params) + if method == protocol.RESIDENT_WORK_METHOD: + return self._resident_work() if method in UNIMPLEMENTED_METHODS: raise AcpMethodError(protocol.METHOD_NOT_FOUND, f"{method} is not implemented") raise AcpMethodError(protocol.METHOD_NOT_FOUND, f"unknown method {method}") # -- handlers --------------------------------------------------------- + def _resident_work(self) -> dict[str, Any]: + """What this agent holds that runs with nobody asking (``_raven/residentWork``). + + An armed cron job of this agent's partition fires only from this + process, and a plugin service lives only as long as it does, so a client + about to close the agent for idleness asks first. Read at the moment of + asking: a one-shot wake that has fired no longer counts. + """ + loop = self._agent_loop + cron = getattr(loop, "cron_service", None) + jobs = len(cron.armed_jobs()) if cron is not None else 0 + services = len(getattr(loop, "running_plugin_services", ()) or ()) + return {"resident": bool(jobs or services), "cronJobs": jobs, "services": services} + def _initialize(self, params: dict[str, Any]) -> dict[str, Any]: """Answer the handshake and keep what the client declared. diff --git a/raven/acp/protocol.py b/raven/acp/protocol.py index 887f36f15..f80e42ef0 100644 --- a/raven/acp/protocol.py +++ b/raven/acp/protocol.py @@ -78,6 +78,14 @@ capability is advertised by presence under ``agentCapabilities._meta``. """ +RESIDENT_WORK_METHOD = "_raven/residentWork" +RESIDENT_WORK_CAPABILITY = "raven.residentWork" +"""raven's resident-work extension: what this agent holds that runs with nobody asking. + +Asked by raven's own client before it closes the agent for idleness; same names +there, advertised the same way as the steer extension. +""" + STOP_REASONS = frozenset( { "end_turn", diff --git a/raven/acp_client/capabilities.py b/raven/acp_client/capabilities.py index 147563194..8fb5d3388 100644 --- a/raven/acp_client/capabilities.py +++ b/raven/acp_client/capabilities.py @@ -39,6 +39,7 @@ from raven.acp_client.permissions import auto_approver from raven.acp_client.protocol import ( MENU_UNREAD_META, + RESIDENT_WORK_CAPABILITY, SESSION_MCP_CAPABILITY, STEER_CAPABILITY, AcpConnectionError, @@ -511,6 +512,8 @@ class _Handshake: can_load: bool = False can_delete: bool = False can_steer: bool = False + reports_resident_work: bool = False + """The agent answers ``_raven/residentWork``, so it can be asked before it is closed for idleness.""" mcp_http: bool = False mcp_sse: bool = False session_mcp: bool = False @@ -556,6 +559,7 @@ def _read_initialize(result: Any) -> _Handshake: # The one extension raven itself serves, announced in the schema's # extension carrier; a client that did not read it must not call it. can_steer=STEER_CAPABILITY in _dict(caps.get("_meta")), + reports_resident_work=RESIDENT_WORK_CAPABILITY in _dict(caps.get("_meta")), mcp_http=bool(mcp_caps.get("http")), mcp_sse=bool(mcp_caps.get("sse")), # The same carrier as the steer flag, and read the same way: presence, diff --git a/raven/acp_client/pool.py b/raven/acp_client/pool.py index 3cf134a43..b0fbecb3a 100644 --- a/raven/acp_client/pool.py +++ b/raven/acp_client/pool.py @@ -82,6 +82,10 @@ # the caller would otherwise spend finding out. _PROBE_TIMEOUT_S = 5.0 +# How long an agent about to be closed for idleness has to say it holds resident +# work. Answered from the agent's dispatcher like the probe, so the same budget. +_RESIDENT_ASK_TIMEOUT_S = _PROBE_TIMEOUT_S + SessionSink = Callable[[str, dict[str, Any]], Awaitable[None]] """Receives ``(method, params)`` for one session's notifications.""" @@ -679,6 +683,26 @@ def _closable(connection: _Connection, budget: float, now: float) -> bool: return False return handshake_of(connection.initialize).can_load or not connection.router.seen + @staticmethod + async def _holds_resident_work(connection: _Connection) -> bool: + """Whether the agent says it holds work that runs with nobody asking. + + An armed cron wake or a running plugin service lives in the agent's + process and puts nothing on the wire, so a connection quiet for the + whole budget can still be holding one, and closing it loses it. Asked + only of an agent that declared it answers, and only once closing it is + otherwise settled, so the answer is the one at the moment of closing. + An agent that does not answer in time is closed as before. + """ + if not handshake_of(connection.initialize).reports_resident_work: + return False + try: + answer = await connection.client.request(protocol.RESIDENT_WORK_METHOD, {}, timeout=_RESIDENT_ASK_TIMEOUT_S) + except Exception as exc: # noqa: BLE001 - an agent that cannot answer is closed as before + logger.debug("acp agent {!r}: no resident-work answer: {}", connection.client.name, exc) + return False + return isinstance(answer, dict) and bool(answer.get("resident")) + async def close_idle(self, budget: float | None = None) -> list[str]: """Close every connection idle for ``budget`` seconds; the names closed. @@ -702,9 +726,12 @@ async def close_idle(self, budget: float | None = None) -> list[str]: continue # A connection whose agent already ended goes too: nothing # would hand it out again, and holding it keeps this check running. - if not connection.alive or self._closable(connection, budget, now): - held.pop(bkey, None) - closing.append((name, connection)) + if connection.alive and not self._closable(connection, budget, now): + continue + if connection.alive and await self._holds_resident_work(connection): + continue + held.pop(bkey, None) + closing.append((name, connection)) if not held: self._connections.pop(name, None) for name, connection in closing: diff --git a/raven/acp_client/protocol.py b/raven/acp_client/protocol.py index 3f0983fb2..b2608550b 100644 --- a/raven/acp_client/protocol.py +++ b/raven/acp_client/protocol.py @@ -78,6 +78,13 @@ STEER_METHOD = "_raven/session/steer" STEER_CAPABILITY = "raven.steer" +# And what a raven agent holds that runs with nobody asking -- an armed cron job +# of its own partition, a running plugin service. The pool asks it before it +# closes the agent for idleness: the wire says nothing about either, and +# closing the process would lose both. Declared the same way, by presence. +RESIDENT_WORK_METHOD = "_raven/residentWork" +RESIDENT_WORK_CAPABILITY = "raven.residentWork" + # The other one, and a promise rather than a feature flag: a raven that declares # this says that ``mcpServers`` on ``session/new`` / ``session/load`` / # ``session/resume`` is really connected, and that the tools behind it are diff --git a/raven/agent/loop/wiring.py b/raven/agent/loop/wiring.py index 5a56c1b1f..9b81c37a0 100644 --- a/raven/agent/loop/wiring.py +++ b/raven/agent/loop/wiring.py @@ -1706,6 +1706,11 @@ def _rebind_workdir(self, session_key: str, target: "str | Path") -> Path: _plugin_services_started = False _started_services: tuple = () + @property + def running_plugin_services(self) -> tuple: + """The contributed services started and not yet stopped: work this host runs with nobody asking.""" + return self._started_services + async def start_plugin_services(self) -> None: """Start every contributed background service, once. diff --git a/raven/config/schema.py b/raven/config/schema.py index 49c15bfb1..50a0ad9a0 100644 --- a/raven/config/schema.py +++ b/raven/config/schema.py @@ -2569,8 +2569,9 @@ class SubagentRuntimeConfig(Base): either way before its process is closed. ``None`` or ``0`` keeps every agent until the host exits. Only an agent that can reopen its sessions (``loadSession``) is closed, so a later turn resumes the conversation in a - fresh process; an agent holding a wake is still answering a prompt and is - never idle.""" + fresh process. An agent that holds work nobody is asking for yet -- an armed + cron job of its own, a running plugin service -- says so when asked + (``_raven/residentWork``) and is kept.""" llm_relay: bool = True """Serve a Raven agent's LiteLLM calls from this process instead of loading diff --git a/raven/proactive_engine/schedulers/cron/service.py b/raven/proactive_engine/schedulers/cron/service.py index 1039e63a8..f7df66ee7 100644 --- a/raven/proactive_engine/schedulers/cron/service.py +++ b/raven/proactive_engine/schedulers/cron/service.py @@ -703,6 +703,16 @@ def list_jobs(self, include_disabled: bool = False) -> list[CronJob]: jobs = store.jobs if include_disabled else [j for j in store.jobs if j.enabled] return sorted(jobs, key=lambda j: j.state.next_run_at_ms or float("inf")) + def armed_jobs(self) -> list[CronJob]: + """The enabled jobs this runner will fire, while it runs; none once it is stopped. + + What keeps the process holding this service from being idle: closing it + leaves these to no runner, since another partition never claims them. + """ + if not self._running: + return [] + return [job for job in self.list_jobs() if self._owns_channel(job.payload.channel)] + def list_missed_foreign_oneshots(self, grace_ms: int = _MISSED_GRACE_MS) -> list[CronJob]: """Read-only: one-shot 'at' jobs owned by another partition whose fire time is more than ``grace_ms`` in the past — their owning diff --git a/tests/acp_stub_server.py b/tests/acp_stub_server.py index 92c66d634..4fd58041e 100644 --- a/tests/acp_stub_server.py +++ b/tests/acp_stub_server.py @@ -105,6 +105,10 @@ table to the prompt response, the field ACP reserves for it. Proves the host keeps it on the run record verbatim. - ``silent`` - reads and never answers, for the timeout path. +- ``resident`` - declares ``raven.residentWork``; a prompt arms an in-process + wake (``ACP_STUB_WAKE_S`` later it writes + ``ACP_STUB_WAKE_MARK``) and ends, and ``_raven/residentWork`` + answers resident until the wake has fired. - ``wedged`` - answers the first ``initialize`` and swallows everything after it, pipes open, process alive: the shape where only a probe round trip can tell nobody is home. @@ -122,6 +126,7 @@ import json import os import sys +import threading import time MODE = os.environ.get("ACP_STUB_MODE", "ok") @@ -205,6 +210,24 @@ # not call what was not declared. if os.environ.get("ACP_STUB_NO_STEER"): del CAPABILITIES["agentCapabilities"]["_meta"] +if MODE == "resident": + CAPABILITIES["agentCapabilities"]["_meta"]["raven.residentWork"] = True + +_WAKES: list = [] + + +def arm_wake() -> None: + """Fire once, ``ACP_STUB_WAKE_S`` from now, in this process: the armed cron job of a real agent.""" + + def fire() -> None: + with open(os.environ["ACP_STUB_WAKE_MARK"], "w", encoding="utf-8") as mark: + mark.write("fired") + _WAKES.remove(timer) + + timer = threading.Timer(float(os.environ.get("ACP_STUB_WAKE_S", "1.0")), fire) + timer.daemon = True + _WAKES.append(timer) + timer.start() def send(frame: dict) -> None: @@ -346,6 +369,10 @@ def handle_prompt(request_id, params) -> None: # The first N prompts end with nothing said, whatever the mode. ok(request_id, {"stopReason": "end_turn"}) return + if MODE == "resident": + arm_wake() + ok(request_id, {"stopReason": "end_turn"}) + return if MODE == "dies_mid_turn": # Half an answer, then the process is gone: the prompt went out and # fails on the client side with the connection, not with a reply. @@ -966,6 +993,8 @@ def main() -> None: handle_prompt(request_id, params) elif method == "_raven/session/steer": handle_steer(request_id, params) + elif method == "_raven/residentWork" and MODE == "resident": + ok(request_id, {"resident": bool(_WAKES)}) elif request_id is not None: err(request_id, -32601, f"unknown method {method}") diff --git a/tests/test_acp_methods.py b/tests/test_acp_methods.py index fbb29f0bc..db6c1da9a 100644 --- a/tests/test_acp_methods.py +++ b/tests/test_acp_methods.py @@ -2303,6 +2303,47 @@ async def test_a_turn_that_did_not_end_normally_is_not_held(self, rig): assert response["result"]["stopReason"] == "refusal", "only a turn that ended normally can be a wait" +class TestResidentWork: + """``_raven/residentWork``: what the agent holds that runs with nobody asking. + + A raven client asks it right before closing the agent for idleness: an armed + cron job of this agent's partition fires only from this process, and a + plugin service lives only as long as the process does. + """ + + async def test_it_is_advertised(self, rig): + response = await rig.handshake() + validate_def("InitializeResponse", response["result"]) + assert response["result"]["agentCapabilities"]["_meta"]["raven.residentWork"] is True + + async def test_an_armed_job_of_its_partition_or_a_running_service_is_resident(self, rig, tmp_path): + from raven.proactive_engine.schedulers.cron.service import CronService + from raven.proactive_engine.schedulers.cron.types import CronSchedule + + async def answer() -> dict: + response = await rig.call("_raven/residentWork", {}) + assert "result" in response, response + return response["result"] + + await rig.handshake() + assert await answer() == {"resident": False, "cronJobs": 0, "services": 0} + + cron = CronService(tmp_path / "jobs.json", allowed_channels={"acp"}) + rig.engine.cron_service = cron + cron.add_job("elsewhere", CronSchedule(kind="every", every_ms=60_000), "for feishu", channel="feishu") + await cron.start() + try: + assert (await answer())["resident"] is False, "another partition's job is not this process's to fire" + cron.add_job("mine", CronSchedule(kind="every", every_ms=60_000), "remind me", channel="acp") + assert await answer() == {"resident": True, "cronJobs": 1, "services": 0} + finally: + cron.stop() + assert (await answer())["resident"] is False, "a stopped scheduler fires nothing" + + rig.engine.running_plugin_services = (object(),) + assert await answer() == {"resident": True, "cronJobs": 0, "services": 1} + + class TestSteer: """``_raven/session/steer``: a person's words for a turn already in flight. A held prompt counts -- the agent is asleep between its own turns, and the diff --git a/tests/test_acp_pool_idle_close.py b/tests/test_acp_pool_idle_close.py index 63cf19710..6b8d1ed62 100644 --- a/tests/test_acp_pool_idle_close.py +++ b/tests/test_acp_pool_idle_close.py @@ -18,6 +18,7 @@ _STUB = Path(__file__).with_name("acp_stub_server.py") _LOADS = {"agentCapabilities": {"loadSession": True}} +_RESIDENT = {"agentCapabilities": {"loadSession": True, "_meta": {"raven.residentWork": True}}} _NO_LOAD = {"agentCapabilities": {}} @@ -28,12 +29,20 @@ async def _no_pooled_connections(): class _Client: - def __init__(self, *, idle: float = 0.0, busy: bool = False, alive: bool = True) -> None: + def __init__(self, *, idle: float = 0.0, busy: bool = False, alive: bool = True, resident: Any = None) -> None: self.name = "agent" self._idle_since = time.monotonic() - idle self.busy = busy self.alive = alive self.closed = False + self.resident = resident + self.asked: list[str] = [] + + async def request(self, method: str, params: dict | None = None, *, timeout: float | None = None) -> Any: + self.asked.append(method) + if isinstance(self.resident, BaseException): + raise self.resident + return {"resident": self.resident} @property def idle_seconds(self) -> float: @@ -68,6 +77,61 @@ async def test_a_quiet_agent_that_can_reload_its_sessions_is_closed() -> None: assert pool.live("quiet") == [] +async def test_an_agent_that_says_it_holds_resident_work_is_kept() -> None: + """An armed wake or a running service is work nobody is asking for yet: closing the agent loses it. + + Asked only of an agent that declared it can answer, and only once closing + it is otherwise settled. No answer in time, or none declared, closes it as before. + """ + holding = _Client(idle=3600, resident=True) + idle = _Client(idle=3600, resident=False) + deaf = _Client(idle=3600, resident=TimeoutError()) + undeclared = _Client(idle=3600, resident=True) + busy = _Client(idle=3600, busy=True, resident=True) + pool = _pool_holding( + holding=_held(holding, _RESIDENT), + idle=_held(idle, _RESIDENT), + deaf=_held(deaf, _RESIDENT), + undeclared=_held(undeclared, _LOADS), + busy=_held(busy, _RESIDENT), + ) + + assert sorted(await pool.close_idle(60)) == ["deaf", "idle", "undeclared"] + assert not holding.closed and not busy.closed + assert holding.asked == ["_raven/residentWork"] + assert undeclared.asked == [] and busy.asked == [] + + +async def test_an_agent_holding_a_wake_is_kept_until_the_wake_has_fired(tmp_path: Path) -> None: + """A turn arms an in-process wake and ends; the agent is then silent until the wake fires. + + The real process, through the pool: closed early, the wake would never fire. + """ + mark = tmp_path / "fired" + pool = AcpConnectionPool() + command = f"{command_quote(sys.executable)} {command_quote(str(_STUB))}" + env = {"ACP_STUB_MODE": "resident", "ACP_STUB_WAKE_S": "1.0", "ACP_STUB_WAKE_MARK": str(mark)} + held = await pool.acquire(name="stub", command=command, cwd=str(tmp_path), env=env) + opened = await held.client.request("session/new", {"cwd": str(tmp_path), "mcpServers": []}) + held.router.reserve(opened["sessionId"]) + prompt = [{"type": "text", "text": "remind me in a second"}] + await held.client.request("session/prompt", {"sessionId": opened["sessionId"], "prompt": prompt}) + held.leased_at -= 10 + + await asyncio.sleep(0.3) + assert await pool.close_idle(0.2) == [] + assert held.alive + + for _ in range(100): + if mark.exists(): + break + await asyncio.sleep(0.05) + assert mark.exists(), "the wake fired in the agent it was armed in" + await asyncio.sleep(0.3) + assert await pool.close_idle(0.2) == ["stub"] + assert not held.alive + + async def test_a_connection_whose_agent_already_ended_is_dropped() -> None: gone = _Client(idle=0, alive=False) pool = _pool_holding(gone=_held(gone, _NO_LOAD, leased_ago=0, seen=("s1",))) From 1b6cbc441f7ddd23145923720b68cc1761502fd2 Mon Sep 17 00:00:00 2001 From: Dizhan Xue Date: Thu, 29 Oct 2026 12:17:44 +0000 Subject: [PATCH 14/19] fix(providers): learn a prompt-cache refusal on a non-streamed call too LiteLLMProvider.chat returned a cache_control refusal as an error and left the downgrade to the retry layer, which suppresses the model in its own process. Behind the LLM relay that layer runs in the agent, while the resend it strips is built by the host's provider, which still took the model to accept marks and put them back: the one downgrade was refused again and the call failed. chat now suppresses the model on a refusal, as chat_stream already did, so the host learns from the first refusal and the resend goes out unmarked, the same two requests a direct agent makes. A test runs the agent and the host as separate processes, chat first, which in-process tests cannot show: there the two sides share one suppression set. --- raven/providers/litellm_provider.py | 5 ++ tests/test_provider_prompt_cache.py | 39 ++++++++++++++++ tests/test_provider_relay.py | 71 +++++++++++++++++++++++++++++ 3 files changed, 115 insertions(+) diff --git a/raven/providers/litellm_provider.py b/raven/providers/litellm_provider.py index 556bf024f..f170ba026 100644 --- a/raven/providers/litellm_provider.py +++ b/raven/providers/litellm_provider.py @@ -469,6 +469,11 @@ async def chat( # live exception here (status_code + type) before it's lost to a # string — the retry/fallback layer reads this verdict. classification = self.classify_error(e) + # Learned here as well as in the retry layer, as ``chat_stream`` does: + # the retry layer may sit in another process (an agent behind the LLM + # relay), and its stripped resend is built by this provider. + if classification.refuses_prompt_cache: + prompt_cache.suppress(original_model) head = self._provider_name or (self._gateway.name if self._gateway else None) return LLMResponse( content=format_llm_error(e, classification, provider=head), diff --git a/tests/test_provider_prompt_cache.py b/tests/test_provider_prompt_cache.py index ffdfb50e2..3d932cf19 100644 --- a/tests/test_provider_prompt_cache.py +++ b/tests/test_provider_prompt_cache.py @@ -592,6 +592,45 @@ def __init__(self): ) +async def test_a_bare_chat_learns_a_refusal_as_a_stream_does(monkeypatch): + """`chat` with no retry layer in front of it learns the downgrade, as `chat_stream` does. + + The LLM relay's host calls it that way, on behalf of an agent whose retry + layer suppresses in the agent's own process: without learning here the host + marks the agent's stripped resend again and the one downgrade is refused. + An unrelated bad request is not taken for a refusal. + """ + import litellm + + from raven.providers import litellm_provider + from raven.providers.litellm_provider import LiteLLMProvider + + model = "openrouter/anthropic/claude-3-haiku" + marked: list[bool] = [] + + async def refuse_marks(**kwargs): + marked.append("cache_control" in str(kwargs["messages"])) + if marked[-1]: + raise litellm.BadRequestError(message="400 cache_control is not permitted", model=model, llm_provider="x") + raise litellm.BadRequestError(message="400 context_length_exceeded", model=model, llm_provider="x") + + monkeypatch.setattr(litellm_provider, "acompletion", refuse_marks) + client = LiteLLMProvider(api_key="k", default_model=model, provider_name="openrouter") + messages = [{"role": "system", "content": "You are a careful assistant. " * 200}, {"role": "user", "content": "q"}] + + first = await client.chat(messages=messages) + assert first.finish_reason == "error" and first.error_classification.refuses_prompt_cache + assert prompt_cache.is_suppressed(model) + await client.chat(messages=messages) + assert marked == [True, False] + + other = "openrouter/anthropic/claude-3-opus" + await LiteLLMProvider(api_key="k", default_model=other, provider_name="openrouter").chat( + messages=[{"role": "user", "content": "q"}] + ) + assert not prompt_cache.is_suppressed(other) + + # --- The probe the construction sites hand to the strategies --- diff --git a/tests/test_provider_relay.py b/tests/test_provider_relay.py index db412f331..95ff765c8 100644 --- a/tests/test_provider_relay.py +++ b/tests/test_provider_relay.py @@ -499,6 +499,77 @@ def test_the_drivers_directory_is_found_without_importing_litellm() -> None: assert out.stdout.split() == ["True", "False"] +def test_an_agent_recovers_from_a_cache_refusal_through_a_host_in_another_process() -> None: + """A downgrade learned in the agent must not be undone by the host that builds its resend. + + Two processes, as in production: each has its own set of suppressed models, + which an in-process relay would share. The agent's first call is a plain + chat, the order in which only the host can make the resend carry marks again. + """ + host = textwrap.dedent( + """ + import asyncio, json, sys + import litellm + from raven.providers import litellm_provider + from raven.providers.relay_server import RelayServer + + wire = [] + + async def refuse_marks(**kwargs): + wire.append("marked" if "cache_control" in str(kwargs["messages"]) else "unmarked") + if wire[-1] == "marked": + raise litellm.BadRequestError(message="400 cache_control is not permitted", model="m", llm_provider="x") + return litellm.ModelResponse( + choices=[{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}], + usage={"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, + ) + + litellm_provider.acompletion = refuse_marks + + async def main(): + server = RelayServer() + await server.start() + print(json.dumps(server.env()), flush=True) + await asyncio.get_running_loop().run_in_executor(None, sys.stdin.read) + await server.close() + print(json.dumps(wire), flush=True) + + asyncio.run(main()) + """ + ) + agent = textwrap.dedent( + """ + import asyncio, json + from raven.providers import relay + from raven.providers.relay_provider import RelayProvider + + async def main(): + model = "openrouter/anthropic/claude-3-haiku" + provider = RelayProvider(default_model=model, provider_name="openrouter", api_key="k") + messages = [{"role": "system", "content": "You are a careful assistant. " * 200}, {"role": "user", "content": "q"}] + answer = await provider.chat_with_retry(messages=messages, model=model) + print(json.dumps({"relayed": relay.active(), "finish": answer.finish_reason, "content": answer.content})) + + asyncio.run(main()) + """ + ) + server = subprocess.Popen([sys.executable, "-c", host], stdin=subprocess.PIPE, stdout=subprocess.PIPE, text=True) + try: + assert server.stdout is not None and server.stdin is not None + env = json.loads(server.stdout.readline()) + out = subprocess.run( + [sys.executable, "-c", agent], capture_output=True, text=True, timeout=180, env={**os.environ, **env} + ) + server.stdin.close() + wire = json.loads(server.stdout.readline()) + finally: + server.kill() + server.wait(timeout=30) + assert out.returncode == 0, out.stderr[-2000:] + assert json.loads(out.stdout.strip().splitlines()[-1]) == {"relayed": True, "finish": "stop", "content": "ok"} + assert wire == ["marked", "unmarked"] + + def test_a_relayed_process_never_loads_litellm(tmp_path) -> None: """The point of the relay, end to end: a turn's call and its catalogue lookups, no LiteLLM import.""" host = textwrap.dedent( From a42d1686b1b10ff0bb468d2d7124035daf407cf8 Mon Sep 17 00:00:00 2001 From: Dizhan Xue Date: Thu, 29 Oct 2026 12:17:50 +0000 Subject: [PATCH 15/19] fix(acp_client): keep one agent lock per event loop The pool kept one asyncio.Lock per agent name for every loop. A lock binds to the loop of its first waiter, so a worker loop's idle check (a CLI command the RPC server runs on a thread) that waited on a lock the main loop held bound it to the worker loop, and every later wait on the main loop raised that the lock was bound to a different event loop. Locks are now kept per running loop, as dag_store keeps its own, and a finished loop's table goes with it. --- raven/acp_client/pool.py | 18 +++++++++++++---- tests/test_acp_pool_idle_close.py | 32 +++++++++++++++++++++++++++++++ 2 files changed, 46 insertions(+), 4 deletions(-) diff --git a/raven/acp_client/pool.py b/raven/acp_client/pool.py index b0fbecb3a..b27a8668f 100644 --- a/raven/acp_client/pool.py +++ b/raven/acp_client/pool.py @@ -26,6 +26,7 @@ import hashlib import json import time +import weakref from collections.abc import Awaitable, Callable from typing import Any @@ -402,7 +403,9 @@ class AcpConnectionPool: def __init__(self) -> None: self._connections: dict[str, dict[str, _Connection]] = {} - self._locks: dict[str, asyncio.Lock] = {} + self._locks: weakref.WeakKeyDictionary[asyncio.AbstractEventLoop, dict[str, asyncio.Lock]] = ( + weakref.WeakKeyDictionary() + ) self._idle_tasks: dict[asyncio.AbstractEventLoop, asyncio.Task] = {} def connections(self, name: str) -> list[_Connection]: @@ -410,10 +413,17 @@ def connections(self, name: str) -> list[_Connection]: return list(self._connections.get(name, {}).values()) def _lock_for(self, name: str) -> asyncio.Lock: - lock = self._locks.get(name) + """The running loop's lock for ``name``. + + One table per loop, as ``dag_store`` keeps its own: ``asyncio.Lock`` binds + to the loop of its first waiter, so a lock shared with a loop on another + thread (a CLI command the RPC server runs) could be bound there, and every + later wait on this loop would raise. A finished loop's table goes with it. + """ + per_loop = self._locks.setdefault(asyncio.get_running_loop(), {}) + lock = per_loop.get(name) if lock is None: - lock = asyncio.Lock() - self._locks[name] = lock + lock = per_loop[name] = asyncio.Lock() return lock def live(self, name: str) -> "list[_Connection]": diff --git a/tests/test_acp_pool_idle_close.py b/tests/test_acp_pool_idle_close.py index 6b8d1ed62..ebaf405a7 100644 --- a/tests/test_acp_pool_idle_close.py +++ b/tests/test_acp_pool_idle_close.py @@ -3,6 +3,7 @@ from __future__ import annotations import asyncio +import contextlib import json import sys import time @@ -318,6 +319,37 @@ async def beside() -> list[str]: assert here.closed +async def test_an_idle_check_on_another_loop_leaves_an_agents_lock_usable_here() -> None: + """A worker loop's check must not bind an agent's lock while this loop holds it. + + The main loop and a worker loop (a CLI command the RPC server runs) each + hold a binding of one agent. ``asyncio.Lock`` binds to the loop of its first + waiter, so a worker check that waited on the lock this loop holds would leave + every later wait here raising "bound to a different event loop". + """ + here = _Client(idle=0) + there = _Client(idle=3600) + pool = AcpConnectionPool() + beside_connection = _held(there, _LOADS) + pool._connections["agent"] = {"here": _held(here, _LOADS, leased_ago=0), "there": beside_connection} + lock = pool._lock_for("agent") + await lock.acquire() + + async def beside() -> None: + beside_connection.loop = asyncio.get_running_loop() + with contextlib.suppress(TimeoutError): + await asyncio.wait_for(pool.close_idle(60), timeout=0.5) + + await asyncio.to_thread(asyncio.run, beside()) + lock.release() + + async def hold_briefly() -> None: + async with pool._lock_for("agent"): + await asyncio.sleep(0.01) + + await asyncio.gather(hold_briefly(), hold_briefly()) + + async def test_a_failed_idle_check_does_not_end_the_next_one(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setattr(pool_mod, "_IDLE_CHECK_MIN_S", 0.01) monkeypatch.setattr(AcpConnectionPool, "_idle_budget", staticmethod(lambda: 0.04)) From 736cb6518bb7877155128fcd06066782697bf7f7 Mon Sep 17 00:00:00 2001 From: Dizhan Xue Date: Thu, 29 Oct 2026 12:18:02 +0000 Subject: [PATCH 16/19] fix(acp_client): launch an agent with its own llm environment without the relay A relayed call is made in the host, so LiteLLM reads the host's environment for anything the provider does not pass explicitly: other endpoints and API versions, cloud credentials, proxies and CA bundles. An agent whose configured env set one of these, OPENAI_BASE_URL for instance, had it ignored once relayed. Such an agent is now launched without the relay and makes its calls itself. Only the agent's own map is compared with the host: the login shell differs from a host's in ways nobody chose for one agent, and reading those would turn the relay off for every agent. API keys are left out, since every call passes its key. CONTEXT.md and the agent integrations page say so, and CONTEXT.md's LLM relay entry also notes the host learning a prompt-cache refusal. --- CONTEXT.md | 7 +++- docs-site/docs/agent-integrations.md | 2 +- docs-site/docs/agent-integrations.zh.md | 2 +- raven/acp_client/client.py | 54 +++++++++++++++++++++++-- tests/test_acp_fork_template.py | 38 +++++++++++++++++ 5 files changed, 96 insertions(+), 7 deletions(-) diff --git a/CONTEXT.md b/CONTEXT.md index cd1b97dce..dc82c9c2d 100644 --- a/CONTEXT.md +++ b/CONTEXT.md @@ -3425,8 +3425,11 @@ and a loop running beside it in another thread is offered the same socket. A con the token first, alone and under a deadline, so a peer without it is dropped before a frame is read. The host builds an agent's provider without writing the agent's key into its own environment, and makes each call in a context of its own under the window of the agent's turn, -so the output ceiling on the wire is the one the agent's loop reserved. An agent that cannot -reach the host loads LiteLLM and works as before. +so the output ceiling on the wire is the one the agent's loop reserved. A call is made with the +host's environment, so an agent whose own configuration sets something LiteLLM reads from it (an +endpoint, a cloud credential, a proxy or CA bundle) is launched without the relay; and a refusal +of prompt-cache marks is learned on the host too, which builds the agent's resend. An agent that +cannot reach the host loads LiteLLM and works as before. **Fork template** (`raven/acp_client/fork_template.py`, `raven/cli/_fork_template.py`, `subagents.runtime.forkTemplate`): On Linux, one process per host that has imported what Raven agents import, frozen out of the diff --git a/docs-site/docs/agent-integrations.md b/docs-site/docs/agent-integrations.md index 85278c2de..a6647ca19 100644 --- a/docs-site/docs/agent-integrations.md +++ b/docs-site/docs/agent-integrations.md @@ -83,7 +83,7 @@ applies to the next launch. | Key | Default | What it does | | --- | --- | --- | | `idleCloseS` | `600` | Closes an agent's process once it has been idle this many seconds: no request in flight either way and nothing said. An agent that still holds an armed reminder or a running plugin service is kept. Only an agent that can reload its sessions is closed, so the next turn resumes the conversation in a fresh process. `0` or `null` keeps agents until the host exits. | -| `llmRelay` | `true` | The host makes a native agent's LiteLLM calls, so the agent never loads LiteLLM. The agent still builds each request, retries and records usage itself. | +| `llmRelay` | `true` | The host makes a native agent's LiteLLM calls, so the agent never loads LiteLLM. The agent still builds each request, retries and records usage itself. An agent whose own `env` sets an LLM endpoint, a cloud credential, a proxy or a CA bundle makes its calls itself. | | `forkTemplate` | `true` | On Linux, native agents are forked from one process that has already loaded Raven, and share those pages instead of each loading its own copy. Other systems launch agents as before. | The template loads ahead only the modules at least two native agents share -- diff --git a/docs-site/docs/agent-integrations.zh.md b/docs-site/docs/agent-integrations.zh.md index f43872e61..acfb29e7d 100644 --- a/docs-site/docs/agent-integrations.zh.md +++ b/docs-site/docs/agent-integrations.zh.md @@ -71,7 +71,7 @@ launcher。readiness probe 成功或收到委派回执，都不代表任务已 | 键 | 默认值 | 作用 | | --- | --- | --- | | `idleCloseS` | `600` | Agent 空闲达到该秒数（双方都没有在途请求，也没有任何消息）后关闭其进程。仍持有已设定的提醒或正在运行的插件服务的 Agent 会被保留。只关闭能重新加载会话的 Agent，因此下一轮会在新进程中续上原对话。`0` 或 `null` 表示一直保留到宿主退出。 | -| `llmRelay` | `true` | 由宿主代为发出原生 Agent 的 LiteLLM 调用，Agent 自身不再加载 LiteLLM。请求组装、重试和用量记录仍由 Agent 自己完成。 | +| `llmRelay` | `true` | 由宿主代为发出原生 Agent 的 LiteLLM 调用，Agent 自身不再加载 LiteLLM。请求组装、重试和用量记录仍由 Agent 自己完成。若 Agent 自己的 `env` 设置了 LLM 端点、云凭据、代理或 CA 证书，它会自己发起调用。 | | `forkTemplate` | `true` | 在 Linux 上，原生 Agent 从一个已加载 Raven 的进程 fork 出来，共享这些内存页，而不是各自加载一份。其他系统仍按原方式启动。 | 模板只预加载至少两个原生 Agent 共用的模块，某个 Agent 独有的模块由它在 fork 之后自行 diff --git a/raven/acp_client/client.py b/raven/acp_client/client.py index 4c108fdcd..402b6ee84 100644 --- a/raven/acp_client/client.py +++ b/raven/acp_client/client.py @@ -142,8 +142,53 @@ async def _end(pgid: int, proc: Any) -> None: pass +#: What LiteLLM reads from the environment at call time beside the key and base +#: a provider passes explicitly: other endpoints and API versions, cloud +#: credentials, proxies and CA bundles. A relayed call is made in the host, with +#: the host's environment, so an agent configured with its own value for one of +#: these is launched without the relay and makes its calls itself. +_LLM_ENV_SUFFIXES = ("_API_BASE", "_BASE_URL", "_API_VERSION") +_LLM_ENV_PREFIXES = ("AWS_", "AZURE_", "VERTEX", "LITELLM_") +_LLM_ENV_NAMES = frozenset( + { + "GOOGLE_APPLICATION_CREDENTIALS", + "HTTP_PROXY", + "HTTPS_PROXY", + "ALL_PROXY", + "NO_PROXY", + "SSL_CERT_FILE", + "SSL_CERT_DIR", + "REQUESTS_CA_BUNDLE", + "CURL_CA_BUNDLE", + } +) + + +def _own_llm_env(agent_env: dict[str, str]) -> list[str]: + """The variables LiteLLM would read that this agent's own config sets differently from the host. + + Only the agent's configured map: the login shell's environment differs from a + host's in ways nobody chose for one agent (a proxy list's casing), and + reading those would turn the relay off for every agent. + """ + own = [] + for key, value in agent_env.items(): + upper = key.upper() + if upper.startswith("RAVEN_"): + continue + if upper in _LLM_ENV_NAMES or upper.endswith(_LLM_ENV_SUFFIXES) or upper.startswith(_LLM_ENV_PREFIXES): + if os.environ.get(key) != value: + own.append(key) + return sorted(own) + + async def _share_with_raven_agent( - name: str, argv: list[str], env: dict[str, str], cwd: str | None, base_env: dict[str, str] + name: str, + argv: list[str], + env: dict[str, str], + cwd: str | None, + base_env: dict[str, str], + agent_env: dict[str, str] | None = None, ) -> Any: """What this host shares with a Raven agent it launches: its model calls, and its loaded code. @@ -162,7 +207,10 @@ async def _share_with_raven_agent( return env runtime = subagent_runtime(default_live()) relayed = False - if runtime.llm_relay: + own = _own_llm_env(agent_env or {}) + if runtime.llm_relay and own: + logger.info("acp agent {!r}: its own {} would not reach relayed calls; it makes them itself", name, own) + elif runtime.llm_relay: from raven.providers.relay_server import relay_env try: @@ -532,7 +580,7 @@ async def launch( # The agent's own argv, for what this logs and reports: the spawned one # may lead with the wrapper that ties the agent to this host. program_argv = argv - shared = await _share_with_raven_agent(name, argv, child_env, cwd, base_env) + shared = await _share_with_raven_agent(name, argv, child_env, cwd, base_env, env or {}) if isinstance(shared, dict): child_env = shared else: diff --git a/tests/test_acp_fork_template.py b/tests/test_acp_fork_template.py index ec80c2785..8576a8158 100644 --- a/tests/test_acp_fork_template.py +++ b/tests/test_acp_fork_template.py @@ -245,6 +245,44 @@ async def test_an_acp_launch_of_a_raven_agent_is_forked_and_offered_the_relay(tm assert client.endpoint.returncode is not None +@pytest.mark.parametrize( + ("own", "relayed"), + [({"OPENAI_BASE_URL": "http://agent-only.invalid/v1"}, False), ({"SERPER_API_KEY": "for-a-tool"}, True)], +) +async def test_an_agent_configured_with_its_own_llm_environment_makes_its_own_calls( + tmp_path: Path, monkeypatch, own: dict[str, str], relayed: bool +) -> None: + """A relayed call is made in the host, with the host's environment. + + An agent whose own config sets something LiteLLM reads at call time (an + endpoint, a cloud credential, a proxy) would have that setting ignored, so it + is launched without the relay. A key for one of its tools changes nothing. + """ + from raven.acp_client.client import AcpClient + from raven.acp_client.pool import close_pool + + for key in own: + monkeypatch.delenv(key, raising=False) + stub = Path(__file__).with_name("acp_stub_server.py") + seen = tmp_path / "relay.txt" + script = _product( + tmp_path, + f""" + import os, runpy + open({str(seen)!r}, "w").write(os.environ.get("RAVEN_LLM_RELAY", "")) + runpy.run_path({str(stub)!r}, run_name="__main__") + """, + ) + client = await AcpClient.launch(name="own-env", command=[sys.executable, str(script)], cwd=str(tmp_path), env=own) + try: + result = await client.request("initialize", {"protocolVersion": 1, "clientCapabilities": {}}, timeout=60) + assert isinstance(result, dict) + assert bool(seen.read_text()) is relayed + finally: + await client.close() + await close_pool() + + async def test_switching_both_off_launches_a_raven_agent_as_before(tmp_path: Path, monkeypatch) -> None: from raven.acp_client.client import AcpClient from raven.config import live as live_mod From a00a0aae84418debb3f2fcd6aa8729e26a2faabf Mon Sep 17 00:00:00 2001 From: Dizhan Xue Date: Thu, 29 Oct 2026 12:37:02 +0000 Subject: [PATCH 17/19] fix(providers): carry tls_verify across the relay The per-section tls_verify option is part of a provider's construction: a section behind a self-signed gateway turns certificate checks off. The relay's spec, from which the host builds an agent's provider, lists the construction by hand and did not carry it, so a relayed agent's calls were checked on the host and refused by its gateway. The spec now carries it, and a test holds the spec to every argument the construction takes, so an argument added there later cannot be dropped the same way. --- raven/providers/relay_provider.py | 1 + tests/test_provider_relay.py | 25 +++++++++++++++++++++++++ 2 files changed, 26 insertions(+) diff --git a/raven/providers/relay_provider.py b/raven/providers/relay_provider.py index ae644b58d..bf17f5b6d 100644 --- a/raven/providers/relay_provider.py +++ b/raven/providers/relay_provider.py @@ -62,6 +62,7 @@ def spec(self) -> dict[str, Any]: "disable_auto_cache_control": bool(self.disable_auto_cache_control), "extra_body": dict(self.extra_body or {}) or None, "model_overrides": dict(self.model_overrides or {}) or None, + "tls_verify": bool(self.tls_verify), "unparsed_reasoning": self._unparsed_reasoning, "generation": relay_wire.encode(getattr(self, "generation", None) or GenerationSettings()), } diff --git a/tests/test_provider_relay.py b/tests/test_provider_relay.py index 95ff765c8..5c60e55e9 100644 --- a/tests/test_provider_relay.py +++ b/tests/test_provider_relay.py @@ -647,6 +647,31 @@ def test_the_default_builder_makes_the_direct_provider_with_the_agents_settings( assert built.generation == generation and built.default_model == "custom/m" +def test_the_spec_carries_every_construction_argument() -> None: + """The host builds an agent's provider from the spec alone, so whatever the construction takes, the spec carries. + + Listed by hand, the spec misses an argument added to the construction later, + and the host then builds a provider unlike the agent's without a word. + ``export_env`` is the host's own choice, not the agent's. + """ + import inspect + + taken = set(inspect.signature(LiteLLMShape.__init__).parameters) - {"self", "export_env"} + carried = set(relay_provider.RelayProvider(default_model="custom/m").spec()) - {"generation"} + assert carried == taken + + +def test_a_section_without_certificate_checks_keeps_them_off_on_the_host() -> None: + """``tls_verify`` off (a self-signed gateway) crosses with the construction, to the host and to the fallback.""" + agent = relay_provider.RelayProvider( + default_model="custom/m", provider_name="custom", api_base="https://gateway.internal/v1", tls_verify=False + ) + assert relay_server.build_litellm_provider(agent.spec()).tls_verify is False + assert agent._locally().tls_verify is False + checked = relay_provider.RelayProvider(default_model="custom/m", provider_name="custom") + assert relay_server.build_litellm_provider(checked.spec()).tls_verify is True + + @pytest.mark.parametrize( ("provider_name", "model", "names"), [ From 991048b55e0a32a859b4be1a506940868963984d Mon Sep 17 00:00:00 2001 From: Dizhan Xue Date: Sat, 31 Oct 2026 17:00:56 +0000 Subject: [PATCH 18/19] test(acp_client): mark the real-wake idle-close test slow The test starts a real stub ACP process and waits out the one-second wake it arms, so most of its wall clock is idle: 1.35 s to 3.65 s depending on the machine, against the suite's 3.0 s idle ceiling, which CI's unit shards enforce strictly. A real subprocess that cannot be faked and still runs in the default selection is what the slow marker is for. --- tests/test_acp_pool_idle_close.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/test_acp_pool_idle_close.py b/tests/test_acp_pool_idle_close.py index ebaf405a7..3afb1d661 100644 --- a/tests/test_acp_pool_idle_close.py +++ b/tests/test_acp_pool_idle_close.py @@ -103,6 +103,7 @@ async def test_an_agent_that_says_it_holds_resident_work_is_kept() -> None: assert undeclared.asked == [] and busy.asked == [] +@pytest.mark.slow(reason="Starts a real stub ACP process and waits out the wake it arms") async def test_an_agent_holding_a_wake_is_kept_until_the_wake_has_fired(tmp_path: Path) -> None: """A turn arms an in-process wake and ends; the agent is then silent until the wake fires. From 979c0312f9914cb5838035a7cf23e56c50046dd4 Mon Sep 17 00:00:00 2001 From: Dizhan Xue Date: Sat, 31 Oct 2026 17:00:56 +0000 Subject: [PATCH 19/19] fix(acp_client): try the fork template again after a slow start A template that did not report ready within 60 s was marked broken, so every later launch on that loop went the ordinary way, without another word, for the life of the host: one slow start on a busy machine gave up the memory the template shares until a restart. The one warning it left ended in an empty reason, since a TimeoutError prints as nothing. A missed deadline now ends that template, and launches go the ordinary way for ten minutes, at once and without a warning each, before a new one is started. The warning names the deadline and the retry. Launches that waited behind a failed start check again under the lock rather than each wait out a start of their own. A template that reports it cannot serve, or fails to start any other way, is still not used again by that loop. --- CONTEXT.md | 4 ++- raven/acp_client/fork_template.py | 45 ++++++++++++++++++++++--- tests/test_acp_fork_template.py | 55 +++++++++++++++++++++++++++++++ 3 files changed, 99 insertions(+), 5 deletions(-) diff --git a/CONTEXT.md b/CONTEXT.md index dc82c9c2d..2be6e2555 100644 --- a/CONTEXT.md +++ b/CONTEXT.md @@ -3449,7 +3449,9 @@ the template, per product, and the next template loads the modules at least two (one product's own while only one has been seen): a module one product alone imports is shared by no other process, so it is left for that product to import after the fork. Only the host's own interpreter starting `-m raven` or a product launcher is forked; anything else launches the -ordinary way. +ordinary way. A template that has not reported ready within a minute is ended, and launches go +the ordinary way for ten minutes before a new one is started; one that reports it cannot serve, +or fails to start any other way, is not used again by that loop. **Leftover agent** (`raven/acp_client/leftovers.py`): An ACP agent process group whose launching host is no longer running. An ACP agent runs in a diff --git a/raven/acp_client/fork_template.py b/raven/acp_client/fork_template.py index c1a16f343..e601362d3 100644 --- a/raven/acp_client/fork_template.py +++ b/raven/acp_client/fork_template.py @@ -44,6 +44,7 @@ import socket import sys import tempfile +import time from collections.abc import Sequence from pathlib import Path from typing import Any @@ -61,6 +62,9 @@ _TEMPLATE_PROGRAM: tuple[str, ...] = ("-m", "raven", TEMPLATE_ARG) _READY_TIMEOUT_S = 60.0 +#: How long launches go the ordinary way after a template missed the ready +#: deadline, before a new one is started. +_RETRY_AFTER_S = 600.0 _REQUEST_TIMEOUT_S = 30.0 _GRACE_POLL_S = 0.02 _TEARDOWN_GRACE_S = 1.0 @@ -192,6 +196,10 @@ def release(self) -> None: leftovers.forget(self.pid) +class _LateReadyError(Exception): + """The template did not report ready within the deadline.""" + + class ForkTemplate: """The host's handle on one template process: starts it, asks it for agents, closes it when idle.""" @@ -210,11 +218,25 @@ def __init__(self, *, skip: Sequence[str] = ()) -> None: """The loop its process and pipes were opened under; the handle is no use to another.""" self.broken = "" """Why this template will not be used again, once it has failed.""" + self._retry_at: float | None = None @property def running(self) -> bool: return self._proc is not None and self._proc.returncode is None + @property + def resting(self) -> bool: + """Whether a template that missed the ready deadline is still waiting before the next is started.""" + return self._retry_at is not None and time.monotonic() < self._retry_at + + def _refusal(self) -> str: + """Why no agent is forked from this template now, or ``""``.""" + if self.broken: + return self.broken + if self.resting: + return "the fork template missed the ready deadline and is not tried again yet" + return "" + @property def children(self) -> int: return len(self._children) @@ -251,7 +273,10 @@ async def _start(self, env: dict[str, str]) -> None: leftovers.record(self._proc.pid, "fork-template") self._log_task = asyncio.get_running_loop().create_task(self._log_stderr(self._proc), name="fork-template-log") self._hold_task = asyncio.get_running_loop().create_task(self._hold(self._proc), name="fork-template-hold") - ready = await asyncio.wait_for(self._ready_report(self._proc), timeout=_READY_TIMEOUT_S) + try: + ready = await asyncio.wait_for(self._ready_report(self._proc), timeout=_READY_TIMEOUT_S) + except TimeoutError: + raise _LateReadyError(f"the template did not report ready within {_READY_TIMEOUT_S:g} s") from None if not ready.get("ready"): raise RuntimeError(str(ready.get("reason") or "the template exited before it was ready")) from loguru import logger @@ -290,15 +315,27 @@ async def spawn( self, argv: list[str], env: dict[str, str], cwd: str | None, base_env: dict[str, str] ) -> ForkedEndpoint: """An agent forked from the template, starting the template first if needed. Raises when it cannot.""" - if self.broken: - raise RuntimeError(self.broken) + if refusal := self._refusal(): + raise RuntimeError(refusal) async with self._lock: + # Asked again: a launch that waited here behind a start that failed + # must not wait out a start of its own. + if refusal := self._refusal(): + raise RuntimeError(refusal) if not self.running: if self._proc is not None: # A template that ended on its own: clear what it left before the next. await self.close() try: await self._start(base_env) + except _LateReadyError as exc: + # A busy machine more often than a broken template: launches go + # the ordinary way for a while, then a new template is tried. + self._retry_at = time.monotonic() + _RETRY_AFTER_S + await self.close() + raise RuntimeError( + f"the fork template did not start: {exc}; a new one is tried in {_RETRY_AFTER_S:g} s" + ) from exc except BaseException as exc: self.broken = f"the fork template did not start: {exc}" await self.close() @@ -507,7 +544,7 @@ async def fork_launch( if _TEMPLATE is not None: await _TEMPLATE.close() _TEMPLATE = ForkTemplate(skip=skip) - if _TEMPLATE.broken: + if _TEMPLATE.broken or _TEMPLATE.resting: return None try: return await _TEMPLATE.spawn(argv, env, cwd, base_env) diff --git a/tests/test_acp_fork_template.py b/tests/test_acp_fork_template.py index 8576a8158..b4636e0e4 100644 --- a/tests/test_acp_fork_template.py +++ b/tests/test_acp_fork_template.py @@ -405,6 +405,61 @@ async def silent(proc): await template.spawn([sys.executable, str(script)], base, str(tmp_path), base) +async def test_a_template_late_to_report_ready_is_tried_again_after_a_rest(tmp_path: Path, monkeypatch) -> None: + """A missed ready deadline is a busy machine, not a broken template. + + The launch that waited goes the ordinary way and says why; launches during + the rest go the ordinary way at once and quietly, neither starting a + template nor waiting for one; after the rest a new template starts and forks. + """ + from loguru import logger + + warnings: list[str] = [] + sink = logger.add(lambda message: warnings.append(message.record["message"]), level="WARNING") + script = _product(tmp_path, "pass\n") + base = dict(os.environ) + argv = [sys.executable, str(script)] + try: + monkeypatch.setattr(fork_template, "_READY_TIMEOUT_S", 0.01) + assert await fork_launch(argv, base, str(tmp_path), base_env=base) is None + assert len(warnings) == 1 and "did not report ready within 0.01 s" in warnings[0], warnings + template = fork_template._TEMPLATE + assert template is not None and not template.broken and not template.running + + monkeypatch.setattr(fork_template, "_READY_TIMEOUT_S", 60.0) + started = time.monotonic() + assert await fork_launch(argv, base, str(tmp_path), base_env=base) is None + assert time.monotonic() - started < 5 and not template.running + assert len(warnings) == 1, warnings + finally: + logger.remove(sink) + + monkeypatch.setattr(template, "_retry_at", time.monotonic()) + endpoint = await fork_launch(argv, base, str(tmp_path), base_env=base) + assert endpoint is not None and "forked" in endpoint.describe() + await asyncio.wait_for(endpoint.wait(), timeout=30) + endpoint.release() + + +async def test_launches_waiting_on_a_late_template_do_not_each_start_another(tmp_path: Path, monkeypatch) -> None: + starts = 0 + real_start = fork_template.ForkTemplate._start + + async def counted(self, env): + nonlocal starts + starts += 1 + await real_start(self, env) + + monkeypatch.setattr(fork_template.ForkTemplate, "_start", counted) + monkeypatch.setattr(fork_template, "_READY_TIMEOUT_S", 0.01) + script = _product(tmp_path, "pass\n") + base = dict(os.environ) + argv = [sys.executable, str(script)] + launches = [fork_launch(argv, base, str(tmp_path), base_env=base) for _ in range(3)] + assert await asyncio.gather(*launches) == [None, None, None] + assert starts == 1 + + async def test_the_ready_line_is_found_among_what_an_import_printed() -> None: from types import SimpleNamespace