import hashlib import copy import json import os import socket import struct import sys import tempfile import threading import time import unittest from types import SimpleNamespace from unittest import mock APP_DIR = os.path.abspath(os.path.join(os.path.dirname(__file__), '..', 'app')) if APP_DIR not in sys.path: sys.path.insert(0, APP_DIR) import worker_supervisor import remote_worker_client from process_identity import current_process_identity from runtime_security import ( PrivateFileLock, atomic_write_private_json, ensure_private_directory, write_private_json_exclusive, ) from worker_local_state import WorkerLocalState from worker_supervisor import ( WorkerAlreadyRunning, WorkerSupervisor, WorkerSupervisorError, build_instance_record, classify_instance, encode_frame, instance_path, load_startup_result, load_shutdown_receipt, receive_frame, send_control_request, ) PACKAGE = { 'schema': 3, 'manifest_sha256': 'a' * 64, 'code_manifest_sha256': 'b' * 64, 'platform_tag': 'windows-x86_64' if os.name == 'nt' else 'linux-x86_64', } PROTOCOL = { 'worker_protocol': 2, 'bundle_format': 2, 'event': 1, 'control': 1, 'projection': 1, } UTC = '2026-09-23T12:00:00Z' def fixture_args(root): return SimpleNamespace( state_dir=os.path.join(root, 'state'), bundle_dir=os.path.join(root, 'data', 'bundles'), work_dir=os.path.join(root, 'data', 'work'), package_manifest=os.path.join(root, 'worker-package.json'), server='https://worker.example', token='x' * 32, parallelism=1, poll_seconds=0.01, error_delay_seconds=0.01, http_timeout=3, ) class WorkerSupervisorTests(unittest.TestCase): @staticmethod def compatibility(): return { 'protocol_version': 2, 'bundle_format_version': 2, 'platform_tag': PACKAGE['platform_tag'], 'code_manifest_sha256': 'a' * 64, 'detector_policy_sha256': 'b' * 64, } def test_recovered_stale_cleanup_emits_idle_before_next_claiming(self): class API: def claim(self, _request_id, _compatibility): return {'retry_after_seconds': 1, 'reason': 'empty_queue'} with tempfile.TemporaryDirectory() as root: local = WorkerLocalState(root) terminal = [] def event(value): local.emit_phase( 'instance-1', value['slot_id'], value['phase'], reservation_id=value.get('reservation_id'), source=value.get('source'), progress=value.get('progress'), assignment_deadline_at=value.get('assignment_deadline_at'), ) slot = remote_worker_client.WorkerSlot( 0, API(), self.compatibility(), root, root, event_callback=event, terminal_callback=terminal.append, ) assignment = {'reservation': { 'reservation_id': 7, 'bundle_id': 'b' * 32, 'scan_event_id': 'c' * 32, 'source': 'gitlab', 'attempts': 1, 'remote_expires_at': '2999-01-01T00:00:00+00:00', }} atomic_write_private_json(slot.state_path, { 'phase': 'assigned', 'assignment': assignment, }) atomic_write_private_json(slot.stale_path, { 'schema': 1, 'outcome': 'discarded_stale', 'reservation_id': 7, 'bundle_id': 'b' * 32, 'scan_event_id': 'c' * 32, 'status_code': 410, 'code': 'assignment_expired', 'recorded_at': UTC, }) self.assertTrue(slot.step()) self.assertEqual(local.snapshot()['slots'][0]['phase'], 'idle') self.assertEqual(len(terminal), 1) self.assertFalse(slot.step()) phases = [item['phase'] for item in local.events_after(0, 100)] idle_index = max(index for index, phase in enumerate(phases) if phase == 'idle') claiming_index = max(index for index, phase in enumerate(phases) if phase == 'claiming') self.assertLess(idle_index, claiming_index) def test_runner_timestamps_drive_history_durations_not_controller_receipt_delay(self): with tempfile.TemporaryDirectory() as root, mock.patch.object( worker_supervisor, '_package_identity', return_value=(PACKAGE, PROTOCOL), ): runtime = WorkerSupervisor(fixture_args(root)) runtime._initialize_local_state() runtime._event({ 'slot_id': 0, 'phase': 'scanning', 'reservation_id': 17, 'source': 'gitlab', 'progress': {}, 'timestamp': '2026-09-23T12:00:00Z', 'phase_started_at': '2026-09-23T12:00:00Z', }) runtime._event({ 'slot_id': 0, 'phase': 'filtering', 'reservation_id': 17, 'source': 'gitlab', 'progress': {}, 'timestamp': '2026-09-23T12:00:01Z', 'phase_started_at': '2026-09-23T12:00:01Z', }) runtime._terminal({ 'history_id': 'a' * 64, 'slot_id': 0, 'reservation_id': 17, 'source': 'gitlab', 'outcome': 'bundle_accepted', 'receipt': {}, 'started_at': '2026-09-23T12:00:00Z', 'completed_at': '2026-09-23T12:00:02Z', 'duration_seconds': 2.0, 'first_sequence': None, 'diagnostics': [], }) record = runtime.local.history(limit=1)[0] self.assertEqual(record['phase_durations']['scanning'], 1.0) self.assertEqual(record['phase_durations']['filtering'], 1.0) def test_supervisor_restart_recovers_high_slot_then_records_terminal_history(self): with tempfile.TemporaryDirectory() as root, mock.patch.object( worker_supervisor, '_package_identity', return_value=(PACKAGE, PROTOCOL), ): args = fixture_args(root) args.parallelism = 1 first = WorkerSupervisor(args) first._initialize_local_state() assignment = {'reservation': { 'reservation_id': 77, 'bundle_id': 'b' * 32, 'scan_event_id': 'c' * 32, 'source': 'gitlab', 'platform': 'gitlab', 'attempts': 1, 'remote_issued_at': '2026-09-23T12:00:00Z', 'remote_expires_at': '2999-01-01T00:00:00Z', }} first.local.emit_phase( 'old-instance', 7, remote_worker_client.WorkerPhase.ASSIGNED, reservation_id=77, source='gitlab', timestamp='2026-09-23T12:00:00Z', ) atomic_write_private_json( os.path.join(args.state_dir, 'slot-7.json'), {'phase': 'assigned', 'assignment': assignment}, ) atomic_write_private_json( os.path.join(args.state_dir, 'slot-7-stale.json'), { 'schema': 1, 'outcome': 'discarded_stale', 'reservation_id': 77, 'bundle_id': 'b' * 32, 'scan_event_id': 'c' * 32, 'status_code': 410, 'code': 'assignment_expired', 'recorded_at': '2026-09-23T12:00:01Z', }, ) restarted = WorkerSupervisor(args) restarted._initialize_local_state() self.assertEqual(restarted._public_worker()['recovery_slots'], 1) slot = remote_worker_client.WorkerSlot( 7, object(), self.compatibility(), args.state_dir, ensure_private_directory(args.bundle_dir, reject_reparse=True), work_root=ensure_private_directory(args.work_dir, reject_reparse=True), claim_enabled=False, event_callback=restarted._event, terminal_callback=restarted._terminal, ) for name in ('tmp', 'ready', 'quarantine'): ensure_private_directory(os.path.join(args.bundle_dir, name), reject_reparse=True) self.assertTrue(slot.step()) self.assertEqual(restarted._public_worker()['recovery_slots'], 0) history = restarted.local.history(limit=1) self.assertEqual(history[0]['reservation_id'], 77) self.assertEqual(history[0]['slot_id'], 7) def test_run_client_twice_recovers_slot_above_lowered_cap_and_wires_history(self): with tempfile.TemporaryDirectory() as root, mock.patch.object( worker_supervisor, '_package_identity', return_value=(PACKAGE, PROTOCOL), ): args = fixture_args(root) args.parallelism = 8 first = WorkerSupervisor(args) first._initialize_local_state() assignment = {'reservation': { 'reservation_id': 87, 'bundle_id': 'e' * 32, 'scan_event_id': 'f' * 32, 'source': 'gitlab', 'platform': 'gitlab', 'attempts': 1, 'remote_issued_at': '2026-09-23T12:00:00Z', 'remote_expires_at': '2999-01-01T00:00:00Z', }} first.local.emit_phase( 'engine-before-restart', 7, remote_worker_client.WorkerPhase.ASSIGNED, reservation_id=87, source='gitlab', timestamp='2026-09-23T12:00:00Z', ) atomic_write_private_json( os.path.join(args.state_dir, 'slot-7.json'), {'phase': 'assigned', 'assignment': assignment}, ) atomic_write_private_json( os.path.join(args.state_dir, 'slot-7-stale.json'), { 'schema': 1, 'outcome': 'discarded_stale', 'reservation_id': 87, 'bundle_id': 'e' * 32, 'scan_event_id': 'f' * 32, 'status_code': 410, 'code': 'assignment_expired', 'recorded_at': '2026-09-23T12:00:01Z', }, ) package = { 'manifest': {'capabilities': [{ 'source': 'gitlab', 'platform': 'gitlab', 'planning_kind': 'exact_git_v1', }]}, 'build_compatibility': self.compatibility(), 'runtime_trees': {}, 'code_manifest': {}, 'code_manifest_sha256': 'a' * 64, 'trufflehog_path': 'trufflehog', 'git_path': 'git', 'detector_policy_path': 'detectors.yaml', } class API: @staticmethod def claim(_request_id, _compatibility): return None real_slot = remote_worker_client.WorkerSlot first_flags = {} class FirstSlot(real_slot): def __init__(self, slot_id, *slot_args, **kwargs): first_flags[slot_id] = kwargs.get('claim_enabled') super().__init__(slot_id, *slot_args, **kwargs) stopped = threading.Event() stopped.set() with mock.patch.object( remote_worker_client, 'verify_worker_package', side_effect=lambda _path: copy.deepcopy(package), ), mock.patch.object( remote_worker_client, 'WorkerHTTPClient', return_value=API(), ), mock.patch.object( remote_worker_client, 'WorkerSlot', FirstSlot, ): self.assertEqual(remote_worker_client.run_client( args, stop_event=stopped, acquire_lock=False, ), 2) self.assertTrue(first_flags[7]) self.assertIn(7, remote_worker_client.persisted_slot_ids(args.state_dir)) args.parallelism = 1 restarted = WorkerSupervisor(args) restarted._initialize_local_state() second_flags = {} drain = threading.Event() engine_stop = threading.Event() failures = [] class SecondSlot(real_slot): def __init__(self, slot_id, *slot_args, **kwargs): second_flags[slot_id] = kwargs.get('claim_enabled') super().__init__(slot_id, *slot_args, **kwargs) def step(self): try: return super().step() except BaseException as exc: failures.append((self.slot_id, exc)) engine_stop.set() raise def terminal(value): restarted._terminal(value) def drain_after_recovery_retirement(): deadline = time.monotonic() + 5 while time.monotonic() < deadline: slots = restarted.local.snapshot()['slots'] recovered = next( (item for item in slots if item['slot_id'] == 7), None, ) if recovered is not None and recovered['phase'] == 'stopped': drain.set() return time.sleep(0.01) engine_stop.set() threading.Thread( target=drain_after_recovery_retirement, daemon=True, ).start() with mock.patch.object( remote_worker_client, 'verify_worker_package', side_effect=lambda _path: copy.deepcopy(package), ), mock.patch.object( remote_worker_client, 'WorkerHTTPClient', return_value=API(), ), mock.patch.object( remote_worker_client, 'WorkerSlot', SecondSlot, ): exit_code = remote_worker_client.run_client( args, drain_event=drain, stop_event=engine_stop, event_callback=restarted._event, terminal_callback=terminal, acquire_lock=False, ) self.assertEqual(exit_code, 0, repr(failures)) self.assertEqual(failures, []) self.assertFalse(second_flags[7]) self.assertNotIn(7, remote_worker_client.persisted_slot_ids(args.state_dir)) slot = next( item for item in restarted.local.snapshot()['slots'] if item['slot_id'] == 7 ) self.assertEqual(slot['phase'], 'stopped') self.assertEqual( slot['progress']['reason'], 'lowered_parallelism_recovery_complete', ) history = restarted.local.history(limit=1) self.assertEqual(history[0]['reservation_id'], 87) self.assertEqual(history[0]['slot_id'], 7) self.assertEqual( {item['instance_id'] for item in history[0]['timeline']}, {'engine-before-restart', restarted.instance_id}, ) def test_terminal_failure_archives_actual_material_and_passes_history_references(self): class Failure(RuntimeError): status_code = 500 body = b'actual body' stdout = b'actual stdout\n' stderr = b'actual stderr\n' returncode = 9 class API: def status(self, _reservation_id): return {'state': 'scanning', 'expires_at': '2999-01-01T00:00:00Z'} def terminal(self, reservation_id, report): payload = json.dumps( report, ensure_ascii=True, sort_keys=True, separators=(',', ':'), ).encode('utf-8') return { 'reservation_id': reservation_id, 'bundle_id': 'b' * 32, 'scan_event_id': 'c' * 32, 'resolution': 'prebundle_report', 'receipt_id': 'd' * 64, 'failure_code': report['failure_code'], 'payload_sha256': hashlib.sha256(payload).hexdigest(), } with tempfile.TemporaryDirectory() as root: local = WorkerLocalState(root) terminal = [] slot = remote_worker_client.WorkerSlot( 0, API(), self.compatibility(), root, root, event_callback=lambda value: local.emit_phase( 'instance-1', value['slot_id'], value['phase'], reservation_id=value.get('reservation_id'), source=value.get('source'), progress=value.get('progress'), assignment_deadline_at=value.get('assignment_deadline_at'), ), terminal_callback=terminal.append, diagnostic_callback=local.archive_diagnostic, ) atomic_write_private_json(slot.state_path, { 'phase': 'assigned', 'assignment': {'reservation': { 'reservation_id': 7, 'bundle_id': 'b' * 32, 'scan_event_id': 'c' * 32, 'source': 'dockerhub', 'platform': 'docker', 'attempts': 1, 'remote_issued_at': '2026-09-23T12:00:00Z', 'remote_expires_at': '2999-01-01T00:00:00Z', }}, }) with mock.patch.object(slot, '_execute', side_effect=Failure('actual exception')): self.assertTrue(slot.step()) self.assertEqual(len(terminal[0]['diagnostics']), 1) reference = terminal[0]['diagnostics'][0] self.assertTrue(reference['artifacts']['stdout']['path'].endswith('.stdout.log')) self.assertEqual( open(os.path.join(root, *reference['artifacts']['body']['path'].split('/')), 'rb').read(), b'actual body', ) def test_diagnostic_combined_log_budget_and_full_local_archive(self): class Failure(RuntimeError): body = b'b' * 20000 stdout = b'o' * 40000 stderr = b'e' * 40000 returncode = 9 with tempfile.TemporaryDirectory() as root: local = WorkerLocalState(root) captured = [] def archive(envelope, materials): captured.append(envelope) return local.archive_diagnostic(envelope, materials) slot = remote_worker_client.WorkerSlot( 0, object(), self.compatibility(), root, root, diagnostic_callback=archive, ) slot._event_phase = remote_worker_client.WorkerPhase.SCANNING state = {'assignment': {'reservation': { 'reservation_id': 7, 'scan_event_id': 'c' * 32, 'source': 'dockerhub', 'platform': 'docker', 'attempts': 1, }}} reference = slot._capture_failure_diagnostic( state, Failure('full material'), 'client_process_failed', remote_worker_client.DiagnosticCategory.INTERNAL, ) report = slot._fitted_terminal_report( 'client_process_failed', 'local worker execution failed', ) slot._archive_fitted_terminal_diagnostics(report) reference = slot._diagnostics[0] process = captured[0].process self.assertLessEqual( process.stdout.stored_size + process.stderr.stored_size, remote_worker_client.MAX_DIAGNOSTIC_LOG_BYTES, ) self.assertTrue(process.stdout.truncated) self.assertTrue(process.stderr.truncated) self.assertEqual( os.path.getsize(os.path.join(root, *reference['artifacts']['stdout']['path'].split('/'))), 40000, ) self.assertEqual( os.path.getsize(os.path.join(root, *reference['artifacts']['stderr']['path'].split('/'))), 40000, ) self.assertEqual( os.path.getsize(os.path.join(root, *reference['artifacts']['body']['path'].split('/'))), 20000, ) def test_diagnostic_archive_failure_does_not_block_terminal_reconciliation(self): class Failure(RuntimeError): stdout = b'actual stdout' class API: def status(self, _reservation_id): return {'state': 'scanning', 'expires_at': '2999-01-01T00:00:00Z'} def terminal(self, reservation_id, report): payload = json.dumps( report, ensure_ascii=True, sort_keys=True, separators=(',', ':'), ).encode('utf-8') return { 'reservation_id': reservation_id, 'bundle_id': 'b' * 32, 'scan_event_id': 'c' * 32, 'resolution': 'prebundle_report', 'receipt_id': 'd' * 64, 'failure_code': report['failure_code'], 'payload_sha256': hashlib.sha256(payload).hexdigest(), } with tempfile.TemporaryDirectory() as root: ensure_private_directory(root, reject_reparse=True) terminal = [] slot = remote_worker_client.WorkerSlot( 0, API(), self.compatibility(), root, root, terminal_callback=terminal.append, diagnostic_callback=mock.Mock(side_effect=OSError('archive unavailable')), ) atomic_write_private_json(slot.state_path, { 'phase': 'assigned', 'assignment': {'reservation': { 'reservation_id': 7, 'bundle_id': 'b' * 32, 'scan_event_id': 'c' * 32, 'source': 'dockerhub', 'platform': 'docker', 'attempts': 1, 'remote_issued_at': UTC, 'remote_expires_at': '2999-01-01T00:00:00Z', }}, }) with mock.patch.object(slot, '_execute', side_effect=Failure('failure')): self.assertTrue(slot.step()) self.assertEqual(len(terminal), 1) self.assertEqual(terminal[0]['diagnostics'], []) def test_client_drain_finishes_persisted_assignment_without_new_claim(self): stepped = [] class DrainingSlot: def __init__(self, slot_id, _api, _compatibility, state_dir, *_args, **kwargs): self.slot_id = slot_id self.state_path = os.path.join(state_dir, f'slot-{slot_id}.json') self.claim_enabled = kwargs['claim_enabled'] self.retry_after_seconds = None def step(self): stepped.append((self.slot_id, self.claim_enabled)) os.remove(self.state_path) return True package = { 'manifest': {'capabilities': [{ 'source': 'gitlab', 'platform': 'gitlab', 'planning_kind': 'exact_git_v1', }]}, 'build_compatibility': { 'protocol_version': 2, 'bundle_format_version': 2, 'platform_tag': PACKAGE['platform_tag'], 'code_manifest_sha256': 'a' * 64, 'detector_policy_sha256': 'b' * 64, }, 'code_manifest': {}, 'code_manifest_sha256': 'a' * 64, 'trufflehog_path': 'trufflehog', 'git_path': 'git', 'detector_policy_path': 'detectors.yaml', 'runtime_trees': {}, } with tempfile.TemporaryDirectory() as root: args = fixture_args(root) ensure_private_directory(args.state_dir, reject_reparse=True) with open(os.path.join(args.state_dir, 'slot-0.json'), 'w', encoding='ascii') as handle: handle.write('{}') drain = threading.Event() drain.set() with mock.patch.object( remote_worker_client, 'verify_worker_package', return_value=package, ), mock.patch.object( remote_worker_client, 'WorkerHTTPClient', return_value=object(), ), mock.patch.object(remote_worker_client, 'WorkerSlot', DrainingSlot): remote_worker_client.run_client( args, drain_event=drain, acquire_lock=False, ) self.assertEqual(stepped, [(0, False)]) def test_client_startup_waits_for_every_slot_thread_before_ready(self): order = [] class ReadySlot: def __init__(self, slot_id, _api, _compatibility, state_dir, *_args, **kwargs): order.append(f'init-{slot_id}') self.slot_id = slot_id self.state_path = os.path.join(state_dir, f'slot-{slot_id}.json') self.claim_enabled = kwargs['claim_enabled'] self.retry_after_seconds = None package = { 'manifest': {'capabilities': [{ 'source': 'gitlab', 'platform': 'gitlab', 'planning_kind': 'exact_git_v1', }]}, 'build_compatibility': self.compatibility(), 'code_manifest': {}, 'code_manifest_sha256': 'a' * 64, 'trufflehog_path': 'trufflehog', 'git_path': 'git', 'detector_policy_path': 'detectors.yaml', 'runtime_trees': {}, } with tempfile.TemporaryDirectory() as root: args = fixture_args(root) args.parallelism = 2 drain = threading.Event() drain.set() with mock.patch.object( remote_worker_client, 'verify_worker_package', return_value=package, ), mock.patch.object( remote_worker_client, 'WorkerHTTPClient', return_value=object(), ), mock.patch.object(remote_worker_client, 'WorkerSlot', ReadySlot): remote_worker_client.run_client( args, drain_event=drain, acquire_lock=False, started_callback=lambda: order.append('ready'), ) self.assertEqual(set(order[:2]), {'init-0', 'init-1'}) self.assertEqual(order[2], 'ready') def test_client_slot_initialization_failure_never_publishes_ready(self): package = { 'manifest': {'capabilities': [{ 'source': 'gitlab', 'platform': 'gitlab', 'planning_kind': 'exact_git_v1', }]}, 'build_compatibility': self.compatibility(), 'code_manifest': {}, 'code_manifest_sha256': 'a' * 64, 'trufflehog_path': 'trufflehog', 'git_path': 'git', 'detector_policy_path': 'detectors.yaml', 'runtime_trees': {}, } with tempfile.TemporaryDirectory() as root: args = fixture_args(root) ready = mock.Mock() with mock.patch.object( remote_worker_client, 'verify_worker_package', return_value=package, ), mock.patch.object( remote_worker_client, 'WorkerHTTPClient', return_value=object(), ), mock.patch.object( remote_worker_client, 'WorkerSlot', side_effect=RuntimeError('slot init'), ), self.assertRaisesRegex(remote_worker_client.WorkerClientError, 'startup failed'): remote_worker_client.run_client( args, acquire_lock=False, started_callback=ready, ) ready.assert_not_called() def test_client_forced_stop_returns_without_waiting_for_blocked_slot(self): entered = threading.Event() release = threading.Event() class BlockingSlot: def __init__(self, slot_id, _api, _compatibility, state_dir, *_args, **kwargs): self.slot_id = slot_id self.state_path = os.path.join(state_dir, f'slot-{slot_id}.json') self.claim_enabled = kwargs['claim_enabled'] self.retry_after_seconds = None atomic_write_private_json(self.state_path, { 'phase': 'assigned', 'assignment': {'reservation': {'reservation_id': 7}}, }) def step(self): entered.set() release.wait(5) return False package = { 'manifest': {'capabilities': [{ 'source': 'gitlab', 'platform': 'gitlab', 'planning_kind': 'exact_git_v1', }]}, 'build_compatibility': self.compatibility(), 'code_manifest': {}, 'code_manifest_sha256': 'a' * 64, 'trufflehog_path': 'trufflehog', 'git_path': 'git', 'detector_policy_path': 'detectors.yaml', 'runtime_trees': {}, } with tempfile.TemporaryDirectory() as root: args = fixture_args(root) stopping = threading.Event() def started(): def stop_after_entry(): entered.wait(5) stopping.set() threading.Thread(target=stop_after_entry, daemon=True).start() with mock.patch.object( remote_worker_client, 'verify_worker_package', return_value=package, ), mock.patch.object( remote_worker_client, 'WorkerHTTPClient', return_value=object(), ), mock.patch.object(remote_worker_client, 'WorkerSlot', BlockingSlot): before = time.monotonic() code = remote_worker_client.run_client( args, stop_event=stopping, acquire_lock=False, started_callback=started, ) elapsed = time.monotonic() - before self.assertEqual(code, 2) self.assertLess(elapsed, 2.0) self.assertTrue(os.path.exists(os.path.join(args.state_dir, 'slot-0.json'))) release.set() time.sleep(0.1) def test_slot_event_callback_canonicalizes_server_deadline(self): observed = [] with tempfile.TemporaryDirectory() as root: slot = remote_worker_client.WorkerSlot( 0, object(), { 'protocol_version': 2, 'bundle_format_version': 2, 'platform_tag': PACKAGE['platform_tag'], 'code_manifest_sha256': 'a' * 64, 'detector_policy_sha256': 'b' * 64, }, root, root, event_callback=observed.append, ) slot._emit('assigned', { 'assignment': {'reservation': { 'reservation_id': 7, 'source': 'gitlab', 'attempts': 0, 'remote_expires_at': '2026-09-23T12:00:00+00:00', }}, }) self.assertEqual( observed[0]['assignment_deadline_at'], '2026-09-23T12:00:00Z', ) def test_length_prefixed_frames_require_bounded_canonical_json(self): left, right = socket.socketpair() try: left.sendall(encode_frame({'a': 1, 'b': 2}, 1024)) self.assertEqual(receive_frame(right, 1024), {'a': 1, 'b': 2}) finally: left.close() right.close() left, right = socket.socketpair() try: payload = b'{"b":2,"a":1}' left.sendall(struct.pack('!I', len(payload)) + payload) with self.assertRaisesRegex(WorkerSupervisorError, 'canonical'): receive_frame(right, 1024) finally: left.close() right.close() def test_control_handshake_authenticates_per_instance_token(self): with tempfile.TemporaryDirectory() as root, mock.patch.object( worker_supervisor, '_package_identity', return_value=(PACKAGE, PROTOCOL), ): runtime = WorkerSupervisor(fixture_args(root)) runtime._publish_instance() try: result = send_control_request(runtime.record, 'handshake', {}) self.assertEqual(result['instance_id'], runtime.instance_id) classification = classify_instance(runtime.state_dir) self.assertEqual(classification['state'], 'starting') wrong = dict(runtime.record) wrong['token'] = 'z' * 48 with self.assertRaisesRegex(WorkerSupervisorError, 'authentication'): send_control_request(wrong, 'handshake', {}) finally: runtime._remove_instance() runtime._shutdown_control() def test_progress_cursor_preparation_precedes_maintenance_and_client_start(self): order = [] def fake_client(_args, **kwargs): order.append('client') kwargs['started_callback']() return 0 with tempfile.TemporaryDirectory() as root, mock.patch.object( worker_supervisor, '_package_identity', return_value=(PACKAGE, PROTOCOL), ), mock.patch.object( worker_supervisor, 'prepare_progress_outbox_cursor', side_effect=lambda *_args, **_kwargs: order.append('cursor'), ), mock.patch.object( worker_supervisor, 'run_client', side_effect=fake_client, ): runtime = WorkerSupervisor(fixture_args(root)) original_start = runtime._start_maintenance def start_maintenance(): order.append('maintenance') original_start() runtime._start_maintenance = start_maintenance self.assertEqual(runtime.run(), 0) self.assertLess(order.index('cursor'), order.index('maintenance')) self.assertLess(order.index('cursor'), order.index('client')) def test_single_instance_graceful_drain_and_shutdown_receipt(self): entered = threading.Event() def fake_client(args, **kwargs): kwargs['started_callback']() entered.set() kwargs['event_callback']({'slot_id': 0, 'phase': 'idle', 'progress': {}}) self.assertTrue(kwargs['drain_event'].wait(5)) return 0 with tempfile.TemporaryDirectory() as root, mock.patch.object( worker_supervisor, '_package_identity', return_value=(PACKAGE, PROTOCOL), ), mock.patch.object(worker_supervisor, 'run_client', side_effect=fake_client): runtime = WorkerSupervisor(fixture_args(root)) startup_file = os.path.join(runtime.state_dir, 'control', 'startup.json') thread = threading.Thread( target=lambda: runtime.run(startup_file, 'launch-nonce'), ) thread.start() self.assertTrue(entered.wait(5)) startup = load_startup_result(startup_file, 'launch-nonce') self.assertEqual(startup['outcome'], 'ready') self.assertEqual(startup['instance_id'], runtime.instance_id) classification = classify_instance(runtime.state_dir) self.assertEqual(classification['state'], 'running') projection_before = open(runtime.local.status_path, 'rb').read() second = WorkerSupervisor(fixture_args(root)) with self.assertRaises(WorkerAlreadyRunning): second.run() self.assertEqual(open(runtime.local.status_path, 'rb').read(), projection_before) accepted = send_control_request( classification['record'], 'stop', {'timeout_seconds': 5.0}, ) self.assertTrue(accepted['accepted']) thread.join(timeout=10) self.assertFalse(thread.is_alive()) receipt = load_shutdown_receipt(runtime.state_dir, runtime.instance_id) self.assertEqual(receipt['exit_code'], 0) self.assertTrue(receipt['drained']) def test_stale_is_removed_only_for_exact_dead_or_reused_identity(self): with tempfile.TemporaryDirectory() as root: state_dir = ensure_private_directory(os.path.join(root, 'state'), reject_reparse=True) ensure_private_directory(os.path.join(state_dir, 'control'), reject_reparse=True) identity = current_process_identity() record = build_instance_record( instance_id='stale-instance', token='t' * 48, identity=identity, package=PACKAGE, protocol=PROTOCOL, control_port=12345, foreground=False, ) write_private_json_exclusive(instance_path(state_dir), record) unknown = classify_instance(state_dir, identity_state=lambda *_args: 'unknown') self.assertEqual(unknown['state'], 'unverifiable') self.assertTrue(os.path.exists(instance_path(state_dir))) stale = classify_instance(state_dir, identity_state=lambda *_args: 'reused') self.assertEqual(stale['state'], 'stale') self.assertFalse(stale['removed']) self.assertTrue(os.path.exists(instance_path(state_dir))) recovered = classify_instance( state_dir, identity_state=lambda *_args: 'reused', remove_stale=True, ) self.assertTrue(recovered['removed']) self.assertFalse(os.path.exists(instance_path(state_dir))) def test_drain_deadline_forces_nonclean_exit_and_preserves_slot_authority(self): entered = threading.Event() result = [] controller_release = threading.Event() controller_threads = [] hard_exit_receipts = [] def blocked_client(args, **kwargs): kwargs['started_callback']() atomic_write_private_json(os.path.join(args.state_dir, 'slot-0.json'), { 'phase': 'assigned', 'assignment': {'reservation': {'reservation_id': 7}}, }) controller = threading.Thread( target=lambda: controller_release.wait(10), daemon=True, ) controller.start() controller_threads.append(controller) kwargs['event_callback']({ 'slot_id': 0, 'phase': 'idle', 'progress': {}, }) entered.set() self.assertTrue(kwargs['stop_event'].wait(5)) return 2 def hard_exit(code): receipt = load_shutdown_receipt(runtime.state_dir, runtime.instance_id) hard_exit_receipts.append((code, receipt)) with tempfile.TemporaryDirectory() as root, mock.patch.object( worker_supervisor, '_package_identity', return_value=(PACKAGE, PROTOCOL), ), mock.patch.object(worker_supervisor, 'run_client', side_effect=blocked_client): runtime = WorkerSupervisor( fixture_args(root), retention_interval_seconds=60, hard_exit_hook=hard_exit, ) thread = threading.Thread(target=lambda: result.append(runtime.run())) thread.start() self.assertTrue(entered.wait(5)) classification = classify_instance(runtime.state_dir) self.assertEqual(classification['state'], 'running') send_control_request( classification['record'], 'stop', {'timeout_seconds': 0.1}, ) thread.join(timeout=10) self.assertFalse(thread.is_alive()) self.assertEqual(result, [2]) receipt = load_shutdown_receipt(runtime.state_dir, runtime.instance_id) self.assertEqual(receipt['exit_code'], 2) self.assertFalse(receipt['drained']) self.assertEqual(hard_exit_receipts, [(2, receipt)]) self.assertTrue(os.path.exists(os.path.join(runtime.state_dir, 'slot-0.json'))) self.assertTrue(controller_threads[0].is_alive()) competing = PrivateFileLock(worker_supervisor.worker_lock_path(runtime.state_dir)) with self.assertRaises(BlockingIOError): competing.acquire() controller_release.set() controller_threads[0].join(timeout=5) runtime._shutdown_control() runtime._remove_instance() runtime._lock.release() def test_second_interrupt_routes_through_durable_forced_hard_exit(self): hard_exit_receipts = [] def client(_args, **kwargs): kwargs['started_callback']() kwargs['event_callback']({ 'slot_id': 0, 'phase': 'idle', 'progress': {}, }) self.assertEqual(runtime.request_interrupt(2), 'draining') self.assertEqual(runtime.request_interrupt(2), 'forced') self.assertTrue(runtime._deadline_escalated) self.assertTrue(kwargs['stop_event'].is_set()) return 2 def hard_exit(code): hard_exit_receipts.append(( code, load_shutdown_receipt(runtime.state_dir, runtime.instance_id), )) with tempfile.TemporaryDirectory() as root, mock.patch.object( worker_supervisor, '_package_identity', return_value=(PACKAGE, PROTOCOL), ), mock.patch.object(worker_supervisor, 'run_client', side_effect=client): runtime = WorkerSupervisor( fixture_args(root), hard_exit_hook=hard_exit, ) self.assertEqual(runtime.run(), 2) self.assertEqual(hard_exit_receipts[0][0], 2) self.assertFalse(hard_exit_receipts[0][1]['drained']) self.assertEqual(hard_exit_receipts[0][1]['exit_code'], 2) self.assertTrue(runtime._lock.acquired) runtime._shutdown_control() runtime._remove_instance() runtime._lock.release() def test_control_server_thread_start_failure_closes_socket_and_record(self): with tempfile.TemporaryDirectory() as root, mock.patch.object( worker_supervisor, '_package_identity', return_value=(PACKAGE, PROTOCOL), ): runtime = WorkerSupervisor(fixture_args(root)) with mock.patch.object(threading.Thread, 'start', side_effect=RuntimeError('start failed')): with self.assertRaisesRegex(RuntimeError, 'start failed'): runtime._publish_instance() self.assertIsNone(runtime.server) self.assertIsNone(runtime.record) self.assertFalse(os.path.exists(instance_path(runtime.state_dir))) runtime._shutdown_control() def test_retention_is_throttled_and_deadline_tick_escalates(self): clock = [0.0] with tempfile.TemporaryDirectory() as root, mock.patch.object( worker_supervisor, '_package_identity', return_value=(PACKAGE, PROTOCOL), ): runtime = WorkerSupervisor( fixture_args(root), monotonic=lambda: clock[0], wall_time=lambda: 1000.0, retention_interval_seconds=10, ) runtime._initialize_local_state() runtime._next_retention_cleanup = 10.0 with mock.patch.object(runtime.local, 'cleanup_retention') as cleanup: runtime._maintenance_tick() clock[0] = 10.0 runtime._maintenance_tick() clock[0] = 11.0 runtime._maintenance_tick() cleanup.assert_called_once() runtime._drain_deadline_monotonic = 11.0 runtime._maintenance_tick() self.assertTrue(runtime.stop_event.is_set()) self.assertTrue(runtime._deadline_escalated) def test_malformed_handshake_is_structured_nonremovable_stale(self): with tempfile.TemporaryDirectory() as root: state_dir = ensure_private_directory(os.path.join(root, 'state'), reject_reparse=True) ensure_private_directory(os.path.join(state_dir, 'control'), reject_reparse=True) record = build_instance_record( instance_id='fixture', token='t' * 48, identity=current_process_identity(), package=PACKAGE, protocol=PROTOCOL, control_port=12345, foreground=False, ) write_private_json_exclusive(instance_path(state_dir), record) classification = classify_instance( state_dir, identity_state=lambda *_args: 'alive', request=lambda *_args: {'instance_id': 'fixture'}, ) self.assertEqual(classification['state'], 'stale') self.assertEqual(classification['reason'], 'control_handshake_invalid') self.assertFalse(classification['removable']) self.assertTrue(os.path.exists(instance_path(state_dir))) def test_live_exact_process_with_control_failure_is_nonremovable_stale(self): with tempfile.TemporaryDirectory() as root: state_dir = ensure_private_directory(os.path.join(root, 'state'), reject_reparse=True) ensure_private_directory(os.path.join(state_dir, 'control'), reject_reparse=True) record = build_instance_record( instance_id='fixture', token='t' * 48, identity=current_process_identity(), package=PACKAGE, protocol=PROTOCOL, control_port=12345, foreground=False, ) write_private_json_exclusive(instance_path(state_dir), record) classification = classify_instance( state_dir, identity_state=lambda *_args: 'alive', request=mock.Mock(side_effect=OSError('control refused')), remove_stale=True, ) self.assertEqual(classification['state'], 'stale') self.assertEqual(classification['reason'], 'control_handshake_failed') self.assertFalse(classification['removable']) self.assertFalse(classification['removed']) self.assertTrue(os.path.exists(instance_path(state_dir))) def test_terminal_history_reconstructs_timeline_across_instances(self): with tempfile.TemporaryDirectory() as root, mock.patch.object( worker_supervisor, '_package_identity', return_value=(PACKAGE, PROTOCOL), ): runtime = WorkerSupervisor(fixture_args(root)) local = runtime._initialize_local_state() local.emit_phase('old-instance', 0, 'claiming', reservation_id=7, timestamp=UTC) local.emit_phase( 'old-instance', 0, 'assigned', reservation_id=7, timestamp='2026-09-23T12:00:01Z', ) local.emit_phase( runtime.instance_id, 0, 'claiming', reservation_id=7, timestamp='2026-09-23T12:00:02Z', progress={'recovered': True}, ) local.emit_phase( runtime.instance_id, 0, 'assigned', reservation_id=7, timestamp='2026-09-23T12:00:03Z', progress={'recovered': True}, ) runtime._terminal({ 'history_id': 'history-7', 'slot_id': 0, 'reservation_id': 7, 'source': 'gitlab', 'outcome': 'discarded_stale', 'receipt': {}, 'started_at': UTC, 'completed_at': '2026-09-23T12:00:04Z', 'duration_seconds': 4.0, 'first_sequence': None, 'diagnostics': [], }) history = local.history()[0] self.assertEqual( [event['instance_id'] for event in history['timeline']], ['old-instance', 'old-instance', runtime.instance_id, runtime.instance_id], ) self.assertTrue(history['timeline'][-1]['progress']['recovered']) def test_startup_wait_reports_early_process_exit(self): process = mock.Mock() process.poll.return_value = 9 with tempfile.TemporaryDirectory() as root: with self.assertRaisesRegex(WorkerSupervisorError, 'exited before'): worker_supervisor.wait_for_startup( os.path.join(root, 'missing.json'), 'nonce', process, timeout=0.2, ) def test_detached_platform_actions_are_explicit_and_injectable(self): popen = mock.Mock(return_value=object()) worker_supervisor.spawn_detached(['worker'], popen=popen, platform_name='nt') windows = popen.call_args.kwargs self.assertIn('creationflags', windows) self.assertNotIn('start_new_session', windows) popen.reset_mock() worker_supervisor.spawn_detached(['worker'], popen=popen, platform_name='posix') linux = popen.call_args.kwargs self.assertTrue(linux['start_new_session']) self.assertNotIn('creationflags', linux) def test_failed_start_termination_requires_exact_spawned_identity(self): process = mock.Mock() process.pid = 41 process.poll.side_effect = [None, 0] identity = { 'pid': 41, 'creation_time': 'created-1', 'executable': 'python', } self.assertTrue(worker_supervisor.terminate_spawned_process( process, identity, identity_state=lambda *_args: 'alive', timeout=1, )) process.terminate.assert_called_once() process.wait.assert_called_once_with(timeout=1.0) with self.assertRaisesRegex(WorkerSupervisorError, 'exact identity'): worker_supervisor.terminate_spawned_process( mock.Mock(pid=42, poll=mock.Mock(return_value=None)), identity, identity_state=lambda *_args: 'reused', ) if __name__ == '__main__': unittest.main()