AI Flow

#!/usr/bin/env python3
"""Opt-in Codex continuation supervisor. Never edits engineering acceptance."""
import argparse
import contextlib
import fcntl
import hashlib
import json
import os
from pathlib import Path
import re
import subprocess
import time
import uuid


def read(path):
    return json.loads(Path(path).read_text())


def write(path, value):
    path = Path(path)
    tmp = path.with_suffix('.tmp')
    tmp.write_text(json.dumps(value, indent=2) + '\n')
    tmp.replace(path)


def decision(status, turn, queue, guard, previous):
    if status.get('type') != 'idle':
        return 'waiting-for-idle'
    if queue:
        return 'user-queue-pending'
    if turn.get('status') != 'completed':
        return 'paused-interrupted-or-failed-turn'
    if not turn.get('id'):
        return 'error-missing-turn-id'
    if previous.get('attemptedTurn') == turn['id']:
        return 'waiting-for-queued-turn'
    texts = [c.get('text', '') for i in turn.get('items', [])
             if i.get('type') == 'userMessage' for c in i.get('content', [])]
    # Conservative extra defense, not a substitute for the explicit pause control.
    if any(re.search(r'(?im)^\s*(?:(?:please|can (?:you|u))\s+)?(?:stop\b|pause\b|cancel\b|do not continue\b|don.t continue\b)', t) for t in texts):
        return 'paused-user-instruction'
    if guard.get('decision') == 'complete' and guard.get('mayStop') is True:
        return 'complete'
    if guard.get('decision') != 'continue' or guard.get('mayStop') is not False:
        return 'error-invalid-guard'
    return 'continue'


class Codex:
    def __init__(self, socket):
        from websockets.sync.client import unix_connect
        self.ws = unix_connect(str(socket), compression=None, max_size=None,
                               open_timeout=5, close_timeout=1)
        self.serial = 0
        self.rpc('initialize', {'clientInfo': {'name': 'adom-aiflow-supervisor', 'version': '1'},
                                'capabilities': {'experimentalApi': True}})
        self.ws.send(json.dumps({'method': 'initialized'}))

    def rpc(self, method, params):
        self.serial += 1
        self.ws.send(json.dumps({'id': self.serial, 'method': method, 'params': params}))
        deadline = time.monotonic() + 15
        while True:
            reply = json.loads(self.ws.recv(timeout=max(.01, deadline-time.monotonic())))
            if reply.get('id') == self.serial:
                if 'error' in reply:
                    raise RuntimeError(str(reply['error']))
                return reply['result']

    def close(self):
        self.ws.close()


def discover(thread):
    candidates = []
    root = Path.home()/'.local/state/adom/codex-runtime'
    for meta in root.glob('*/server.json'):
        try:
            data = read(meta)
            with contextlib.closing(Codex(data['socket'])) as client:
                t = client.rpc('thread/read', {'threadId': thread, 'includeTurns': False})['thread']
                if t['id'] == thread and t['status']['type'] != 'notLoaded':
                    candidates.append((data, t))
        except Exception:
            continue
    if len(candidates) != 1:
        raise RuntimeError(f'Expected exactly one loaded runtime for thread; found {len(candidates)}')
    return candidates[0]


def guard(run):
    result = subprocess.run(['adom-aiflow', '--run', str(run), 'finish-guard', '--json'],
                            capture_output=True, text=True, timeout=60,
                            env={**os.environ, 'AIFLOW_DASHBOARD': 'off'}, cwd=run.parent.parent)
    data = json.loads(result.stdout)
    if (result.returncode, data.get('decision')) not in [(0, 'complete'), (2, 'continue')]:
        raise RuntimeError('Finish guard failed or returned an unsupported decision')
    return data


def progress_hash(run, report):
    # Exclude activity timestamps and supervisor receipts: activity is not progress.
    data = read(run/'run.json')
    payload = {k: data.get(k) for k in ['currentStep', 'outcomes', 'boards', 'delivery', 'continuation']}
    payload['remaining'] = report.get('remaining')
    # Component review is real engineering progress before any board change.
    # Keep activity-only stamps out of the comparison and fail closed on bad JSON.
    register = run/'components.json'
    if register.exists():
        def evidence(value):
            if isinstance(value, dict):
                return {k: evidence(v) for k, v in value.items()
                        if k not in {'checkedAt', 'updatedAt', 'createdAt', 'timestamp', 'lastActivityAt'}}
            if isinstance(value, list):
                return [evidence(v) for v in value]
            return value
        payload['componentReview'] = evidence(read(register))
    return hashlib.sha256(json.dumps(payload, sort_keys=True).encode()).hexdigest()


