mirror of
https://github.com/saphid/frame-control.git
synced 2026-10-06 01:00:18 +02:00
254 lines
13 KiB
Python
254 lines
13 KiB
Python
"""MCP protocol, exact-action approvals and explicit assistant data sharing."""
|
|
import io
|
|
import json
|
|
import os
|
|
import shutil
|
|
from pathlib import Path
|
|
import subprocess
|
|
import sys
|
|
import tempfile
|
|
import threading
|
|
import unittest
|
|
from unittest import mock
|
|
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
|
|
|
|
sys.path.insert(0, str(Path(__file__).resolve().parents[1] / 'ui'))
|
|
import frame_agent as agent
|
|
import frame_assistant as assistant
|
|
import frame_mcp as mcp
|
|
import server
|
|
|
|
|
|
class Approvals(unittest.TestCase):
|
|
def test_requires_human_decision_exact_action_and_single_use(self):
|
|
gate = agent.Approvals()
|
|
action = {'name': 'power', 'arguments': {'action': 'reboot'}}
|
|
token = gate.request(action)['confirmation']
|
|
with self.assertRaises(ValueError):
|
|
gate.consume(token, action)
|
|
gate.decide(token, True)
|
|
with self.assertRaises(ValueError):
|
|
gate.consume(token, {'name': 'power', 'arguments': {'action': 'poweroff'}})
|
|
gate.consume(token, action)
|
|
with self.assertRaises(ValueError):
|
|
gate.consume(token, action)
|
|
|
|
def test_expiry_rejection_and_non_boolean_approval(self):
|
|
gate = agent.Approvals()
|
|
token = gate.request({})['confirmation']
|
|
gate.decide(token, 'true')
|
|
with self.assertRaises(ValueError):
|
|
gate.inspect(token)
|
|
token = gate.request({})['confirmation']
|
|
with mock.patch.object(agent.time, 'monotonic', return_value=float('inf')):
|
|
with self.assertRaises(ValueError):
|
|
gate.decide(token, True)
|
|
|
|
def test_concurrent_consumption_executes_once(self):
|
|
gate = agent.Approvals()
|
|
token = gate.request({})['confirmation']
|
|
gate.decide(token, True)
|
|
results = []
|
|
def consume():
|
|
try:
|
|
gate.consume(token, {})
|
|
results.append(True)
|
|
except ValueError:
|
|
results.append(False)
|
|
threads = [threading.Thread(target=consume) for _ in range(8)]
|
|
for thread in threads: thread.start()
|
|
for thread in threads: thread.join()
|
|
self.assertEqual(results.count(True), 1)
|
|
|
|
def test_action_never_runs_before_approval(self):
|
|
with mock.patch.object(agent, 'approvals', agent.Approvals()), mock.patch.object(server, 'flatpak') as install:
|
|
body = {'name': 'install', 'arguments': {'id': 'org.example.App'}}
|
|
result = agent.call(server, body)
|
|
install.assert_not_called()
|
|
body['confirmation'] = result['confirmation']
|
|
with self.assertRaises(ValueError): agent.call(server, body)
|
|
agent.approvals.decide(body['confirmation'], True)
|
|
agent.call(server, body)
|
|
install.assert_called_once_with({'id': 'org.example.App', 'action': 'install'})
|
|
with self.assertRaises(ValueError): agent.call(server, body)
|
|
|
|
def test_file_content_change_invalidates_approval(self):
|
|
with tempfile.TemporaryDirectory() as tmp, mock.patch.object(agent, 'approvals', agent.Approvals()), mock.patch.object(server, 'push_file') as push:
|
|
path = Path(tmp) / 'note.txt'
|
|
path.write_text('first')
|
|
body = {'name': 'send_file', 'arguments': {'path': str(path)}}
|
|
result = agent.call(server, body)
|
|
agent.approvals.decide(result['confirmation'], True)
|
|
body['confirmation'] = result['confirmation']
|
|
path.write_text('second')
|
|
with self.assertRaises(ValueError): agent.call(server, body)
|
|
push.assert_not_called()
|
|
|
|
def test_no_arbitrary_commands_or_arguments(self):
|
|
for name, args in [('shell', {'command': 'true'}), ('panel', {'id': 'org.example.App', 'args': '--evil'}),
|
|
('power', {'action': 'factory-reset'}), ('send_text', {'text': ''})]:
|
|
with self.assertRaises(ValueError): agent.call(server, {'name': name, 'arguments': args})
|
|
|
|
|
|
class Assistant(unittest.TestCase):
|
|
def setUp(self):
|
|
self.received = []
|
|
owner = self
|
|
class Endpoint(BaseHTTPRequestHandler):
|
|
def log_message(self, *args): pass
|
|
def do_POST(self):
|
|
owner.received.append((dict(self.headers), json.loads(self.rfile.read(int(self.headers['Content-Length'])))))
|
|
if self.path == '/redirect':
|
|
self.send_response(302)
|
|
self.send_header('Location', '/other')
|
|
self.end_headers()
|
|
return
|
|
data = json.dumps({'choices': [{'message': {'content': '<script>not executed</script>'}}]}).encode()
|
|
self.send_response(200)
|
|
self.send_header('Content-Length', str(len(data)))
|
|
self.end_headers()
|
|
self.wfile.write(data)
|
|
self.httpd = ThreadingHTTPServer(('127.0.0.1', 0), Endpoint)
|
|
self.thread = threading.Thread(target=self.httpd.serve_forever, daemon=True)
|
|
self.thread.start()
|
|
self.body = {'endpoint': 'http://127.0.0.1:%d/chat' % self.httpd.server_port, 'model': 'local', 'prompt': 'Hello', 'consent': True}
|
|
|
|
def tearDown(self):
|
|
self.httpd.shutdown()
|
|
self.httpd.server_close()
|
|
self.thread.join()
|
|
|
|
def test_no_opt_in_no_request_or_capture(self):
|
|
capture = mock.Mock()
|
|
for consent in (False, None, 'true', 1):
|
|
with self.assertRaises(ValueError): assistant.chat({**self.body, 'consent': consent, 'screenshot': True}, capture)
|
|
capture.assert_not_called()
|
|
self.assertEqual(self.received, [])
|
|
|
|
def test_text_only_keyless_and_optional_screenshot(self):
|
|
capture = mock.Mock(return_value=b'png')
|
|
self.assertIn('script', assistant.chat(self.body, capture)['reply'])
|
|
capture.assert_not_called()
|
|
headers, body = self.received[-1]
|
|
self.assertNotIn('Authorization', headers)
|
|
self.assertEqual(body['messages'], [{'role': 'user', 'content': 'Hello'}])
|
|
assistant.chat({**self.body, 'screenshot': True, 'key': 'test-key'}, capture)
|
|
capture.assert_called_once()
|
|
headers, body = self.received[-1]
|
|
self.assertEqual(headers['Authorization'], 'Bearer test-key')
|
|
self.assertEqual(body['messages'][0]['content'][1]['image_url']['url'], 'data:image/png;base64,cG5n')
|
|
|
|
def test_redirects_do_not_forward_context_or_credentials(self):
|
|
with self.assertRaises(ValueError):
|
|
assistant.chat({**self.body, 'endpoint': self.body['endpoint'].replace('/chat', '/redirect'), 'key': 'secret'}, mock.Mock())
|
|
self.assertEqual(len(self.received), 1)
|
|
|
|
def test_bad_urls_fail_before_capture(self):
|
|
for url in ('file:///etc/passwd', 'http://example.com/chat', 'https://user:pass@example.com', 'https://example.com?key=secret'):
|
|
capture = mock.Mock()
|
|
with self.assertRaises(ValueError): assistant.chat({**self.body, 'endpoint': url, 'screenshot': True}, capture)
|
|
capture.assert_not_called()
|
|
|
|
|
|
class AssistantPage(unittest.TestCase):
|
|
@unittest.skipUnless(shutil.which('node'), 'Node is required for the page script regression')
|
|
def test_approval_navigation_races(self):
|
|
root = Path(__file__).resolve().parents[1]
|
|
result = subprocess.run(['node', str(root / 'tests/assistant_ui.cjs'), str(root / 'ui/assistant.html')],
|
|
capture_output=True, text=True, timeout=10)
|
|
self.assertEqual(result.returncode, 0, result.stdout + result.stderr)
|
|
|
|
|
|
class Protocol(unittest.TestCase):
|
|
def test_stdio_initialize_list_call_errors_and_eof(self):
|
|
messages = [
|
|
{'jsonrpc': '2.0', 'id': 1, 'method': 'initialize', 'params': {'protocolVersion': '2025-06-18'}},
|
|
{'jsonrpc': '2.0', 'method': 'notifications/initialized'},
|
|
{'jsonrpc': '2.0', 'id': 2, 'method': 'tools/list'},
|
|
{'jsonrpc': '2.0', 'id': 3, 'method': 'tools/call', 'params': {'name': 'shell'}},
|
|
{'jsonrpc': '2.0', 'id': 4, 'method': 'ping'},
|
|
]
|
|
result = subprocess.run([sys.executable, str(Path(mcp.__file__))], input='\n'.join(map(json.dumps, messages)) + '\n', text=True, capture_output=True, timeout=10)
|
|
self.assertEqual(result.returncode, 0, result.stderr)
|
|
replies = list(map(json.loads, result.stdout.splitlines()))
|
|
self.assertEqual([r['id'] for r in replies], [1, 2, 3, 4])
|
|
self.assertEqual(replies[0]['result']['protocolVersion'], '2025-06-18')
|
|
self.assertIn('screenshot', [t['name'] for t in replies[1]['result']['tools']])
|
|
self.assertTrue(replies[2]['result']['isError'])
|
|
|
|
def test_mcp_cannot_approve_and_returns_review_url(self):
|
|
client = mock.Mock(url='http://127.0.0.1:47810')
|
|
client.request.return_value = {'approvalPath': '/assistant#confirm=token'}
|
|
result = mcp.call(client, 'power', {'action': 'reboot'})
|
|
self.assertIn('http://127.0.0.1:47810/assistant', result['content'][0]['text'])
|
|
with self.assertRaises(ValueError): mcp.call(client, 'approve', {'confirmation': 'token'})
|
|
with self.assertRaises(ValueError): mcp.call(client, 'status', {'path': '/api/open'})
|
|
|
|
def test_loopback_only_backend(self):
|
|
for url in ('https://example.com', 'http://127.0.0.1/api', 'http://secret@localhost:1234', 'file:///tmp/x'):
|
|
with self.assertRaises(ValueError): mcp.Client(url)
|
|
|
|
|
|
class ManagedBackend(unittest.TestCase):
|
|
def test_private_backend_auth_and_cleanup(self):
|
|
from urllib.error import HTTPError, URLError
|
|
from urllib.request import urlopen
|
|
with mock.patch.dict(os.environ, {'FRAME_ALIAS': 'frame-control-test.invalid'}):
|
|
with mcp.backend() as client:
|
|
url = client.url
|
|
self.assertIn('os', client.request('/api/host'))
|
|
with self.assertRaises(HTTPError) as error:
|
|
urlopen(url + '/api/host', timeout=2)
|
|
self.assertEqual(error.exception.code, 403)
|
|
error.exception.close()
|
|
# A second client has its own backend and key.
|
|
with mcp.backend() as other:
|
|
self.assertNotEqual(client.url, other.url)
|
|
self.assertNotEqual(client.key, other.key)
|
|
self.assertIn('os', client.request('/api/host'))
|
|
with self.assertRaises(URLError):
|
|
urlopen(url + '/', timeout=2)
|
|
|
|
def test_private_ssh_socket_is_not_the_desktop_socket(self):
|
|
with mock.patch.object(server.frame_host, 'MUX', True), \
|
|
mock.patch.object(server.frame_host.os, 'getuid', return_value=501, create=True), \
|
|
mock.patch.object(server.frame_host.os, 'getpid', return_value=123):
|
|
self.assertEqual(server.frame_host.control_path(), '/tmp/frame-ui-501-%C')
|
|
self.assertEqual(server.frame_host.control_path(private=True), '/tmp/frame-ui-501-123-%C')
|
|
|
|
|
|
class ComputerState(unittest.TestCase):
|
|
def test_gamescope_triplets_and_empty_focus(self):
|
|
import frame_computer
|
|
parsed = frame_computer.parse_windows('GAMESCOPE_FOCUSABLE_WINDOWS(CARDINAL) = 16, 42, 123, 32, 55, 999\nGAMESCOPE_FOCUSED_APP(CARDINAL) = \n')
|
|
self.assertEqual(parsed['windows'], [{'windowId': '0x10', 'appid': 42, 'pid': 123}, {'windowId': '0x20', 'appid': 55, 'pid': 999}])
|
|
self.assertIsNone(parsed['focusedApp'])
|
|
with self.assertRaises(ValueError):
|
|
frame_computer.parse_windows('GAMESCOPE_FOCUSABLE_WINDOWS(CARDINAL) = 1, 2')
|
|
with self.assertRaises(ValueError):
|
|
frame_computer.parse_windows('GAMESCOPE_FOCUSABLE_WINDOWS(CARDINAL) = untrusted')
|
|
with self.assertRaises(ValueError):
|
|
frame_computer.parse_windows('GAMESCOPE_FOCUSABLE_WINDOWS: no such atom on any window.')
|
|
|
|
def test_partial_snapshot_reports_failure_not_empty_success(self):
|
|
import frame_computer
|
|
with mock.patch.object(frame_computer.subprocess, 'run', side_effect=OSError('no display')), \
|
|
mock.patch.object(frame_computer, 'accessibility', side_effect=OSError('no AT-SPI')):
|
|
result = frame_computer.snapshot()
|
|
self.assertIn('windowError', result)
|
|
self.assertIn('accessibilityError', result)
|
|
self.assertFalse(result['inputEnabled'])
|
|
self.assertNotIn('windows', result)
|
|
|
|
def test_mcp_computer_state_is_read_only(self):
|
|
client = mock.Mock()
|
|
client.request.return_value = {'windows': []}
|
|
mcp.call(client, 'computer_state', {})
|
|
client.request.assert_called_once_with('/api/computer/state')
|
|
spec = next(t for t in mcp.TOOLS if t['name'] == 'computer_state')
|
|
self.assertTrue(spec['annotations']['readOnlyHint'])
|
|
|
|
|
|
if __name__ == '__main__':
|
|
unittest.main()
|