fix(agent): scope skill index to owner (#2404)

Co-authored-by: Kenny Van de Maele <kenny@kvandemaele.be>
This commit is contained in:
nubs
2026-06-09 07:51:29 +00:00
committed by GitHub
parent 0aba00f4cf
commit f1cda91683
2 changed files with 59 additions and 3 deletions
+5 -3
View File
@@ -855,7 +855,7 @@ def _build_system_prompt(
_ov_sig = _hl.sha256(_json.dumps(get_builtin_overrides() or {}, sort_keys=True).encode()).hexdigest() _ov_sig = _hl.sha256(_json.dumps(get_builtin_overrides() or {}, sort_keys=True).encode()).hexdigest()
except Exception: except Exception:
_ov_sig = "" _ov_sig = ""
cache_key = (frozenset(disabled_tools or []), bool(mcp_mgr), needs_admin, _rt_key, compact, _ov_sig, suppress_local_context) cache_key = (frozenset(disabled_tools or []), bool(mcp_mgr), needs_admin, _rt_key, compact, _ov_sig, owner, suppress_local_context)
if _cached_base_prompt and _cached_base_prompt_key == cache_key and not active_document: if _cached_base_prompt and _cached_base_prompt_key == cache_key and not active_document:
agent_prompt = _cached_base_prompt agent_prompt = _cached_base_prompt
# Skill index is user-editable (name + description), so it must never # Skill index is user-editable (name + description), so it must never
@@ -863,7 +863,7 @@ def _build_system_prompt(
# when the cache hits. # when the cache hits.
_, _skill_index_block = _build_base_prompt( _, _skill_index_block = _build_base_prompt(
disabled_tools, mcp_mgr, needs_admin, relevant_tools, disabled_tools, mcp_mgr, needs_admin, relevant_tools,
mcp_disabled_map=mcp_disabled_map, compact=compact, mcp_disabled_map=mcp_disabled_map, compact=compact, owner=owner,
suppress_local_context=suppress_local_context, suppress_local_context=suppress_local_context,
) )
else: else:
@@ -874,6 +874,7 @@ def _build_system_prompt(
relevant_tools, relevant_tools,
mcp_disabled_map=mcp_disabled_map, mcp_disabled_map=mcp_disabled_map,
compact=compact, compact=compact,
owner=owner,
suppress_local_context=suppress_local_context, suppress_local_context=suppress_local_context,
) )
if not active_document: if not active_document:
@@ -1246,6 +1247,7 @@ def _build_base_prompt(
relevant_tools=None, relevant_tools=None,
mcp_disabled_map=None, mcp_disabled_map=None,
compact: bool = False, compact: bool = False,
owner: Optional[str] = None,
suppress_local_context: bool = False, suppress_local_context: bool = False,
): ):
"""Build the agent prompt with only relevant tools included. """Build the agent prompt with only relevant tools included.
@@ -1299,7 +1301,7 @@ def _build_base_prompt(
from src.constants import DATA_DIR from src.constants import DATA_DIR
_sm = SkillsManager(DATA_DIR) _sm = SkillsManager(DATA_DIR)
active_tools = list(set(TOOL_SECTIONS.keys()) - set(disabled or [])) active_tools = list(set(TOOL_SECTIONS.keys()) - set(disabled or []))
skill_idx = _sm.index_for(owner=None, active_toolsets=active_tools) skill_idx = _sm.index_for(owner=owner, active_toolsets=active_tools)
if skill_idx: if skill_idx:
lines = ["## Available skills", lines = ["## Available skills",
"Procedures the assistant should consult before doing domain work. " "Procedures the assistant should consult before doing domain work. "
@@ -76,6 +76,23 @@ def _seed_index_skill(tmp_path: Path) -> Path:
return data_dir return data_dir
def _write_index_skill(data_dir: Path, name: str, description: str, owner: str) -> None:
skill_dir = data_dir / "skills" / owner / name
skill_dir.mkdir(parents=True, exist_ok=True)
(skill_dir / "SKILL.md").write_text(
"---\n"
f"name: {name}\n"
f"description: {description}\n"
"when_to_use: when this owner needs a private workflow\n"
"category: private\n"
"status: published\n"
f"owner: {owner}\n"
"---\n\n"
f"# {name}\n",
encoding="utf-8",
)
def _patch_prefs(monkeypatch, data_dir): def _patch_prefs(monkeypatch, data_dir):
"""Mirror the helpers from test_skill_prompt_injection.py: point """Mirror the helpers from test_skill_prompt_injection.py: point
`src.constants.DATA_DIR` at our tmp, and patch the prefs loader so `src.constants.DATA_DIR` at our tmp, and patch the prefs loader so
@@ -152,3 +169,40 @@ def test_skill_index_lands_in_untrusted_user_message(tmp_path, monkeypatch):
) )
assert untrusted[0]["role"] == "user" assert untrusted[0]["role"] == "user"
assert "Source: skills" in untrusted[0]["content"] assert "Source: skills" in untrusted[0]["content"]
def test_skill_index_is_owner_scoped_across_prompt_cache_hits(tmp_path, monkeypatch):
"""Authenticated users must not receive another user's skill index.
This calls the prompt builder twice without clearing the base-prompt cache,
so the second call exercises the cache-hit path as well as owner scoping.
"""
data_dir = tmp_path / "data"
_write_index_skill(data_dir, "alice-only", "Alice private procedure", "alice")
_write_index_skill(data_dir, "bob-only", "Bob private procedure", "bob")
_patch_prefs(monkeypatch, data_dir)
from src.agent_loop import _build_system_prompt # noqa: WPS433
messages = [{"role": "user", "content": "use my workflow"}]
alice_out, _ = _build_system_prompt(
messages=messages, model="test-model",
active_document=None, mcp_mgr=None, owner="alice",
)
bob_out, _ = _build_system_prompt(
messages=messages, model="test-model",
active_document=None, mcp_mgr=None, owner="bob",
)
alice_text = "\n".join(m.get("content", "") or "" for m in alice_out)
bob_text = "\n".join(m.get("content", "") or "" for m in bob_out)
assert "alice-only" in alice_text
assert "Alice private procedure" in alice_text
assert "bob-only" not in alice_text
assert "Bob private procedure" not in alice_text
assert "bob-only" in bob_text
assert "Bob private procedure" in bob_text
assert "alice-only" not in bob_text
assert "Alice private procedure" not in bob_text