confluence_web/agent/supervisor

#!/usr/bin/env python3
"""Immutable Linux PID-1 supervisor; host mode is development-only."""
import ctypes
import json
import os
import select
import signal
import socket
import subprocess
import sys
import time
from pathlib import Path

libc = ctypes.CDLL(None, use_errno=True)
# Adopt double-forked descendants on hosts, as well as under container PID 1.
if libc.prctl(36, 1, 0, 0, 0) != 0:  # PR_SET_CHILD_SUBREAPER
    sys.exit(1)
if libc.prctl(4, 0, 0, 0, 0) != 0:  # PR_SET_DUMPABLE: deny same-UID proc memory/fd access
    sys.exit(1)
container = os.getpid() == 1
if not container and '--dev' not in sys.argv:
    sys.stderr.write('Supervisor requires container PID 1; use --dev only for host tests.\n')
    sys.exit(1)
# PID 1 ignores namespace-local SIGKILL/SIGSTOP. Other signals terminate the
# run or are ignored; none can pause or extend the deadline. No exec after prctl.
for sig in (signal.SIGTERM, signal.SIGINT, signal.SIGHUP, signal.SIGUSR1, signal.SIGUSR2):
    signal.signal(sig, signal.SIG_IGN)


def max_run_seconds():
    """Protocol maximum for one query, set by the backend at container start.
    Matches LIMITS.MAX_DEADLINE_MS in the bridge; invalid or absent values fall back to 900."""
    raw = os.environ.get('CONFLUENCE_WEB_MAX_DEADLINE_SECONDS', '')
    try:
        value = int(raw)
    except ValueError:
        return 900
    return value if 60 <= value <= 3600 else 900


MAX_RUN_SECONDS = max_run_seconds()
START_GRACE_SECONDS = 180  # a container that never receives a start frame ends here
started = time.monotonic()
deadline = started + MAX_RUN_SECONDS
parent, child = socket.socketpair()
child.set_inheritable(True)
root = Path(__file__).resolve().parent
entry = root / 'dist' / 'bridge.js'
args = ['node', str(entry)]
env = {'PATH': '/usr/local/bin:/usr/bin:/bin', 'HOME': '/home/agent', 'LANG': 'C.UTF-8',
       'AGENT_SUPERVISOR_FD': str(child.fileno()),
       'CONFLUENCE_WEB_MAX_DEADLINE_SECONDS': str(MAX_RUN_SECONDS)}
# Per-query call budgets, forwarded verbatim; the bridge validates them and falls back to
# the contract defaults (100 Confluence / 50 model) when absent or invalid.
for name in ('CONFLUENCE_WEB_MAX_CONFLUENCE_CALLS', 'CONFLUENCE_WEB_MAX_MODEL_CALLS'):
    if name in os.environ:
        env[name] = os.environ[name]
if '--dev' in sys.argv:
    env['PATH'] = os.environ.get('PATH', env['PATH'])
    env['HOME'] = os.environ.get('HOME', '/tmp')
    env['AGENT_WORK_DIR'] = os.environ.get('AGENT_WORK_DIR', '/work')
bridge = subprocess.Popen(args, env=env, pass_fds=(child.fileno(),))
child.close()
parent.setblocking(False)
protected = {os.getpid(), bridge.pid}
buffer = b''


def descendants():
    """Never signal unrelated host processes. Adopted children retain this ancestry."""
    parents = {}
    for item in Path('/proc').iterdir():
        if item.name.isdecimal():
            try:
                fields = (item / 'stat').read_text().rsplit(')', 1)[1].split()
                parents[int(item.name)] = int(fields[1])
            except (OSError, ValueError, IndexError):
                pass
    owned = {os.getpid()}
    while True:
        extra = {pid for pid, ppid in parents.items() if ppid in owned} - owned
        if not extra:
            break
        owned.update(extra)
    return owned - protected


def reap():
    # Re-scan until every descendant is gone, including children forked during
    # termination. Reap only adopted children, leaving Popen to reap the bridge.
    while time.monotonic() < deadline:
        victims = descendants()
        if not victims:
            return
        for pid in victims:
            try:
                os.kill(pid, signal.SIGKILL)
            except ProcessLookupError:
                pass
            try:
                os.waitpid(pid, os.WNOHANG)
            except ChildProcessError:
                pass
        time.sleep(0.01)
    raise TimeoutError()


code = 1
try:
    deadline_set = False
    collecting = False
    while time.monotonic() < deadline:
        if not deadline_set and time.monotonic() >= started + START_GRACE_SECONDS:
            deadline = time.monotonic()  # missing start frame: treat as expired
            break
        if bridge.poll() is not None:
            code = bridge.returncode
            break
        readable, _, _ = select.select([parent], [], [], min(0.05, max(0, deadline - time.monotonic())))
        if not readable:
            continue
        data = parent.recv(4096)
        if not data:
            try:
                code = bridge.wait(timeout=min(1, max(0.01, deadline - time.monotonic())))
            except subprocess.TimeoutExpired:
                pass
            break
        buffer += data
        if len(buffer) > 8192:
            break
        while b'\n' in buffer:
            line, buffer = buffer.split(b'\n', 1)
            msg = json.loads(line)
            if not isinstance(msg, dict):
                raise ValueError()
            if msg.get('type') == 'set_deadline' and not deadline_set:
                ms = msg.get('remaining_ms')
                if type(ms) is not int or not 1 <= ms <= MAX_RUN_SECONDS * 1000:
                    raise ValueError()
                deadline = min(deadline, time.monotonic() + ms / 1000)
                deadline_set = True
            elif msg == {'type': 'reap_descendants'} and not collecting:
                reap()
                collecting = True
                parent.sendall(b'{"type":"reaped"}\n')
            else:
                raise ValueError()
except (OSError, ValueError, TimeoutError):
    pass
finally:
    if container and time.monotonic() >= deadline:
        sys.exit(1)  # Kernel destroys the namespace; do not add cleanup time.
    # Killing PID 1 ends the container namespace. Host development cleanup is
    # still restricted to descendants of this supervisor.
    protected.discard(bridge.pid)
    deadline = max(deadline, time.monotonic() + 1)
    try:
        reap()
    except (OSError, TimeoutError):
        pass
    try:
        bridge.wait(timeout=1)
    except subprocess.TimeoutExpired:
        bridge.kill()
    parent.close()
sys.exit(code if code >= 0 else 1)