Artemis-Eapil/artemis/runtime/task_control.py

"""Cooperative, session-scoped task guidance and device lease parking."""

import asyncio
import json
import logging
import os
from pathlib import Path
import threading
import time
import uuid

from artemis.config import TRACES_PATH
from artemis.runtime.device_lock import DeviceExecutionLock

logger = logging.getLogger(__name__)


def control_dir(session_id: str) -> Path:
    if not session_id or not all(c.isalnum() or c in "-_" for c in session_id):
        raise ValueError("无效的任务编号")
    return TRACES_PATH / session_id


def _write_atomic(path: Path, data: dict) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    temporary = path.with_name(f".{path.name}.{uuid.uuid4().hex}.tmp")
    temporary.write_text(json.dumps(data, ensure_ascii=False), encoding="utf-8")
    os.replace(temporary, path)


def read_state(session_id: str) -> dict:
    try:
        return json.loads(
            (control_dir(session_id) / "task_control_state.json").read_text(encoding="utf-8")
        )
    except (OSError, ValueError):
        return {}


def request_pause(session_id: str) -> None:
    (control_dir(session_id) / "resume.request").unlink(missing_ok=True)
    _write_atomic(control_dir(session_id) / "pause.request", {"created_at": time.time()})


def request_resume(session_id: str) -> None:
    _write_atomic(control_dir(session_id) / "resume.request", {"created_at": time.time()})


def add_instruction(
    session_id: str, instruction: str, *, release_loop: bool = False, source: str = "user"
) -> None:
    if not instruction.strip() or len(instruction) > 10000:
        raise ValueError("指导内容必须在 1–10000 字之间")
    directory = control_dir(session_id) / "instructions"
    directory.mkdir(parents=True, exist_ok=True)
    # Unique names make simultaneous web and MCP submissions additive, never last-write-wins.
    path = directory / f"{time.time_ns():020d}-{uuid.uuid4().hex}.json"
    _write_atomic(
        path,
        {
            "instruction": instruction.strip(),
            "release_loop": bool(release_loop),
            "source": source,
            "timestamp": time.time(),
        },
    )


def consume_instructions(base_dir: str | Path) -> dict | None:
    directory = Path(base_dir) / "instructions"
    found = []
    for path in sorted(directory.glob("*.json")) if directory.exists() else []:
        try:
            claimed = path.with_suffix(".reading")
            path.rename(claimed)
            found.append(json.loads(claimed.read_text(encoding="utf-8")))
            claimed.unlink()
        except (OSError, ValueError):
            continue
    legacy = Path(base_dir) / "injected_instruction.json"
    if legacy.exists():
        try:
            claimed = legacy.with_suffix(".reading")
            legacy.rename(claimed)
            found.append(json.loads(claimed.read_text(encoding="utf-8")))
            claimed.unlink()
        except (OSError, ValueError):
            pass
    if not found:
        return None
    return {
        "instruction": "\n\n".join(str(item.get("instruction") or "") for item in found).strip(),
        "release_loop": any(bool(item.get("release_loop")) for item in found),
    }


class TaskControl:
    """Yield a device lease at the next safe turn and reacquire at the FIFO tail."""

    def __init__(self, session_id: str, device_lock: DeviceExecutionLock):
        self.session_id = session_id
        self.device_lock = device_lock
        self.directory = control_dir(session_id)

    def _state(self, status: str) -> None:
        _write_atomic(
            self.directory / "task_control_state.json",
            {"status": status, "updated_at": time.time(), "pid": os.getpid()},
        )

    async def checkpoint(self, *, on_pause=None, on_resume=None) -> bool:
        pause = self.directory / "pause.request"
        if not pause.exists():
            return False
        pause.unlink(missing_ok=True)
        if on_pause is not None:
            await on_pause()
        self.device_lock.release()
        self._state("task_paused")
        resume = self.directory / "resume.request"
        while not resume.exists():
            await asyncio.sleep(0.2)
        resume.unlink(missing_ok=True)
        self._state("resuming")
        previous = self.device_lock
        self.device_lock = DeviceExecutionLock(
            previous.device_id,
            description=previous.description,
            max_concurrency=previous.max_concurrency,
            concurrency_mode=previous.concurrency_mode,
            session_id=self.session_id,
            ingress=previous.ingress,
            lock_scope=previous.lock_scope,
        )
        cancelled = threading.Event()
        acquiring = asyncio.create_task(
            asyncio.to_thread(self.device_lock.acquire, cancel_event=cancelled)
        )
        try:
            await asyncio.shield(acquiring)
        except asyncio.CancelledError:
            cancelled.set()
            try:
                await asyncio.shield(acquiring)
            except Exception:
                logger.warning("Resume acquisition ended during cancellation", exc_info=True)
            self.device_lock.release()
            raise
        if on_resume is not None:
            await on_resume()
        self._state("running")
        return True

    def release(self) -> None:
        self.device_lock.release()
        (self.directory / "task_control_state.json").unlink(missing_ok=True)