203 lines
7.4 KiB
Python
203 lines
7.4 KiB
Python
"""Cooperative, session-scoped task guidance and device lease parking."""
|
||
|
||
import asyncio
|
||
from collections.abc import Awaitable
|
||
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"
|
||
) -> dict:
|
||
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"
|
||
record = {
|
||
"id": path.stem,
|
||
"instruction": instruction.strip(),
|
||
"release_loop": bool(release_loop),
|
||
"source": source,
|
||
"timestamp": time.time(),
|
||
}
|
||
_write_atomic(path, record)
|
||
return record
|
||
|
||
|
||
def list_instructions(session_id: str) -> list[dict]:
|
||
"""Show submissions separately from guidance the runner actually consumed."""
|
||
directory = control_dir(session_id) / "instructions"
|
||
records = []
|
||
for folder, status in ((directory, "pending"), (directory / "consumed", "consumed")):
|
||
for path in folder.glob("*.json") if folder.exists() else []:
|
||
try:
|
||
record = json.loads(path.read_text(encoding="utf-8"))
|
||
if isinstance(record, dict) and record.get("instruction"):
|
||
records.append(
|
||
{**record, "id": record.get("id") or path.stem, "status": status}
|
||
)
|
||
except (OSError, ValueError):
|
||
continue
|
||
return sorted(records, key=lambda record: (record.get("timestamp", 0), record["id"]))
|
||
|
||
|
||
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)
|
||
record = json.loads(claimed.read_text(encoding="utf-8"))
|
||
record["consumed_at"] = time.time()
|
||
_write_atomic(directory / "consumed" / path.name, record)
|
||
claimed.unlink()
|
||
found.append(record)
|
||
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
|
||
|
||
async def run_interruptibly[T](
|
||
self, operation: Awaitable[T], *, on_pause=None, on_resume=None
|
||
) -> tuple[bool, T | None]:
|
||
"""Park a task during a long model/helper call, then discard its stale result.
|
||
|
||
Only use for operations whose cancellation cannot leave a device action
|
||
half executed. The caller must return to fresh perception after resume.
|
||
"""
|
||
running = asyncio.ensure_future(operation)
|
||
pause = self.directory / "pause.request"
|
||
try:
|
||
while True:
|
||
if pause.exists():
|
||
running.cancel()
|
||
await asyncio.gather(running, return_exceptions=True)
|
||
await self.checkpoint(on_pause=on_pause, on_resume=on_resume)
|
||
return True, None
|
||
if running.done():
|
||
return False, running.result()
|
||
await asyncio.wait({running}, timeout=0.2)
|
||
finally:
|
||
if not running.done():
|
||
running.cancel()
|
||
await asyncio.gather(running, return_exceptions=True)
|
||
|
||
def release(self) -> None:
|
||
self.device_lock.release()
|
||
(self.directory / "task_control_state.json").unlink(missing_ok=True)
|