From 73b410023ea7e73bb45d691679560407b2207436 Mon Sep 17 00:00:00 2001 From: sami jaghouar Date: Thu, 11 Jun 2026 00:29:28 +0000 Subject: [PATCH] Add uploaded skill allowlist --- README.md | 3 +++ install.sh | 33 +++++++++++++++++++++++++++++++-- src/rlm/engine.py | 6 +++--- src/rlm/tools/skills.py | 28 ++++++++++++++++++++++++++-- tests/test_skills.py | 32 +++++++++++++++++++++++++++++++- 5 files changed, 94 insertions(+), 8 deletions(-) diff --git a/README.md b/README.md index 0e9df31..03a180c 100644 --- a/README.md +++ b/README.md @@ -61,6 +61,7 @@ All configuration is via environment variables: | `PRIME_API_KEY` | — | PI Inference pair: targets `https://api.pinference.ai/api/v1` and forwards `PRIME_TEAM_ID` as `X-Prime-Team-ID` when set. | | `OPENAI_API_KEY` / `OPENAI_BASE_URL` | resolved by SDK | OpenAI pair — when `OPENAI_API_KEY` is set, AsyncOpenAI's native env handling is used (covers OpenAI direct and verifiers' rollout tunnel both). Provider precedence: explicit → PI → OpenAI. Keys are scoped to their own base URL so an `OPENAI_API_KEY` lying around can't leak to PI Inference. | | `RLM_TOOLS` | `ipython` | Comma-separated subset of builtin tools (`ipython`, `bash`, `edit`) to enable. Empty string = no tools. Unknown names raise. | +| `RLM_SKILLS` | all installed skills | Comma-separated allowlist of uploaded skills to install/expose/pre-import. Unset = all skills, empty string = no skills. | | `RLM_MAX_DEPTH` | `0` | Max recursion depth (`0` means no sub-agents) | | `RLM_EXEC_TIMEOUT` | `300` | Seconds per IPython execution | | `RLM_MAX_OUTPUT` | `-1` | Max chars returned from a tool call (`-1` disables truncation; `0` is invalid) | @@ -121,6 +122,8 @@ These artifacts are consumable for debugging, visualization, or training-data ex `rlm` itself ships no skills. Skills are supplied by the host environment: before `install.sh` runs, the environment places skill packages under `/task/rlm-skills//`, and `install.sh` installs them alongside `rlm` so they're both importable and on `$PATH`. +Set `RLM_SKILLS` to restrict this surface. Unset installs and exposes every uploaded skill, `RLM_SKILLS=""` disables uploaded skills, and a comma-separated list such as `RLM_SKILLS=websearch,search_docs` installs/exposes only those skills. The same filter controls system-prompt skill text and IPython pre-imports. + From IPython, import a skill and call its async `run(...)` entrypoint: ```python diff --git a/install.sh b/install.sh index 45b3e1a..d2e9519 100644 --- a/install.sh +++ b/install.sh @@ -121,19 +121,48 @@ fi # Skills are owned by the environment (e.g. ComposableEnv uploads them to # /task/rlm-skills before this script runs). Discover and install any # that are present so they're both importable and on PATH. +skill_enabled() { + local tool_name="$1" + + # Unset means install every uploaded skill. Set to an empty string to + # install none, or to a comma-separated allowlist of import/CLI names. + if [ -z "${RLM_SKILLS+x}" ]; then + return 0 + fi + if [ -z "$RLM_SKILLS" ]; then + return 1 + fi + + local IFS="," + local allowed + for allowed in $RLM_SKILLS; do + allowed="$(echo "$allowed" | tr -d '[:space:]' | tr '-' '_')" + if [ "$allowed" = "$tool_name" ]; then + return 0 + fi + done + return 1 +} + SKILL_ARGS="" SKILL_TOOL_NAMES="" if [ -d /task/rlm-skills ]; then for skill_dir in /task/rlm-skills/*/; do [ -f "$skill_dir/pyproject.toml" ] || continue skill_name=$(grep '^name' "$skill_dir/pyproject.toml" | head -1 | sed 's/.*"\(.*\)".*/\1/') - SKILL_ARGS="$SKILL_ARGS --with-editable $skill_dir --with-executables-from $skill_name" + tool_name="" for candidate in "$skill_dir"/src/*; do [ -d "$candidate" ] || continue [ "$(basename "$candidate")" = "__pycache__" ] && continue - SKILL_TOOL_NAMES="$SKILL_TOOL_NAMES $(basename "$candidate")" + tool_name="$(basename "$candidate")" break done + [ -n "$tool_name" ] || continue + if ! skill_enabled "$tool_name"; then + continue + fi + SKILL_ARGS="$SKILL_ARGS --with-editable $skill_dir --with-executables-from $skill_name" + SKILL_TOOL_NAMES="$SKILL_TOOL_NAMES $tool_name" done fi diff --git a/src/rlm/engine.py b/src/rlm/engine.py index 5deec84..8455f07 100644 --- a/src/rlm/engine.py +++ b/src/rlm/engine.py @@ -26,7 +26,6 @@ ) from rlm.types import CompactionApplied, RLMMetrics, RLMResult, TokenUsage - # Injected as a user message when the branch's context size reaches the # compaction threshold. The model's next reply is expected to be a # plain-text handoff summary; any tool calls it emits are ignored and @@ -519,10 +518,11 @@ def _load_system_prompt( ) -> str: if self.system_prompt_path: return Path(self.system_prompt_path).read_text() + installed_skills = get_installed_skills() system_prompt = build_system_prompt( self.cwd, - str(SKILLS_DIR) if SKILLS_DIR is not None else None, - get_installed_skills(), + str(SKILLS_DIR) if SKILLS_DIR is not None and installed_skills else None, + installed_skills, messages_path, allow_recursion=self.depth < self.max_depth, active_tools=active_tools, diff --git a/src/rlm/tools/skills.py b/src/rlm/tools/skills.py index 6cac7ce..50bb48c 100644 --- a/src/rlm/tools/skills.py +++ b/src/rlm/tools/skills.py @@ -2,10 +2,10 @@ from __future__ import annotations +import os from importlib import metadata from pathlib import Path - TASK_SKILLS_DIR = Path("/task/rlm-skills") @@ -22,6 +22,18 @@ def _normalize_skill_name(name: str) -> str: return name.replace("-", "_") +def _allowed_skill_names() -> list[str] | None: + """Return the RLM_SKILLS allowlist, or None when all skills are enabled.""" + raw = os.environ.get("RLM_SKILLS") + if raw is None: + return None + return [ + _normalize_skill_name(token.strip()) + for token in raw.split(",") + if token.strip() + ] + + def get_installed_skills() -> list[str]: """Return installed skill names discovered from distribution metadata.""" skills: set[str] = set() @@ -30,4 +42,16 @@ def get_installed_skills() -> list[str]: name = dist.metadata.get("Name", "") if name.startswith(prefix): skills.add(_normalize_skill_name(name[len(prefix) :])) - return sorted(skills) + + installed = sorted(skills) + allowed = _allowed_skill_names() + if allowed is None: + return installed + + unknown = sorted(set(allowed).difference(installed)) + if unknown: + raise ValueError( + f"RLM_SKILLS contains unknown skill(s): {unknown}. " + f"Installed: {installed}" + ) + return [skill for skill in installed if skill in allowed] diff --git a/tests/test_skills.py b/tests/test_skills.py index d349e43..0eeea9c 100644 --- a/tests/test_skills.py +++ b/tests/test_skills.py @@ -21,8 +21,38 @@ show_tool_result, tool_result, ) - from rlm.engine import RLMEngine +from rlm.tools.skills import get_installed_skills + + +def test_installed_skills_default_all(monkeypatch): + monkeypatch.delenv("RLM_SKILLS", raising=False) + + assert get_installed_skills() == ["boom", "say"] + + +def test_installed_skills_allowlist(monkeypatch): + monkeypatch.setenv("RLM_SKILLS", "say") + + assert get_installed_skills() == ["say"] + + +def test_installed_skills_empty_allowlist(monkeypatch): + monkeypatch.setenv("RLM_SKILLS", "") + + assert get_installed_skills() == [] + + +def test_installed_skills_unknown_allowlist(monkeypatch): + monkeypatch.setenv("RLM_SKILLS", "missing") + + try: + get_installed_skills() + except ValueError as exc: + assert "RLM_SKILLS contains unknown skill" in str(exc) + assert "missing" in str(exc) + else: + raise AssertionError("expected unknown RLM_SKILLS entry to raise") async def test_python_skill_valid(session):