"""Synthetic RAG expansion boundary. Standard library; no model or network. Run: python3 rag-neighbour-authorisation.py [--test] The in-memory store and policy are trusted fixture inputs, not an auth service. """ import argparse from dataclasses import dataclass, replace from typing import FrozenSet, Optional import unittest @dataclass(frozen=True) class Principal: tenant: str groups: FrozenSet[str] @dataclass(frozen=True) class Chunk: id: str tenant: str parent: str position: int groups: Optional[FrozenSet[str]] text: str class PolicyUnavailable(RuntimeError): pass class Store: def __init__(self, chunks): self.chunks = {} for chunk in chunks: if chunk.id in self.chunks: raise ValueError("duplicate chunk ID") self.chunks[chunk.id] = chunk self.body_reads = [] def read_body(self, chunk_id): self.body_reads.append(chunk_id) return self.chunks[chunk_id].text class Policy: def __init__(self, available=True): self.available = available def allows(self, principal, chunk): if not self.available: raise PolicyUnavailable("policy unavailable") return (chunk.tenant == principal.tenant and chunk.groups is not None and bool(principal.groups.intersection(chunk.groups))) def candidates(store, seed, radius): # This fixture intentionally does not rely on parent IDs being global. return sorted((c for c in store.chunks.values() if c.parent == seed.parent and abs(c.position - seed.position) <= radius), key=lambda c: (c.position, c.id)) def assemble(store, principal, seed_ids, policy, radius=1, unsafe=False): if type(radius) is not int or radius < 0: raise ValueError("radius must be a non-negative integer") selected = {} for seed_id in seed_ids: seed = store.chunks.get(seed_id) if seed is None: raise ValueError("unknown seed") if not policy.allows(principal, seed): continue for chunk in candidates(store, seed, radius): # Unsafe branch is retained ONLY as a counterexample. if unsafe or policy.allows(principal, chunk): selected[chunk.id] = chunk # Finish all policy decisions before reading any body. A policy exception # aborts the request instead of returning a partially assembled context. return [(cid, store.read_body(cid)) for cid in selected] def corpus(): return [ Chunk("a0", "alpha", "guide", 0, frozenset({"staff"}), "Public-to-staff setup instructions."), Chunk("a1", "alpha", "guide", 1, frozenset({"finance"}), "SYNTHETIC_FINANCE_ONLY"), Chunk("a2", "alpha", "guide", 2, frozenset({"staff"}), "Public-to-staff troubleshooting."), Chunk("b0", "beta", "guide", 0, frozenset({"staff"}), "SYNTHETIC_OTHER_TENANT"), ] def demo(): p = Principal("alpha", frozenset({"staff"})) for label, unsafe in [("seed_only", True), ("per_chunk", False)]: store = Store(corpus()) rows = assemble(store, p, ["a0"], Policy(), unsafe=unsafe) print(label + "=" + ",".join(cid for cid, _ in rows)) print(label + "_body_reads=" + ",".join(store.body_reads)) class Tests(unittest.TestCase): def setUp(self): self.p = Principal("alpha", frozenset({"staff"})) self.store = Store(corpus()) self.policy = Policy() def run_safe(self, seeds=("a0",), **kwargs): return assemble(self.store, self.p, seeds, self.policy, **kwargs) def test_unsafe_path_discloses_denied_neighbours(self): rows = assemble(self.store, self.p, ["a0"], self.policy, unsafe=True) self.assertEqual([cid for cid, _ in rows], ["a0", "b0", "a1"]) def test_safe_path_never_reads_denied_bodies(self): self.assertEqual(self.run_safe(), [("a0", corpus()[0].text)]) self.assertEqual(self.store.body_reads, ["a0"]) def test_tenant_check_even_when_group_names_match(self): self.run_safe() self.assertNotIn("b0", self.store.body_reads) def test_missing_policy_metadata_denied(self): self.store.chunks["a1"] = replace(self.store.chunks["a1"], groups=None) self.run_safe() self.assertNotIn("a1", self.store.body_reads) def test_empty_acl_denied(self): self.store.chunks["a1"] = replace(self.store.chunks["a1"], groups=frozenset()) self.run_safe() self.assertEqual(self.store.body_reads, ["a0"]) def test_denied_seed_cannot_start_expansion(self): self.assertEqual(self.run_safe(("a1",)), []) self.assertEqual(self.store.body_reads, []) def test_authorised_neighbour_is_retained(self): self.store.chunks["a1"] = replace(self.store.chunks["a1"], groups=frozenset({"staff"})) self.assertEqual([cid for cid, _ in self.run_safe()], ["a0", "a1"]) def test_policy_outage_aborts_without_body_reads(self): self.policy.available = False with self.assertRaises(PolicyUnavailable): self.run_safe() self.assertEqual(self.store.body_reads, []) def test_outage_after_seed_check_aborts_before_body_reads(self): class FailsLater(Policy): calls = 0 def allows(inner, principal, chunk): inner.calls += 1 if inner.calls == 3: raise PolicyUnavailable("synthetic mid-request outage") return super().allows(principal, chunk) self.policy = FailsLater() with self.assertRaises(PolicyUnavailable): self.run_safe() self.assertEqual(self.store.body_reads, []) def test_repeated_seeds_do_not_duplicate_context(self): self.run_safe(("a0", "a0")) self.assertEqual(self.store.body_reads, ["a0"]) def test_missing_seed_aborts_before_body_reads(self): with self.assertRaises(ValueError): self.run_safe(("a0", "missing")) self.assertEqual(self.store.body_reads, []) def test_invalid_radius_rejected(self): for radius in [-1, True, 1.5]: with self.subTest(radius=radius), self.assertRaises(ValueError): self.run_safe(radius=radius) def test_duplicate_ids_rejected(self): with self.assertRaises(ValueError): Store([corpus()[0], corpus()[0]]) def test_empty_seed_set(self): self.assertEqual(self.run_safe(()), []) self.assertEqual(self.store.body_reads, []) if __name__ == "__main__": parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--test", action="store_true") args = parser.parse_args() if args.test: unittest.main(argv=[__file__], verbosity=2) else: demo()