| 1 | from __future__ import annotations |
| 2 | |
| 3 | import json |
| 4 | import logging |
| 5 | import os |
| 6 | import re |
| 7 | from contextlib import contextmanager |
| 8 | from datetime import datetime |
| 9 | from functools import wraps |
| 10 | from pathlib import Path |
| 11 | from typing import Any |
| 12 | |
| 13 | try: |
| 14 | import fcntl |
| 15 | except ImportError: # pragma: no cover - non-POSIX platforms |
| 16 | fcntl = None |
| 17 | |
| 18 | |
| 19 | STALE_KEYS = ["story", "characters", "script", "storyboard", "shot_descriptions", "camera_tree", "frames", "clips", "final_video"] |
| 20 | |
| 21 | |
| 22 | def _synchronized(method): |
| 23 | """Hold the index file lock across a read-modify-write cycle. |
| 24 | |
| 25 | Every mutator loads the whole sessions file, edits it, and saves it back; |
| 26 | without a lock, two concurrent writers (threads or processes) silently |
| 27 | drop each other's updates. |
| 28 | """ |
| 29 | |
| 30 | @wraps(method) |
| 31 | def wrapper(self, *args, **kwargs): |
| 32 | with self._locked(): |
| 33 | return method(self, *args, **kwargs) |
| 34 | |
| 35 | return wrapper |
| 36 | |
| 37 | |
| 38 | class SessionIndex: |
| 39 | def __init__(self, workspace_root: str | Path) -> None: |
| 40 | self.workspace_root = Path(workspace_root).resolve() |
| 41 | self.vimax_dir = self.workspace_root / ".vimax" |
| 42 | self.sessions_path = self.vimax_dir / "sessions.json" |
| 43 | self.memory_path = self.vimax_dir / "memory.md" |
| 44 | self.logs_dir = self.vimax_dir / "logs" |
| 45 | self.working_root = self.workspace_root / ".working_dir" |
| 46 | self.vimax_dir.mkdir(parents=True, exist_ok=True) |
| 47 | self.logs_dir.mkdir(parents=True, exist_ok=True) |
| 48 | self.working_root.mkdir(parents=True, exist_ok=True) |
| 49 | if not self.memory_path.exists(): |
| 50 | self.memory_path.write_text("# User Preferences\n", encoding="utf-8") |
| 51 | if not self.sessions_path.exists(): |
| 52 | self.save({"active_session_id": "", "sessions": {}}) |
| 53 | |
| 54 | @contextmanager |
| 55 | def _locked(self): |
| 56 | if fcntl is None: |
| 57 | yield |
| 58 | return |
| 59 | lock_path = self.vimax_dir / "sessions.lock" |
| 60 | with open(lock_path, "a+", encoding="utf-8") as handle: |
| 61 | fcntl.flock(handle, fcntl.LOCK_EX) |
| 62 | try: |
| 63 | yield |
| 64 | finally: |
| 65 | fcntl.flock(handle, fcntl.LOCK_UN) |
| 66 | |
| 67 | def load(self) -> dict[str, Any]: |
| 68 | try: |
| 69 | return json.loads(self.sessions_path.read_text(encoding="utf-8")) |
| 70 | except FileNotFoundError: |
| 71 | return {"active_session_id": "", "sessions": {}} |
| 72 | except json.JSONDecodeError: |
| 73 | # A corrupt file usually means a crash mid-write. Returning empty |
| 74 | # state is fine for this call, but the next save() would overwrite |
| 75 | # the file and destroy every session — keep the evidence first. |
| 76 | backup = self.sessions_path.with_name(f"sessions.json.corrupt-{datetime.now().strftime('%Y%m%d-%H%M%S-%f')}") |
| 77 | try: |
| 78 | os.replace(self.sessions_path, backup) |
| 79 | logging.error("sessions.json was corrupt; preserved it at %s and starting with empty state", backup) |
| 80 | except OSError: |
| 81 | logging.error("sessions.json is corrupt and could not be backed up; starting with empty state") |
| 82 | return {"active_session_id": "", "sessions": {}} |
| 83 | |
| 84 | def save(self, data: dict[str, Any]) -> None: |
| 85 | tmp_path = self.sessions_path.with_name("sessions.json.tmp") |
| 86 | tmp_path.write_text(json.dumps(data, ensure_ascii=False, indent=2), encoding="utf-8") |
| 87 | os.replace(tmp_path, self.sessions_path) |
| 88 | |
| 89 | def active(self) -> dict[str, Any] | None: |
| 90 | data = self.load() |
| 91 | session_id = str(data.get("active_session_id", "")) |
| 92 | if not session_id: |
| 93 | return None |
| 94 | record = data.get("sessions", {}).get(session_id) |
| 95 | return self._with_session_defaults(record) if isinstance(record, dict) else None |
| 96 | |
| 97 | def get(self, session_id: str) -> dict[str, Any] | None: |
| 98 | normalized = self._normalize_session_id(session_id) |
| 99 | record = self.load().get("sessions", {}).get(normalized) |
| 100 | return self._with_session_defaults(record) if isinstance(record, dict) else None |
| 101 | |
| 102 | @_synchronized |
| 103 | def create( |
| 104 | self, |
| 105 | idea: str = "", |
| 106 | user_requirement: str = "", |
| 107 | style: str = "", |
| 108 | session_id: str | None = None, |
| 109 | project_name: str = "", |
| 110 | ) -> dict[str, Any]: |
| 111 | data = self.load() |
| 112 | sessions = data.setdefault("sessions", {}) |
| 113 | clean_project_name = str(project_name or "").strip()[:64] |
| 114 | final_id = self._normalize_session_id(session_id) if session_id else self._new_session_id(clean_project_name or idea or user_requirement or "vimax", sessions) |
| 115 | if final_id in sessions: |
| 116 | final_id = self._dedupe_session_id(final_id, sessions) |
| 117 | now = datetime.now().isoformat(timespec="seconds") |
| 118 | working_dir = self._working_dir_for_id(final_id) |
| 119 | (working_dir / "idea2video").mkdir(parents=True, exist_ok=True) |
| 120 | (working_dir / "script2video").mkdir(parents=True, exist_ok=True) |
| 121 | record = { |
| 122 | "session_id": final_id, |
| 123 | "project_name": clean_project_name, |
| 124 | "working_dir": str(working_dir.relative_to(self.workspace_root)), |
| 125 | "idea": idea, |
| 126 | "user_requirement": user_requirement, |
| 127 | "style": style, |
| 128 | "stage": "created", |
| 129 | "summary": "", |
| 130 | "stale": {key: False for key in STALE_KEYS}, |
| 131 | "recent_turn_records": [], |
| 132 | "compacted_summary": "", |
| 133 | "compacted_turns": 0, |
| 134 | "compaction_snapshots": [], |
| 135 | "last_compaction_reason": "", |
| 136 | "last_compaction_at": "", |
| 137 | "created_at": now, |
| 138 | "updated_at": now, |
| 139 | } |
| 140 | sessions[final_id] = record |
| 141 | data["active_session_id"] = final_id |
| 142 | self.save(data) |
| 143 | return record |
| 144 | |
| 145 | def get_or_create_active(self, idea: str = "", user_requirement: str = "", style: str = "") -> dict[str, Any]: |
| 146 | active = self.active() |
| 147 | if active is not None: |
| 148 | return active |
| 149 | return self.create(idea=idea, user_requirement=user_requirement, style=style) |
| 150 | |
| 151 | @_synchronized |
| 152 | def set_active(self, session_id: str) -> dict[str, Any]: |
| 153 | normalized = self._normalize_session_id(session_id) |
| 154 | data = self.load() |
| 155 | if normalized not in data.get("sessions", {}): |
| 156 | raise KeyError(f"Unknown session_id: {session_id}") |
| 157 | data["active_session_id"] = normalized |
| 158 | self.save(data) |
| 159 | return dict(data["sessions"][normalized]) |
| 160 | |
| 161 | @_synchronized |
| 162 | def update_stage(self, session_id: str, stage: str, summary: str = "") -> None: |
| 163 | data = self.load() |
| 164 | record = data.get("sessions", {}).get(session_id) |
| 165 | if not isinstance(record, dict): |
| 166 | raise KeyError(f"Unknown session_id: {session_id}") |
| 167 | record["stage"] = stage |
| 168 | if summary: |
| 169 | record["summary"] = summary |
| 170 | record["updated_at"] = datetime.now().isoformat(timespec="seconds") |
| 171 | self.save(data) |
| 172 | |
| 173 | @_synchronized |
| 174 | def mark_stale(self, session_id: str, keys: list[str]) -> None: |
| 175 | data = self.load() |
| 176 | record = data.get("sessions", {}).get(session_id) |
| 177 | if not isinstance(record, dict): |
| 178 | raise KeyError(f"Unknown session_id: {session_id}") |
| 179 | stale = record.setdefault("stale", {key: False for key in STALE_KEYS}) |
| 180 | for key in keys: |
| 181 | stale[key] = True |
| 182 | record["updated_at"] = datetime.now().isoformat(timespec="seconds") |
| 183 | self.save(data) |
| 184 | |
| 185 | @_synchronized |
| 186 | def update_compaction(self, session_id: str, result: dict[str, Any]) -> None: |
| 187 | data = self.load() |
| 188 | session = data.get("sessions", {}).get(session_id) |
| 189 | if not isinstance(session, dict): |
| 190 | raise KeyError(f"Unknown session_id: {session_id}") |
| 191 | summary = str(result.get("summary", "") or "") |
| 192 | compacted_count = int(result.get("compacted_message_count", 0) or 0) |
| 193 | snapshot = { |
| 194 | "level": len(session.get("compaction_snapshots", []) or []) + 1, |
| 195 | "reason": str(result.get("reason", "manual") or "manual"), |
| 196 | "mode": str(result.get("mode", "unknown") or "unknown"), |
| 197 | "summary": summary, |
| 198 | "preserved_messages": int(result.get("preserved_message_count", 0) or 0), |
| 199 | "compacted_message_count": compacted_count, |
| 200 | "estimated_tokens_before": int(result.get("estimated_tokens_before", 0) or 0), |
| 201 | "estimated_tokens_after": int(result.get("estimated_tokens_after", 0) or 0), |
| 202 | "created_at": str(result.get("created_at", "") or datetime.now().isoformat(timespec="seconds")), |
| 203 | } |
| 204 | session["compacted_summary"] = summary |
| 205 | session["compacted_turns"] = int(session.get("compacted_turns", 0) or 0) + max(1, compacted_count // 2) |
| 206 | snapshots = list(session.get("compaction_snapshots", []) or []) |
| 207 | snapshots.append(snapshot) |
| 208 | session["compaction_snapshots"] = snapshots[-8:] |
| 209 | session["last_compaction_reason"] = snapshot["reason"] |
| 210 | session["last_compaction_at"] = snapshot["created_at"] |
| 211 | session["updated_at"] = datetime.now().isoformat(timespec="seconds") |
| 212 | self.save(data) |
| 213 | self.append_log("loop_history", {"session_id": session_id, "event": "context_compacted", "compaction": snapshot}) |
| 214 | |
| 215 | def compacted_summary(self, session_id: str | None = None) -> str: |
| 216 | record = self.get(session_id) if session_id else self.active() |
| 217 | return str((record or {}).get("compacted_summary", "") or "") |
| 218 | |
| 219 | @_synchronized |
| 220 | def append_turn_record(self, session_id: str, record: dict[str, Any]) -> None: |
| 221 | data = self.load() |
| 222 | session = data.get("sessions", {}).get(session_id) |
| 223 | if isinstance(session, dict): |
| 224 | recent = session.setdefault("recent_turn_records", []) |
| 225 | recent.append({ |
| 226 | "turn_id": record.get("turn_id", ""), |
| 227 | "status": record.get("status", ""), |
| 228 | "tool_round_count": len(record.get("tool_rounds", [])), |
| 229 | "final_preview": str(record.get("final_assistant_text", ""))[:240], |
| 230 | "created_at": record.get("created_at", ""), |
| 231 | }) |
| 232 | session["recent_turn_records"] = recent[-6:] |
| 233 | session["updated_at"] = datetime.now().isoformat(timespec="seconds") |
| 234 | self.save(data) |
| 235 | self.append_log("loop_history", {"session_id": session_id, **record}) |
| 236 | |
| 237 | def working_dir(self, session_id: str | None = None) -> Path: |
| 238 | record = self.get(session_id) if session_id else self.active() |
| 239 | if record is None: |
| 240 | record = self.create() |
| 241 | path = (self.workspace_root / str(record["working_dir"])).resolve() |
| 242 | if path != self.working_root and self.working_root not in path.parents: |
| 243 | raise ValueError(f"Session working_dir escapes .working_dir: {record.get('working_dir')}") |
| 244 | path.mkdir(parents=True, exist_ok=True) |
| 245 | return path |
| 246 | |
| 247 | def artifact_checklist(self, session_id: str | None = None) -> dict[str, bool]: |
| 248 | root = self.working_dir(session_id) |
| 249 | idea_dir = root / "idea2video" |
| 250 | idea_scene_dirs = sorted(path for path in idea_dir.glob("scene_*") if path.is_dir()) if idea_dir.exists() else [] |
| 251 | idea_scene_storyboards = [path / "storyboard.json" for path in idea_scene_dirs] |
| 252 | idea_scene_camera_trees = [path / "camera_tree.json" for path in idea_scene_dirs] |
| 253 | idea_scene_shot_desc_groups = [list((scene / "shots").glob("*/shot_description.json")) for scene in idea_scene_dirs] |
| 254 | idea_scene_selector_outputs = [output for scene in idea_scene_dirs for output in (scene / "shots").glob("*/*_selector_output.json")] |
| 255 | |
| 256 | script_shots = root / "script2video" / "shots" |
| 257 | script_shot_descs = list(script_shots.glob("*/shot_description.json")) if script_shots.exists() else [] |
| 258 | script_selector_outputs = list(script_shots.glob("*/*_selector_output.json")) if script_shots.exists() else [] |
| 259 | |
| 260 | novel_dir = root / "novel2video" |
| 261 | novel_events = list((novel_dir / "events").glob("event_*.json")) if novel_dir.exists() else [] |
| 262 | novel_relevant_chunks = [path for path in (novel_dir / "relevant_chunks").glob("event_*/*") if path.is_file()] if novel_dir.exists() else [] |
| 263 | novel_scenes = list((novel_dir / "scenes").glob("event_*/scene_*.json")) if novel_dir.exists() else [] |
| 264 | novel_event_chars = list((novel_dir / "global_information" / "characters" / "event_level").glob("event_*_characters.json")) if novel_dir.exists() else [] |
| 265 | novel_level_chars = list((novel_dir / "global_information" / "characters" / "novel_level").glob("novel_characters_after_event_*.json")) if novel_dir.exists() else [] |
| 266 | return { |
| 267 | "idea2video/story.txt": (idea_dir / "story.txt").exists(), |
| 268 | "idea2video/characters.json": (idea_dir / "characters.json").exists(), |
| 269 | "idea2video/script.json": (idea_dir / "script.json").exists(), |
| 270 | "idea2video/scene_*/storyboard.json": bool(idea_scene_storyboards) and all(path.exists() for path in idea_scene_storyboards), |
| 271 | "idea2video/scene_*/camera_tree.json": bool(idea_scene_camera_trees) and all(path.exists() for path in idea_scene_camera_trees), |
| 272 | "idea2video/scene_*/shots/*/shot_description.json": bool(idea_scene_shot_desc_groups) and all(idea_scene_shot_desc_groups), |
| 273 | "idea2video/scene_*/shots/*/*_selector_output.json": bool(idea_scene_selector_outputs), |
| 274 | "idea2video/final_video.mp4": (idea_dir / "final_video.mp4").exists(), |
| 275 | "script2video/script.txt": (root / "script2video" / "script.txt").exists(), |
| 276 | "script2video/characters.json": (root / "script2video" / "characters.json").exists(), |
| 277 | "script2video/storyboard.json": (root / "script2video" / "storyboard.json").exists(), |
| 278 | "script2video/shots/*/shot_description.json": bool(script_shot_descs), |
| 279 | "script2video/camera_tree.json": (root / "script2video" / "camera_tree.json").exists(), |
| 280 | "script2video/shots/*/*_selector_output.json": bool(script_selector_outputs), |
| 281 | "script2video/final_video.mp4": (root / "script2video" / "final_video.mp4").exists(), |
| 282 | "novel2video/novel/novel.txt": (novel_dir / "novel" / "novel.txt").exists(), |
| 283 | "novel2video/novel/novel_compressed.txt": (novel_dir / "novel" / "novel_compressed.txt").exists(), |
| 284 | "novel2video/events/event_*.json": bool(novel_events), |
| 285 | "novel2video/relevant_chunks/event_*": bool(novel_relevant_chunks), |
| 286 | "novel2video/scenes/event_*/scene_*.json": bool(novel_scenes), |
| 287 | "novel2video/global_information/characters/event_level/*.json": bool(novel_event_chars), |
| 288 | "novel2video/global_information/characters/novel_level/*.json": bool(novel_level_chars), |
| 289 | } |
| 290 | |
| 291 | def memory_text(self) -> str: |
| 292 | return self.memory_path.read_text(encoding="utf-8") if self.memory_path.exists() else "" |
| 293 | |
| 294 | def write_memory(self, text: str) -> None: |
| 295 | self.memory_path.write_text(text, encoding="utf-8") |
| 296 | |
| 297 | def append_log(self, name: str, payload: dict[str, Any]) -> None: |
| 298 | event = {"timestamp": datetime.now().isoformat(timespec="seconds"), **payload} |
| 299 | path = self.logs_dir / f"{name}.jsonl" |
| 300 | with path.open("a", encoding="utf-8") as f: |
| 301 | f.write(json.dumps(event, ensure_ascii=False, default=str) + "\n") |
| 302 | |
| 303 | def snapshot(self) -> dict[str, Any]: |
| 304 | active = self.active() |
| 305 | if active is None: |
| 306 | return {"active_session_id": "", "session": None} |
| 307 | return {"active_session_id": active["session_id"], "session": active, "artifact_checklist": self.artifact_checklist(active["session_id"])} |
| 308 | |
| 309 | def _with_session_defaults(self, record: dict[str, Any]) -> dict[str, Any]: |
| 310 | item = dict(record) |
| 311 | item.setdefault("project_name", "") |
| 312 | item.setdefault("compacted_summary", "") |
| 313 | item.setdefault("compacted_turns", 0) |
| 314 | item.setdefault("compaction_snapshots", []) |
| 315 | item.setdefault("last_compaction_reason", "") |
| 316 | item.setdefault("last_compaction_at", "") |
| 317 | item.setdefault("recent_turn_records", []) |
| 318 | return item |
| 319 | |
| 320 | def _new_session_id(self, source: str, sessions: dict[str, Any]) -> str: |
| 321 | stamp = datetime.now().strftime("%Y%m%d-%H%M%S") |
| 322 | slug = (re.sub(r"[^a-zA-Z0-9]+", "-", source.lower()).strip("-")[:32].strip("-") or "vimax") |
| 323 | return self._dedupe_session_id(f"{stamp}-{slug}", sessions) |
| 324 | |
| 325 | def _dedupe_session_id(self, base: str, sessions: dict[str, Any]) -> str: |
| 326 | candidate = base |
| 327 | counter = 2 |
| 328 | while candidate in sessions: |
| 329 | candidate = f"{base}-{counter}" |
| 330 | counter += 1 |
| 331 | return candidate |
| 332 | |
| 333 | def _normalize_session_id(self, session_id: str | None) -> str: |
| 334 | raw = str(session_id or "").strip() |
| 335 | if not raw: |
| 336 | raise ValueError("session_id cannot be empty") |
| 337 | normalized = re.sub(r"[^a-zA-Z0-9]+", "-", raw).strip("-")[:96] |
| 338 | if not normalized: |
| 339 | raise ValueError(f"Invalid session_id: {session_id}") |
| 340 | return normalized |
| 341 | |
| 342 | def _working_dir_for_id(self, session_id: str) -> Path: |
| 343 | path = (self.working_root / session_id).resolve() |
| 344 | if path != self.working_root and self.working_root not in path.parents: |
| 345 | raise ValueError(f"Session path escapes .working_dir: {session_id}") |
| 346 | return path |
| 347 |