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