fn_registry/python/functions/infra/mark_claude_role.py

"""Marca el role de una sesion de Claude Code resolviendo PID -> sessionId.

El meta-orquestador de flota arranca un Claude Code y necesita marcarlo con su
role (orchestrator | executor) en su `goal.json` para que la TUI fleetview lo
clasifique (p. ej. pinear el orquestador arriba). El problema es de timing: el
`sessionId` no se conoce al arrancar el proceso; Claude Code escribe el archivo
`~/.claude/sessions/<PID>.json` (que contiene el `sessionId`) unos segundos
despues. Esta funcion espera ese archivo sondeando, resuelve el sessionId y
escribe SOLO la clave `role` en `~/.claude/goals/<sessionId>.json` preservando
todos los demas campos del goal.
"""

import json
import os
import time

_VALID_ROLES = ("orchestrator", "executor")


def mark_claude_role(
    pid: int,
    role: str,
    wait_s: float = 10.0,
    sessions_dir: str | None = None,
    goals_dir: str | None = None,
) -> dict:
    """Resuelve PID -> sessionId esperando el sessions/<pid>.json y marca el role.

    Sondea `<sessions_dir>/<pid>.json` cada ~0.25s hasta `wait_s` segundos (medido
    con time.monotonic) esperando que exista y contenga un `sessionId` no vacio.
    Si lo encuentra, abre `<goals_dir>/<sessionId>.json` (dict vacio si no existe),
    setea SOLO la clave `role` preservando el resto (goal, phase, dod, etc.), y lo
    escribe atomicamente (tmp + os.replace). Si expira, devuelve un dict de error
    sin lanzar excepcion (el launcher decide que hacer).

    Args:
        pid: PID del proceso Claude Code recien arrancado.
        role: "orchestrator" o "executor". Cualquier otro valor lanza ValueError
            sin tocar disco.
        wait_s: segundos maximos a esperar el sessions/<pid>.json. Default 10.0.
        sessions_dir: directorio de los sessions JSON. Default ~/.claude/sessions.
        goals_dir: directorio de los goal JSON. Default ~/.claude/goals.

    Returns:
        Si tuvo exito: dict con ok=True, pid, session_id, role y path (ruta del
        goal escrito). Si expiro el poll: dict con ok=False, error (mensaje) y pid.

    Raises:
        ValueError: si `role` no esta en ("orchestrator", "executor").
    """
    if role not in _VALID_ROLES:
        raise ValueError(
            f"role invalido: {role!r}. Debe ser uno de {_VALID_ROLES}"
        )

    home = os.path.expanduser("~")
    if sessions_dir is None:
        sessions_dir = os.path.join(home, ".claude", "sessions")
    if goals_dir is None:
        goals_dir = os.path.join(home, ".claude", "goals")

    session_path = os.path.join(sessions_dir, f"{pid}.json")

    # POLL: espera hasta wait_s segundos a que exista sessions/<pid>.json con un
    # sessionId no vacio. Bucle corto con time.monotonic como deadline (no un
    # unico sleep largo) para reaccionar en cuanto aparezca el archivo.
    deadline = time.monotonic() + wait_s
    sid = ""
    while True:
        sid = _read_session_id(session_path)
        if sid:
            break
        if time.monotonic() >= deadline:
            return {
                "ok": False,
                "error": f"timeout esperando sessions/{pid}.json",
                "pid": pid,
            }
        time.sleep(0.25)

    goal_path = os.path.join(goals_dir, f"{sid}.json")

    # Lee el goal existente (dict vacio si no existe o es ilegible/no-dict).
    goal: dict = {}
    if os.path.exists(goal_path):
        try:
            with open(goal_path, "r", encoding="utf-8") as fh:
                loaded = json.load(fh)
            if isinstance(loaded, dict):
                goal = loaded
        except (json.JSONDecodeError, ValueError, OSError):
            goal = {}

    # Setea SOLO la clave role, preservando todo lo demas.
    goal["role"] = role

    # Escritura atomica: tmp + os.replace en el mismo directorio.
    os.makedirs(goals_dir, exist_ok=True)
    tmp_path = goal_path + f".tmp.{os.getpid()}"
    with open(tmp_path, "w", encoding="utf-8") as fh:
        json.dump(goal, fh, ensure_ascii=False, indent=2)
        fh.flush()
        os.fsync(fh.fileno())
    os.replace(tmp_path, goal_path)

    return {
        "ok": True,
        "pid": pid,
        "session_id": sid,
        "role": role,
        "path": goal_path,
    }


def _read_session_id(session_path: str) -> str:
    """Lee el sessionId del sessions/<pid>.json. Devuelve "" si falta/ilegible/vacio."""
    if not os.path.exists(session_path):
        return ""
    try:
        with open(session_path, "r", encoding="utf-8") as fh:
            data = json.load(fh)
    except (json.JSONDecodeError, ValueError, OSError):
        return ""
    if not isinstance(data, dict):
        return ""
    sid = data.get("sessionId")
    if isinstance(sid, str) and sid:
        return sid
    return ""