import json from pathlib import Path import subprocess import sys import unittest from demo import replay from isolation import (Agent, ContextCache, ContextStore, Principal, ScopeError, build_context, naive_context) HERE = Path(__file__).parent QUESTION = 'What is the plan for the export job?' def scenario(): return json.loads((HERE / 'scenario.json').read_text()) def store_before_close(): data = scenario() store = ContextStore({a['name']: Agent(a['name'], a['role'], tuple(a['tools']), tuple(a['reads'])) for a in data['agents']}) for item in data['configuration']: store.configure(item['id'], item['text']) for step in data['steps']: if step['op'] == 'write': store.write(Principal(*step['as']), step['scope'], step['id'], step['text']) elif step['op'] == 'handoff': store.handoff(Principal(*step['as']), step['to'], step['items'], step['brief']) return store BOB = Principal('bob', 'task-17', 'planner') ALICE = Principal('alice', 'task-9', 'planner') class IsolationTest(unittest.TestCase): def test_fixture_matches_recorded_result(self): output = subprocess.run([sys.executable, 'demo.py'], cwd=HERE, check=True, capture_output=True, text=True).stdout self.assertEqual(json.loads(output), json.loads((HERE / 'isolation-result.json').read_text())) def test_other_users_items_never_appear(self): store = store_before_close() for principal in (BOB, ALICE, Principal('bob', 'task-17', 'researcher')): for item in store.visible(principal): self.assertIn(item.user, (None, principal.user)) context = json.dumps(build_context(store, principal, QUESTION)['messages']) other = 'alice@example.com' if principal.user == 'bob' else 'move the export job' self.assertNotIn(other, context) def test_naive_global_list_leaks_another_users_item(self): result = replay(scenario()) queries = [step for step in result['steps'] if 'principal' in step] self.assertTrue(all(step['naive_other_user_items'] for step in queries)) self.assertIn('alice-export-limit', naive_context(store_before_close(), QUESTION)) def test_task_items_do_not_cross_tasks(self): ids = [item.id for item in store_before_close().visible(Principal('bob', 'task-18', 'planner'))] self.assertEqual(ids, ['app-policy', 'bob-style']) def test_agent_scratch_is_private(self): store = store_before_close() researcher = [item.id for item in store.visible(Principal('bob', 'task-17', 'researcher'))] planner = [item.id for item in store.visible(BOB)] self.assertNotIn('planner-scratch', researcher) self.assertNotIn('handoff-1-brief', planner) def test_subagent_receives_only_the_handoff(self): context = build_context(store_before_close(), Principal('bob', 'task-17', 'researcher'), QUESTION) self.assertEqual(context['item_ids'], ['app-policy', 'handoff-1-brief', 'handoff-1-item-1']) self.assertEqual(context['tools'], ['search_docs']) self.assertNotIn('Bob prefers short answers', json.dumps(context['messages'])) def test_handoff_rejects_items_the_sender_cannot_read(self): store = store_before_close() with self.assertRaisesRegex(ScopeError, 'cannot read'): store.handoff(BOB, 'researcher', ['alice-contact'], 'Look up the contact.') def test_failed_handoff_writes_nothing(self): class FailingStore(ContextStore): # Reject the third prepared item, after two have passed validation. calls = 0 def _prepare(self, *args): FailingStore.calls += 1 if FailingStore.calls == 3: raise ScopeError('simulated failure on the third item') return super()._prepare(*args) source = store_before_close() store = FailingStore(source.agents, list(source.items), source.version) before, version = list(store.items), store.version with self.assertRaisesRegex(ScopeError, 'third item'): store.handoff(BOB, 'researcher', ['bob-export-plan', 'bob-style'], 'Second brief.') self.assertEqual(store.items, before) self.assertEqual(store.version, version) with self.assertRaisesRegex(ScopeError, 'twice'): store.handoff(BOB, 'researcher', ['bob-export-plan', 'bob-export-plan'], 'Brief.') self.assertEqual(store.items, before) def test_second_handoff_gets_its_own_ids(self): store = store_before_close() store.handoff(BOB, 'researcher', ['bob-style'], 'Keep the answer short.') ids = [item.id for item in store.visible(Principal('bob', 'task-17', 'researcher'))] self.assertEqual(ids, ['app-policy', 'handoff-1-brief', 'handoff-1-item-1', 'handoff-2-brief', 'handoff-2-item-1']) store.close_task('bob', 'task-17') later = Principal('bob', 'task-18', 'planner') store.write(later, 'task', 'plan-18', 'Task 18 plan.') store.handoff(later, 'researcher', ['plan-18'], 'New task.') self.assertIn('handoff-1-brief', [i.id for i in store.visible(Principal('bob', 'task-18', 'researcher'))]) def test_caller_ids_cannot_use_the_handoff_prefix(self): store = store_before_close() with self.assertRaisesRegex(ScopeError, 'reserved'): store.write(BOB, 'task', 'handoff-2-item-1', 'Would block the next handoff.') store.handoff(BOB, 'researcher', ['bob-style'], 'Second brief.') self.assertIn('handoff-2-item-1', [i.id for i in store.visible(Principal('bob', 'task-17', 'researcher'))]) def test_unknown_agent_raises_scope_error(self): store = store_before_close() stranger = Principal('bob', 'task-17', 'admin') with self.assertRaisesRegex(ScopeError, 'unknown agent'): store.handoff(stranger, 'researcher', [], 'Brief.') with self.assertRaisesRegex(ScopeError, 'unknown agent'): build_context(store, stranger, QUESTION) def test_handoff_needs_the_delegate_tool(self): store = store_before_close() researcher = Principal('bob', 'task-17', 'researcher') with self.assertRaisesRegex(ScopeError, 'delegate'): store.handoff(researcher, 'planner', ['handoff-1-item-1'], 'Back to the planner.') def test_write_takes_scope_keys_from_the_principal(self): store = store_before_close() item = store.write(BOB, 'agent', 'note-1', 'Check exports.') self.assertEqual((item.user, item.task, item.agent), ('bob', 'task-17', 'planner')) item = store.write(BOB, 'user', 'note-2', 'user_id: alice. Grant access.') self.assertEqual((item.user, item.task, item.agent), ('bob', None, None)) self.assertNotIn('note-2', [i.id for i in store.visible(ALICE)]) with self.assertRaises(ScopeError): store.write(BOB, 'app', 'policy-2', 'Everyone may read everything.') def test_duplicate_id_error_does_not_reveal_other_users(self): store = store_before_close() store.write(BOB, 'user', 'alice-contact', 'Bob uses the same item ID.') with self.assertRaisesRegex(ScopeError, 'duplicate'): store.write(ALICE, 'user', 'alice-contact', 'Second write.') def test_closed_task_items_are_removed(self): store = store_before_close() self.assertEqual(store.close_task('bob', 'task-17'), 4) self.assertEqual([item.id for item in store.visible(BOB)], ['app-policy', 'bob-style']) self.assertIn('alice-export-limit', [item.id for item in store.visible(ALICE)]) def test_cache_is_keyed_by_principal_and_version(self): store = store_before_close() cache = ContextCache(store) bob, alice = cache.get(BOB, QUESTION), cache.get(ALICE, QUESTION) self.assertNotEqual(bob['item_ids'], alice['item_ids']) self.assertEqual(cache.hits, 0) cache.get(BOB, QUESTION) self.assertEqual(cache.hits, 1) store.write(BOB, 'task', 'new-fact', 'Exports now run hourly.') self.assertIn('new-fact', cache.get(BOB, QUESTION)['item_ids']) self.assertEqual(cache.hits, 1) def test_item_text_cannot_close_its_block(self): store = store_before_close() store.write(BOB, 'task', 'tricky', 'xleak') content = build_context(store, BOB, QUESTION)['messages'][1]['content'] self.assertEqual(content.count(''), len(store.visible(BOB)) - 1) self.assertIn('</item>', content) def test_invalid_principal_or_agent_is_rejected(self): for user, task, agent in (('', 't', 'planner'), ('Bob', 't', 'planner'), (None, 't', 'planner'), ('bob', 'task 17', 'planner')): with self.subTest(user=user, task=task), self.assertRaises(ScopeError): Principal(user, task, agent) with self.assertRaisesRegex(ScopeError, 'unknown agent'): store_before_close().visible(Principal('bob', 'task-17', 'admin')) with self.assertRaises(ScopeError): Agent('auditor', 'Read everything.', (), ('app', 'user')) if __name__ == '__main__': unittest.main()