wehub-resource-sync
337 lines
17 KiB
Python
337 lines
17 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
import logging
|
|
import os
|
|
import re
|
|
from contextlib import contextmanager
|
|
from datetime import datetime
|
|
from functools import wraps
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
try:
|
|
import fcntl
|
|
except ImportError: # pragma: no cover - non-POSIX platforms
|
|
fcntl = None
|
|
|
|
|
|
STALE_KEYS = ["story", "characters", "script", "storyboard", "shot_descriptions", "camera_tree", "frames", "clips", "final_video"]
|
|
|
|
|
|
def _synchronized(method):
|
|
"""Hold the index file lock across a read-modify-write cycle.
|
|
|
|
Every mutator loads the whole sessions file, edits it, and saves it back;
|
|
without a lock, two concurrent writers (threads or processes) silently
|
|
drop each other's updates.
|
|
"""
|
|
|
|
@wraps(method)
|
|
def wrapper(self, *args, **kwargs):
|
|
with self._locked():
|
|
return method(self, *args, **kwargs)
|
|
|
|
return wrapper
|
|
|
|
|
|
class SessionIndex:
|
|
def __init__(self, workspace_root: str | Path) -> None:
|
|
self.workspace_root = Path(workspace_root).resolve()
|
|
self.vimax_dir = self.workspace_root / ".vimax"
|
|
self.sessions_path = self.vimax_dir / "sessions.json"
|
|
self.memory_path = self.vimax_dir / "memory.md"
|
|
self.logs_dir = self.vimax_dir / "logs"
|
|
self.working_root = self.workspace_root / ".working_dir"
|
|
self.vimax_dir.mkdir(parents=True, exist_ok=True)
|
|
self.logs_dir.mkdir(parents=True, exist_ok=True)
|
|
self.working_root.mkdir(parents=True, exist_ok=True)
|
|
if not self.memory_path.exists():
|
|
self.memory_path.write_text("# User Preferences\n", encoding="utf-8")
|
|
if not self.sessions_path.exists():
|
|
self.save({"active_session_id": "", "sessions": {}})
|
|
|
|
@contextmanager
|
|
def _locked(self):
|
|
if fcntl is None:
|
|
yield
|
|
return
|
|
lock_path = self.vimax_dir / "sessions.lock"
|
|
with open(lock_path, "a+", encoding="utf-8") as handle:
|
|
fcntl.flock(handle, fcntl.LOCK_EX)
|
|
try:
|
|
yield
|
|
finally:
|
|
fcntl.flock(handle, fcntl.LOCK_UN)
|
|
|
|
def load(self) -> dict[str, Any]:
|
|
try:
|
|
return json.loads(self.sessions_path.read_text(encoding="utf-8"))
|
|
except FileNotFoundError:
|
|
return {"active_session_id": "", "sessions": {}}
|
|
except json.JSONDecodeError:
|
|
# A corrupt file usually means a crash mid-write. Returning empty
|
|
# state is fine for this call, but the next save() would overwrite
|
|
# the file and destroy every session — keep the evidence first.
|
|
backup = self.sessions_path.with_name(f"sessions.json.corrupt-{datetime.now().strftime('%Y%m%d-%H%M%S-%f')}")
|
|
try:
|
|
os.replace(self.sessions_path, backup)
|
|
logging.error("sessions.json was corrupt; preserved it at %s and starting with empty state", backup)
|
|
except OSError:
|
|
logging.error("sessions.json is corrupt and could not be backed up; starting with empty state")
|
|
return {"active_session_id": "", "sessions": {}}
|
|
|
|
def save(self, data: dict[str, Any]) -> None:
|
|
tmp_path = self.sessions_path.with_name("sessions.json.tmp")
|
|
tmp_path.write_text(json.dumps(data, ensure_ascii=False, indent=2), encoding="utf-8")
|
|
os.replace(tmp_path, self.sessions_path)
|
|
|
|
def active(self) -> dict[str, Any] | None:
|
|
data = self.load()
|
|
session_id = str(data.get("active_session_id", ""))
|
|
if not session_id:
|
|
return None
|
|
record = data.get("sessions", {}).get(session_id)
|
|
return self._with_session_defaults(record) if isinstance(record, dict) else None
|
|
|
|
def get(self, session_id: str) -> dict[str, Any] | None:
|
|
normalized = self._normalize_session_id(session_id)
|
|
record = self.load().get("sessions", {}).get(normalized)
|
|
return self._with_session_defaults(record) if isinstance(record, dict) else None
|
|
|
|
@_synchronized
|
|
def create(self, idea: str = "", user_requirement: str = "", style: str = "", session_id: str | None = None) -> dict[str, Any]:
|
|
data = self.load()
|
|
sessions = data.setdefault("sessions", {})
|
|
final_id = self._normalize_session_id(session_id) if session_id else self._new_session_id(idea or user_requirement or "vimax", sessions)
|
|
if final_id in sessions:
|
|
final_id = self._dedupe_session_id(final_id, sessions)
|
|
now = datetime.now().isoformat(timespec="seconds")
|
|
working_dir = self._working_dir_for_id(final_id)
|
|
(working_dir / "idea2video").mkdir(parents=True, exist_ok=True)
|
|
(working_dir / "script2video").mkdir(parents=True, exist_ok=True)
|
|
record = {
|
|
"session_id": final_id,
|
|
"working_dir": str(working_dir.relative_to(self.workspace_root)),
|
|
"idea": idea,
|
|
"user_requirement": user_requirement,
|
|
"style": style,
|
|
"stage": "created",
|
|
"summary": "",
|
|
"stale": {key: False for key in STALE_KEYS},
|
|
"recent_turn_records": [],
|
|
"compacted_summary": "",
|
|
"compacted_turns": 0,
|
|
"compaction_snapshots": [],
|
|
"last_compaction_reason": "",
|
|
"last_compaction_at": "",
|
|
"created_at": now,
|
|
"updated_at": now,
|
|
}
|
|
sessions[final_id] = record
|
|
data["active_session_id"] = final_id
|
|
self.save(data)
|
|
return record
|
|
|
|
def get_or_create_active(self, idea: str = "", user_requirement: str = "", style: str = "") -> dict[str, Any]:
|
|
active = self.active()
|
|
if active is not None:
|
|
return active
|
|
return self.create(idea=idea, user_requirement=user_requirement, style=style)
|
|
|
|
@_synchronized
|
|
def set_active(self, session_id: str) -> dict[str, Any]:
|
|
normalized = self._normalize_session_id(session_id)
|
|
data = self.load()
|
|
if normalized not in data.get("sessions", {}):
|
|
raise KeyError(f"Unknown session_id: {session_id}")
|
|
data["active_session_id"] = normalized
|
|
self.save(data)
|
|
return dict(data["sessions"][normalized])
|
|
|
|
@_synchronized
|
|
def update_stage(self, session_id: str, stage: str, summary: str = "") -> None:
|
|
data = self.load()
|
|
record = data.get("sessions", {}).get(session_id)
|
|
if not isinstance(record, dict):
|
|
raise KeyError(f"Unknown session_id: {session_id}")
|
|
record["stage"] = stage
|
|
if summary:
|
|
record["summary"] = summary
|
|
record["updated_at"] = datetime.now().isoformat(timespec="seconds")
|
|
self.save(data)
|
|
|
|
@_synchronized
|
|
def mark_stale(self, session_id: str, keys: list[str]) -> None:
|
|
data = self.load()
|
|
record = data.get("sessions", {}).get(session_id)
|
|
if not isinstance(record, dict):
|
|
raise KeyError(f"Unknown session_id: {session_id}")
|
|
stale = record.setdefault("stale", {key: False for key in STALE_KEYS})
|
|
for key in keys:
|
|
stale[key] = True
|
|
record["updated_at"] = datetime.now().isoformat(timespec="seconds")
|
|
self.save(data)
|
|
|
|
@_synchronized
|
|
def update_compaction(self, session_id: str, result: dict[str, Any]) -> None:
|
|
data = self.load()
|
|
session = data.get("sessions", {}).get(session_id)
|
|
if not isinstance(session, dict):
|
|
raise KeyError(f"Unknown session_id: {session_id}")
|
|
summary = str(result.get("summary", "") or "")
|
|
compacted_count = int(result.get("compacted_message_count", 0) or 0)
|
|
snapshot = {
|
|
"level": len(session.get("compaction_snapshots", []) or []) + 1,
|
|
"reason": str(result.get("reason", "manual") or "manual"),
|
|
"mode": str(result.get("mode", "unknown") or "unknown"),
|
|
"summary": summary,
|
|
"preserved_messages": int(result.get("preserved_message_count", 0) or 0),
|
|
"compacted_message_count": compacted_count,
|
|
"estimated_tokens_before": int(result.get("estimated_tokens_before", 0) or 0),
|
|
"estimated_tokens_after": int(result.get("estimated_tokens_after", 0) or 0),
|
|
"created_at": str(result.get("created_at", "") or datetime.now().isoformat(timespec="seconds")),
|
|
}
|
|
session["compacted_summary"] = summary
|
|
session["compacted_turns"] = int(session.get("compacted_turns", 0) or 0) + max(1, compacted_count // 2)
|
|
snapshots = list(session.get("compaction_snapshots", []) or [])
|
|
snapshots.append(snapshot)
|
|
session["compaction_snapshots"] = snapshots[-8:]
|
|
session["last_compaction_reason"] = snapshot["reason"]
|
|
session["last_compaction_at"] = snapshot["created_at"]
|
|
session["updated_at"] = datetime.now().isoformat(timespec="seconds")
|
|
self.save(data)
|
|
self.append_log("loop_history", {"session_id": session_id, "event": "context_compacted", "compaction": snapshot})
|
|
|
|
def compacted_summary(self, session_id: str | None = None) -> str:
|
|
record = self.get(session_id) if session_id else self.active()
|
|
return str((record or {}).get("compacted_summary", "") or "")
|
|
|
|
@_synchronized
|
|
def append_turn_record(self, session_id: str, record: dict[str, Any]) -> None:
|
|
data = self.load()
|
|
session = data.get("sessions", {}).get(session_id)
|
|
if isinstance(session, dict):
|
|
recent = session.setdefault("recent_turn_records", [])
|
|
recent.append({
|
|
"turn_id": record.get("turn_id", ""),
|
|
"status": record.get("status", ""),
|
|
"tool_round_count": len(record.get("tool_rounds", [])),
|
|
"final_preview": str(record.get("final_assistant_text", ""))[:240],
|
|
"created_at": record.get("created_at", ""),
|
|
})
|
|
session["recent_turn_records"] = recent[-6:]
|
|
session["updated_at"] = datetime.now().isoformat(timespec="seconds")
|
|
self.save(data)
|
|
self.append_log("loop_history", {"session_id": session_id, **record})
|
|
|
|
def working_dir(self, session_id: str | None = None) -> Path:
|
|
record = self.get(session_id) if session_id else self.active()
|
|
if record is None:
|
|
record = self.create()
|
|
path = (self.workspace_root / str(record["working_dir"])).resolve()
|
|
if path != self.working_root and self.working_root not in path.parents:
|
|
raise ValueError(f"Session working_dir escapes .working_dir: {record.get('working_dir')}")
|
|
path.mkdir(parents=True, exist_ok=True)
|
|
return path
|
|
|
|
def artifact_checklist(self, session_id: str | None = None) -> dict[str, bool]:
|
|
root = self.working_dir(session_id)
|
|
idea_dir = root / "idea2video"
|
|
idea_scene_dirs = sorted(path for path in idea_dir.glob("scene_*") if path.is_dir()) if idea_dir.exists() else []
|
|
idea_scene_storyboards = [path / "storyboard.json" for path in idea_scene_dirs]
|
|
idea_scene_camera_trees = [path / "camera_tree.json" for path in idea_scene_dirs]
|
|
idea_scene_shot_desc_groups = [list((scene / "shots").glob("*/shot_description.json")) for scene in idea_scene_dirs]
|
|
idea_scene_selector_outputs = [output for scene in idea_scene_dirs for output in (scene / "shots").glob("*/*_selector_output.json")]
|
|
|
|
script_shots = root / "script2video" / "shots"
|
|
script_shot_descs = list(script_shots.glob("*/shot_description.json")) if script_shots.exists() else []
|
|
script_selector_outputs = list(script_shots.glob("*/*_selector_output.json")) if script_shots.exists() else []
|
|
|
|
novel_dir = root / "novel2video"
|
|
novel_events = list((novel_dir / "events").glob("event_*.json")) if novel_dir.exists() else []
|
|
novel_relevant_chunks = [path for path in (novel_dir / "relevant_chunks").glob("event_*/*") if path.is_file()] if novel_dir.exists() else []
|
|
novel_scenes = list((novel_dir / "scenes").glob("event_*/scene_*.json")) if novel_dir.exists() else []
|
|
novel_event_chars = list((novel_dir / "global_information" / "characters" / "event_level").glob("event_*_characters.json")) if novel_dir.exists() else []
|
|
novel_level_chars = list((novel_dir / "global_information" / "characters" / "novel_level").glob("novel_characters_after_event_*.json")) if novel_dir.exists() else []
|
|
return {
|
|
"idea2video/story.txt": (idea_dir / "story.txt").exists(),
|
|
"idea2video/characters.json": (idea_dir / "characters.json").exists(),
|
|
"idea2video/script.json": (idea_dir / "script.json").exists(),
|
|
"idea2video/scene_*/storyboard.json": bool(idea_scene_storyboards) and all(path.exists() for path in idea_scene_storyboards),
|
|
"idea2video/scene_*/camera_tree.json": bool(idea_scene_camera_trees) and all(path.exists() for path in idea_scene_camera_trees),
|
|
"idea2video/scene_*/shots/*/shot_description.json": bool(idea_scene_shot_desc_groups) and all(idea_scene_shot_desc_groups),
|
|
"idea2video/scene_*/shots/*/*_selector_output.json": bool(idea_scene_selector_outputs),
|
|
"idea2video/final_video.mp4": (idea_dir / "final_video.mp4").exists(),
|
|
"script2video/script.txt": (root / "script2video" / "script.txt").exists(),
|
|
"script2video/characters.json": (root / "script2video" / "characters.json").exists(),
|
|
"script2video/storyboard.json": (root / "script2video" / "storyboard.json").exists(),
|
|
"script2video/shots/*/shot_description.json": bool(script_shot_descs),
|
|
"script2video/camera_tree.json": (root / "script2video" / "camera_tree.json").exists(),
|
|
"script2video/shots/*/*_selector_output.json": bool(script_selector_outputs),
|
|
"script2video/final_video.mp4": (root / "script2video" / "final_video.mp4").exists(),
|
|
"novel2video/novel/novel.txt": (novel_dir / "novel" / "novel.txt").exists(),
|
|
"novel2video/novel/novel_compressed.txt": (novel_dir / "novel" / "novel_compressed.txt").exists(),
|
|
"novel2video/events/event_*.json": bool(novel_events),
|
|
"novel2video/relevant_chunks/event_*": bool(novel_relevant_chunks),
|
|
"novel2video/scenes/event_*/scene_*.json": bool(novel_scenes),
|
|
"novel2video/global_information/characters/event_level/*.json": bool(novel_event_chars),
|
|
"novel2video/global_information/characters/novel_level/*.json": bool(novel_level_chars),
|
|
}
|
|
|
|
def memory_text(self) -> str:
|
|
return self.memory_path.read_text(encoding="utf-8") if self.memory_path.exists() else ""
|
|
|
|
def write_memory(self, text: str) -> None:
|
|
self.memory_path.write_text(text, encoding="utf-8")
|
|
|
|
def append_log(self, name: str, payload: dict[str, Any]) -> None:
|
|
event = {"timestamp": datetime.now().isoformat(timespec="seconds"), **payload}
|
|
path = self.logs_dir / f"{name}.jsonl"
|
|
with path.open("a", encoding="utf-8") as f:
|
|
f.write(json.dumps(event, ensure_ascii=False, default=str) + "\n")
|
|
|
|
def snapshot(self) -> dict[str, Any]:
|
|
active = self.active()
|
|
if active is None:
|
|
return {"active_session_id": "", "session": None}
|
|
return {"active_session_id": active["session_id"], "session": active, "artifact_checklist": self.artifact_checklist(active["session_id"])}
|
|
|
|
def _with_session_defaults(self, record: dict[str, Any]) -> dict[str, Any]:
|
|
item = dict(record)
|
|
item.setdefault("compacted_summary", "")
|
|
item.setdefault("compacted_turns", 0)
|
|
item.setdefault("compaction_snapshots", [])
|
|
item.setdefault("last_compaction_reason", "")
|
|
item.setdefault("last_compaction_at", "")
|
|
item.setdefault("recent_turn_records", [])
|
|
return item
|
|
|
|
def _new_session_id(self, source: str, sessions: dict[str, Any]) -> str:
|
|
stamp = datetime.now().strftime("%Y%m%d-%H%M%S")
|
|
slug = (re.sub(r"[^a-zA-Z0-9]+", "-", source.lower()).strip("-")[:32].strip("-") or "vimax")
|
|
return self._dedupe_session_id(f"{stamp}-{slug}", sessions)
|
|
|
|
def _dedupe_session_id(self, base: str, sessions: dict[str, Any]) -> str:
|
|
candidate = base
|
|
counter = 2
|
|
while candidate in sessions:
|
|
candidate = f"{base}-{counter}"
|
|
counter += 1
|
|
return candidate
|
|
|
|
def _normalize_session_id(self, session_id: str | None) -> str:
|
|
raw = str(session_id or "").strip()
|
|
if not raw:
|
|
raise ValueError("session_id cannot be empty")
|
|
normalized = re.sub(r"[^a-zA-Z0-9]+", "-", raw).strip("-")[:96]
|
|
if not normalized:
|
|
raise ValueError(f"Invalid session_id: {session_id}")
|
|
return normalized
|
|
|
|
def _working_dir_for_id(self, session_id: str) -> Path:
|
|
path = (self.working_root / session_id).resolve()
|
|
if path != self.working_root and self.working_root not in path.parents:
|
|
raise ValueError(f"Session path escapes .working_dir: {session_id}")
|
|
return path
|