diff --git a/CONTEXT.md b/CONTEXT.md index 24bba0e85..2be6e2555 100644 --- a/CONTEXT.md +++ b/CONTEXT.md @@ -3392,6 +3392,67 @@ 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`. 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 +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. 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. 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 +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, 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. 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 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..a6647ca19 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. 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. 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 -- +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..acfb29e7d 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 会被保留。只关闭能重新加载会话的 Agent，因此下一轮会在新进程中续上原对话。`0` 或 `null` 表示一直保留到宿主退出。 | +| `llmRelay` | `true` | 由宿主代为发出原生 Agent 的 LiteLLM 调用，Agent 自身不再加载 LiteLLM。请求组装、重试和用量记录仍由 Agent 自己完成。若 Agent 自己的 `env` 设置了 LLM 端点、云凭据、代理或 CA 证书，它会自己发起调用。 | +| `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/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/client.py b/raven/acp_client/client.py index 0b5db96f2..402b6ee84 100644 --- a/raven/acp_client/client.py +++ b/raven/acp_client/client.py @@ -142,6 +142,92 @@ 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], + 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. + + 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 + 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: + 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 +580,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, env or {}) + 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 +690,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 +949,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..e601362d3 --- /dev/null +++ b/raven/acp_client/fork_template.py @@ -0,0 +1,563 @@ +"""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. + +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. 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 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. +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 +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 asyncio +import contextlib +import json +import os +import signal +import socket +import sys +import tempfile +import time +from collections.abc import Sequence +from pathlib import Path +from typing import Any + +#: 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" + +#: 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 +#: 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 +_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 _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.""" + + 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._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.""" + 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) + + 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", *_TEMPLATE_PROGRAM, "--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") + self._hold_task = asyncio.get_running_loop().create_task(self._hold(self._proc), name="fork-template-hold") + 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 + + 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 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() + 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") + + 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. 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 + + 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) + transport = getattr(proc, "_transport", None) + if transport is not None: + with contextlib.suppress(Exception): + transport.close() + 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 + 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 + + +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(): + 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 + 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 or _TEMPLATE.resting: + 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() diff --git a/raven/acp_client/pool.py b/raven/acp_client/pool.py index f119947fc..b27a8668f 100644 --- a/raven/acp_client/pool.py +++ b/raven/acp_client/pool.py @@ -22,8 +22,11 @@ from __future__ import annotations import asyncio +import contextlib import hashlib import json +import time +import weakref from collections.abc import Awaitable, Callable from typing import Any @@ -57,6 +60,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 @@ -73,6 +83,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.""" @@ -297,6 +311,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.""" @@ -334,11 +360,22 @@ 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 + 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. @@ -366,17 +403,27 @@ 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]: """Every connection held under ``name``, live or not.""" 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]": @@ -444,6 +491,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 +581,7 @@ async def acquire( config=config, ) held[bkey] = connection + self._watch_idle() return connection async def adopt( @@ -582,6 +631,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 +669,108 @@ 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 + + @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. + + ``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. 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 [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 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: + 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 this loop's idle check if it is not running; it ends once the loop holds nothing.""" + loop = asyncio.get_running_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_tasks[loop] = loop.create_task(self._idle_loop(), name="acp-idle-close") + + async def _idle_loop(self) -> None: + 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)) + 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 +780,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. """ + 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 @@ -712,6 +873,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/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/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/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/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/raven/config/live.py b/raven/config/live.py index 43b324385..1c9293dc5 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,31 @@ 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 + + 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 isinstance(section, dict): + 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..50a0ad9a0 100644 --- a/raven/config/schema.py +++ b/raven/config/schema.py @@ -2555,6 +2555,38 @@ 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 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 + 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 +2604,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/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/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..f170ba026 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.""" @@ -768,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), @@ -1260,17 +966,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..4dcb8b790 --- /dev/null +++ b/raven/providers/litellm_shape.py @@ -0,0 +1,359 @@ +"""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, + 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: + 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} + + # 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: + """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..990a795b9 --- /dev/null +++ b/raven/providers/relay.py @@ -0,0 +1,294 @@ +"""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, cast + +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 {"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) + sock.settimeout(_CONNECT_TIMEOUT_S) + 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) + 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]) + if (unreachable := _unreachable_kind(kind)) is not None: + raise unreachable + 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) + 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: + 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 cast(F, call) + + 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..bf17f5b6d --- /dev/null +++ b/raven/providers/relay_provider.py @@ -0,0 +1,195 @@ +"""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.binding import active_binding +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, + "tls_verify": bool(self.tls_verify), + "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]: + # 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, + 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..1f3cd8f95 --- /dev/null +++ b/raven/providers/relay_server.py @@ -0,0 +1,344 @@ +"""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, 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, 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. + +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 contextvars +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.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 +#: 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] + + +def build_litellm_provider(spec: dict[str, Any]) -> LLMProvider: + """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, export_env=False) + 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 = 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 = "" + 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.""" + + @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._accept, 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._accept, 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 + + 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) + 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: + admitted = await self._admit(reader) + if admitted is None: + return + 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) + 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(), 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}) + + 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(), self._turn(request, provider, args): + 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: + 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 + 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..63ccf72de --- /dev/null +++ b/raven/providers/relay_wire.py @@ -0,0 +1,311 @@ +"""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 + + +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() + + +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/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_fork_template.py b/tests/test_acp_fork_template.py new file mode 100644 index 000000000..b4636e0e4 --- /dev/null +++ b/tests/test_acp_fork_template.py @@ -0,0 +1,642 @@ +"""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 + + +@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 + + 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() + + +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_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 + + 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 + + +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() + + +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" + + +# 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]: + 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 + + 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 + assert not Path(f"/proc/{second_pid}").exists() + finally: + 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( + "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. + + 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_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 new file mode 100644 index 000000000..3afb1d661 --- /dev/null +++ b/tests/test_acp_pool_idle_close.py @@ -0,0 +1,371 @@ +"""The ACP pool closes an agent's process once it has been idle for the configured budget.""" + +from __future__ import annotations + +import asyncio +import contextlib +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}} +_RESIDENT = {"agentCapabilities": {"loadSession": True, "_meta": {"raven.residentWork": 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, 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: + 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_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 == [] + + +@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. + + 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",))) + + 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_tasks.get(asyncio.get_running_loop()) + 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 + + +@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_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_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)) + 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_tasks[asyncio.get_running_loop()], timeout=5) + assert len(calls) >= 2 and client.closed 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"]] 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..3d932cf19 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", @@ -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 new file mode 100644 index 000000000..5c60e55e9 --- /dev/null +++ b/tests/test_provider_relay.py @@ -0,0 +1,1046 @@ +"""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 shutil +import subprocess +import sys +import tempfile +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.binding import ModelBinding, active_binding, use_binding +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" + + 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"]): + 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 + + +# 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": []} + + 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_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"}]) + 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 _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") + 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_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( + """ + 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" + + +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"), + [ + ("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") + 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: + reader, writer = await _open(server) + try: + if frame is None: + return None + 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() + + +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"}]) + + +@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) -> 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: + assert await reader.readline() == relay_wire.greeting("t") + await relay_wire.read_frame(reader) + reply = next(replies, None) + if reply is not None: + await relay_wire.write_frame(writer, reply) + writer.close() + + 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() + 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() + shutil.rmtree(folder, ignore_errors=True) + + +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:")) + + +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) + 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() 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}")