yaoyao/integrations/hermes-bots-bridge/tests/test_computer_policy.py

import importlib.util
import json
import sys
import tempfile
import threading
import time
import unittest
from pathlib import Path
from types import SimpleNamespace
from unittest.mock import patch

spec = importlib.util.spec_from_file_location('test_yaoyao_bridge', Path(__file__).parents[1] / 'bridge_runtime.py')
bridge = importlib.util.module_from_spec(spec)
sys.modules[spec.name] = bridge
spec.loader.exec_module(bridge)


class ComputerPolicyTests(unittest.TestCase):
    def setUp(self):
        self.temp = tempfile.TemporaryDirectory()
        self.home = Path(self.temp.name)
        self.agent = SimpleNamespace(session_id='stored', enabled_toolsets=['all'])
        self.session = {'agent': self.agent, 'session_key': 'stored', 'running': True, 'cwd': str(self.home),
                        '_yaoyao_computer_policy': {'mode': 'isolated', 'hostAccess': False}}
        self.server = SimpleNamespace(_sessions={'runtime': self.session}, _sessions_lock=threading.RLock())
        self.binding = bridge.Binding('runtime', 'stored', 'default', self.home, 'generation',
                                     'http://127.0.0.1:12345', 'a' * 40, time.time() * 1000 + 60000,
                                     ('user', 'password'), self.session, self.agent,
                                     computer_policy=dict(self.session['_yaoyao_computer_policy']))
        bridge._bindings.clear(); bridge._generations.clear(); bridge._children.clear()
        bridge._bindings['runtime'] = self.binding
        self.patches = [patch.dict(sys.modules, {'agent.redact': SimpleNamespace(redact_sensitive_text=lambda text, **kw: text)}),
                        patch.object(bridge, '_server', return_value=self.server),
                        patch.object(bridge, '_check_session_mode'),
                        patch.object(bridge, '_transport_identity', return_value=self.binding.owner),
                        patch.object(bridge, '_profile_home', return_value=self.home)]
        for p in self.patches: p.start()

    def tearDown(self):
        for p in reversed(self.patches): p.stop()
        self.temp.cleanup()

    def test_isolated_skills_and_vm_tools_use_native_catalog_but_host_operations_require_approval(self):
        for name in ['skills_list', 'skill_view', 'vision_analyze', 'yaoyao_computer_shell_aabbcc']:
            self.assertIsNone(bridge.computer_directive(session_id='stored', tool_name=name, args={}))
        for name in ['terminal', 'process_manage', 'read_file', 'write_file', 'execute_code', 'delegate_task', 'session_search', 'browser_navigate', 'yaoyao_computer_copy_file_aabbcc']:
            result = bridge.computer_directive(session_id='stored', tool_name=name, args={'path': '/authorized.txt'})
            self.assertEqual(result['action'], 'approve')
            self.assertIn('/authorized.txt', result['message'])
        self.binding.computer_policy['hostAccess'] = True
        self.assertIsNone(bridge.computer_directive(session_id='stored', tool_name='terminal', args={}))
        self.binding.computer_policy.update(mode='profile', hostAccess=False)
        self.assertIsNone(bridge.computer_directive(session_id='stored', tool_name='terminal', args={}))

    def test_stopped_expired_unbound_and_replaced_sessions_cannot_call_native_tools(self):
        for change in ['stop', 'expire', 'unbind', 'replace']:
            with self.subTest(change=change):
                self.session.update(running=True, agent=self.agent)
                self.binding.expires_at = time.time() * 1000 + 60000
                bridge._bindings['runtime'] = self.binding
                if change == 'stop': self.session['running'] = False
                if change == 'expire': self.binding.expires_at = 0
                if change == 'unbind': bridge._bindings.clear()
                if change == 'replace': self.session['agent'] = SimpleNamespace(session_id='stored')
                self.assertEqual(bridge.computer_directive(session_id='runtime', tool_name='skill_view', args={})['action'], 'block')

    def test_ordinary_profile_sessions_are_unaffected(self):
        self.session.pop('_yaoyao_computer_policy')
        self.assertIsNone(bridge.computer_directive(session_id='stored', tool_name='terminal', args={}))

    def test_policy_cannot_be_changed_during_renewal(self):
        self.session['running'] = False
        body = {'session_id': 'runtime', 'stored_session_id': 'stored', 'profile': 'default',
                'generation': 'generation', 'bridge_url': self.binding.bridge_url, 'token': self.binding.token,
                'expires_at': time.time() * 1000 + 60000, 'computer_policy': {'mode': 'profile', 'hostAccess': True}}
        with patch.object(bridge, '_agent_has_tools', return_value=True):
            with self.assertRaises(bridge.BridgeError) as error:
                bridge.bind(self.binding.owner, body)
            self.assertEqual(error.exception.code, 'generation_conflict')

    def test_transfer_is_owned_by_hermes_and_checks_generation_owner_and_protected_paths(self):
        constants = SimpleNamespace(set_hermes_home_override=lambda p: None, reset_hermes_home_override=lambda p: None)
        file_tools = SimpleNamespace(_check_sensitive_path=lambda p, t: 'denied' if p.endswith('protected') else None,
                                     _check_protected_instruction_write=lambda p, t: None,
                                     _check_approval_required_write=lambda p, t: None)
        safety = SimpleNamespace(get_read_block_error=lambda p: 'denied' if p.endswith('protected') else None)
        modules = {'hermes_constants': constants, 'tools.file_tools': file_tools, 'agent.file_safety': safety}
        (self.home / 'authorized').write_bytes(b'authorized-input')
        body = {'session_id': 'runtime', 'generation': 'generation', 'action': 'read', 'path': 'authorized'}
        with patch.dict(sys.modules, modules):
            self.assertEqual(bridge.computer_file(self.binding.owner, body)['data'], 'YXV0aG9yaXplZC1pbnB1dA==')
            for owner, changed in [(('other', 'password'), body), (self.binding.owner, {**body, 'generation': 'stale'}),
                                    (self.binding.owner, {**body, 'path': 'protected'})]:
                with self.assertRaises(bridge.BridgeError): bridge.computer_file(owner, changed)
            bridge.computer_file(self.binding.owner, {**body, 'action': 'write', 'path': 'result', 'data': 'cmVzdWx0'})
            self.assertEqual((self.home / 'result').read_bytes(), b'result')
            self.session['running'] = False
            with self.assertRaises(bridge.BridgeError): bridge.computer_file(self.binding.owner, body)


if __name__ == '__main__':
    unittest.main()