def check(run, directory, thread, state):
    if (directory/'pause.json').exists():
        return 'paused', state
    meta, _ = discover(thread)
    with contextlib.closing(Codex(meta['socket'])) as client:
        status = client.rpc('thread/read', {'threadId': thread, 'includeTurns': False})['thread']['status']
        if status['type'] in ['notLoaded', 'systemError']:
            return 'error-runtime-unavailable', state
        if status['type'] != 'idle':
            return 'waiting-for-idle', state
        turn = client.rpc('thread/turns/list', {'threadId': thread, 'limit': 1,
                         'sortDirection': 'desc', 'itemsView': 'full'})['data'][0]
        queue = client.rpc('thread/queue/list', {'threadId': thread})['data']
        report = guard(run)
        action = decision(status, turn, queue, report, state)
        if action != 'continue':
            return action, state
        fingerprint = progress_hash(run, report)
        unchanged = state.get('unchanged', 0) + 1 if state.get('progressHash') == fingerprint else 0
        if unchanged >= 3:
            return 'paused-no-recorded-progress', state
        # Recheck immediately before dispatch; never interrupt an active turn.
        if (directory/'pause.json').exists():
            return 'paused', state
        fresh = client.rpc('thread/read', {'threadId': thread, 'includeTurns': False})['thread']['status']
        if fresh['type'] != 'idle':
            return 'waiting-for-idle', state
        if client.rpc('thread/queue/list', {'threadId': thread})['data']:
            return 'user-queue-pending', state
        state.update(attemptedTurn=turn['id'], progressHash=fingerprint, unchanged=unchanged,
                     dispatchId=str(uuid.uuid4()), dispatch='intent-recorded')
        # Persist BEFORE sending: after a timeout never blindly enqueue again.
        write(directory/'state.json', state)
        prompt = ('AIFlow continuation authorized by the user for this run only. '
                  'First honor any newer user pause, cancellation, or change of scope. '
                  'Read the project completion plan and current run; continue the next useful authorized task. '
                  'Do not weaken gates, fabricate evidence, bypass approvals, create another supervisor, '
                  'or stop after a progress report. If genuinely blocked after exhausting useful work, '
                  'pause this supervisor and record the concrete dependency. '
                  f'Run: {run}. Supervisor: {directory}. '
                  f'Guard next action: {report.get("nextAction", "Review remaining work")}')
        queued = client.rpc('thread/queue/add', {'threadId': thread,
                    'clientUserMessageId': state['dispatchId'],
                    'input': [{'type': 'text', 'text': prompt}]})['queuedSubmission']
        state.update(queuedId=queued['id'], dispatch='queued')
        write(directory/'state.json', state)
        # A newer human turn may have started meanwhile; remove only our item.
        latest = client.rpc('thread/read', {'threadId': thread, 'includeTurns': False})['thread']['status']
        # queue/add may auto-start immediately in the IDE. Attribute activity
        # by the exact client ID, not by timing or an assumption of a race.
        if latest['type'] == 'active' and not (directory/'pause.json').exists():
            recent = client.rpc('thread/turns/list', {'threadId': thread, 'limit': 1,
                                'sortDirection': 'desc', 'itemsView': 'summary'})['data']
            if recent and any(item.get('type') == 'userMessage'
                              and item.get('clientId') == state['dispatchId']
                              for item in recent[0].get('items', [])):
                state.update(dispatch='auto-start-verified', startedTurn=recent[0]['id'])
                return 'continuation-dispatched', state
        if latest['type'] != 'idle' or (directory/'pause.json').exists():
            client.rpc('thread/queue/delete', {'threadId': thread, 'queuedSubmissionId': queued['id']})
            state['dispatch'] = 'removed-after-state-change'
            return 'paused-dispatch-race', state
        if latest['type'] == 'idle':
            client.rpc('thread/queue/start', {'threadId': thread, 'queuedSubmissionId': queued['id']})
            state['dispatch'] = 'start-acknowledged'
        return 'continuation-dispatched', state


def main():
    p = argparse.ArgumentParser(description=__doc__)
    p.add_argument('command', choices=['watch', 'inspect', 'status', 'pause', 'resume'])
    p.add_argument('--run', type=Path, required=True)
    p.add_argument('--thread', default=os.environ.get('CODEX_THREAD_ID'))
    p.add_argument('--interval', type=int, default=60)
    p.add_argument('--reason', default='Explicit supervisor control')
    a = p.parse_args()
    a.run = a.run.resolve()
    if not (a.run/'run.json').is_file():
        p.error('run.json must exist')
    directory = a.run/'supervisor'
    directory.mkdir(exist_ok=True)
    if a.command == 'status':
        print(json.dumps(read(directory/'status.json'), indent=2)); return
    if a.command == 'pause':
        write(directory/'pause.json', {'reason': a.reason, 'at': time.time()}); return
    if a.command == 'resume':
        # Restart explicitly with watch after inspecting the reason for stopping.
        (directory/'pause.json').unlink(missing_ok=True); return
    if not a.thread or a.interval < 60:
        p.error('Exact thread ID required; interval must be at least 60 seconds')
    meta, t = discover(a.thread)
    if a.command == 'inspect':
        print(json.dumps({'threadId': t['id'], 'name': t.get('name'), 'status': t['status'],
                          'socket': meta['socket'], 'run': str(a.run)}, indent=2)); return
    # One supervisor per thread, even if two runs attempt to bind to it.
    lockdir = Path.home()/'.local/state/adom/aiflow-supervisor'
    lockdir.mkdir(parents=True, exist_ok=True)
    lockname = hashlib.sha256(a.thread.encode()).hexdigest()
    with open(lockdir/lockname, 'w') as lock:
        fcntl.flock(lock, fcntl.LOCK_EX | fcntl.LOCK_NB)
        state = read(directory/'state.json') if (directory/'state.json').exists() else {}
        if state.get('threadId', a.thread) != a.thread:
            raise RuntimeError('Run already bound to another thread')
        state['threadId'] = a.thread
        while True:
            try:
                action, state = check(a.run, directory, a.thread, state)
                error = None
            except Exception as exc:
                action, error = 'error', str(exc)
            write(directory/'state.json', state)
            status = {'status': action, 'error': error, 'checkedAt': time.time(), 'pid': os.getpid(),
                      'threadId': a.thread, 'intervalSeconds': a.interval, 'state': state}
            write(directory/'status.json', status)
            with open(directory/'events.jsonl', 'a') as log:
                log.write(json.dumps(status)+'\n')
            if action.startswith(('paused', 'error')) or action == 'complete':
                break
            time.sleep(a.interval)


if __name__ == '__main__':
    main()