mirror of
https://github.com/overcuriousity/trace.git
synced 2025-12-21 13:32:20 +00:00
Restructure codebase for AI agent optimization
Major refactoring to organize code into focused, single-responsibility modules that are easier for AI coding agents and developers to navigate and modify. **Module Reorganization:** Models Package (trace/models/): - Moved models.py content into models/__init__.py - Extracted IOC extraction into models/extractors/ioc_extractor.py (236 lines) - Extracted tag extraction into models/extractors/tag_extractor.py (34 lines) - Reduced duplication and improved maintainability Storage Package (trace/storage_impl/): - Split storage.py (402 lines) into focused modules: - storage.py: Main Storage class (112 lines) - state_manager.py: StateManager for context/settings (92 lines) - lock_manager.py: Cross-platform file locking (87 lines) - demo_data.py: Demo case creation (143 lines) - Added backward-compatible wrapper at trace/storage.py TUI Utilities (trace/tui/): - Created rendering package: - colors.py: Color pair constants and initialization (43 lines) - text_renderer.py: Text rendering with highlighting (137 lines) - Created handlers package: - export_handler.py: Export functionality (238 lines) - Main tui.py (3307 lines) remains for future refactoring **Benefits:** - Smaller, focused files (most < 250 lines) - Clear single responsibilities - Easier to locate and modify specific functionality - Better separation of concerns - Reduced cognitive load for AI agents - All tests pass, no features removed **Testing:** - All existing tests pass - Imports verified - CLI and storage functionality tested - Backward compatibility maintained Updated CLAUDE.md to document new architecture and AI optimization strategy.
This commit is contained in:
92
trace/storage_impl/state_manager.py
Normal file
92
trace/storage_impl/state_manager.py
Normal file
@@ -0,0 +1,92 @@
|
||||
"""State manager for active context and settings"""
|
||||
|
||||
import json
|
||||
from pathlib import Path
|
||||
from typing import Optional, TYPE_CHECKING
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from .storage import Storage
|
||||
|
||||
DEFAULT_APP_DIR = Path.home() / ".trace"
|
||||
|
||||
|
||||
class StateManager:
|
||||
"""Manages active context and user settings"""
|
||||
|
||||
def __init__(self, app_dir: Path = DEFAULT_APP_DIR):
|
||||
self.app_dir = app_dir
|
||||
self.state_file = self.app_dir / "state"
|
||||
self.settings_file = self.app_dir / "settings.json"
|
||||
self._ensure_app_dir()
|
||||
|
||||
def _ensure_app_dir(self):
|
||||
if not self.app_dir.exists():
|
||||
self.app_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
def set_active(self, case_id: Optional[str] = None, evidence_id: Optional[str] = None):
|
||||
state = self.get_active()
|
||||
state["case_id"] = case_id
|
||||
state["evidence_id"] = evidence_id
|
||||
# Atomic write: write to temp file then rename
|
||||
temp_file = self.state_file.with_suffix(".tmp")
|
||||
with open(temp_file, 'w', encoding='utf-8') as f:
|
||||
json.dump(state, f, ensure_ascii=False)
|
||||
temp_file.replace(self.state_file)
|
||||
|
||||
def get_active(self) -> dict:
|
||||
if not self.state_file.exists():
|
||||
return {"case_id": None, "evidence_id": None}
|
||||
try:
|
||||
with open(self.state_file, 'r', encoding='utf-8') as f:
|
||||
return json.load(f)
|
||||
except (json.JSONDecodeError, IOError):
|
||||
return {"case_id": None, "evidence_id": None}
|
||||
|
||||
def validate_and_clear_stale(self, storage: 'Storage') -> str:
|
||||
"""Validate active state against storage and clear stale references.
|
||||
Returns warning message if state was cleared, empty string otherwise."""
|
||||
state = self.get_active()
|
||||
case_id = state.get("case_id")
|
||||
evidence_id = state.get("evidence_id")
|
||||
warning = ""
|
||||
|
||||
if case_id:
|
||||
case = storage.get_case(case_id)
|
||||
if not case:
|
||||
warning = f"Active case (ID: {case_id[:8]}...) no longer exists. Clearing active context."
|
||||
self.set_active(None, None)
|
||||
return warning
|
||||
|
||||
# Validate evidence if set
|
||||
if evidence_id:
|
||||
_, evidence = storage.find_evidence(evidence_id)
|
||||
if not evidence:
|
||||
warning = f"Active evidence (ID: {evidence_id[:8]}...) no longer exists. Clearing to case level."
|
||||
self.set_active(case_id, None)
|
||||
return warning
|
||||
|
||||
elif evidence_id:
|
||||
# Evidence set but no case - invalid state
|
||||
warning = "Invalid state: evidence set without case. Clearing active context."
|
||||
self.set_active(None, None)
|
||||
return warning
|
||||
|
||||
return warning
|
||||
|
||||
def get_settings(self) -> dict:
|
||||
if not self.settings_file.exists():
|
||||
return {"pgp_enabled": True}
|
||||
try:
|
||||
with open(self.settings_file, 'r', encoding='utf-8') as f:
|
||||
return json.load(f)
|
||||
except (json.JSONDecodeError, IOError):
|
||||
return {"pgp_enabled": True}
|
||||
|
||||
def set_setting(self, key: str, value):
|
||||
settings = self.get_settings()
|
||||
settings[key] = value
|
||||
# Atomic write: write to temp file then rename
|
||||
temp_file = self.settings_file.with_suffix(".tmp")
|
||||
with open(temp_file, 'w', encoding='utf-8') as f:
|
||||
json.dump(settings, f, ensure_ascii=False)
|
||||
temp_file.replace(self.settings_file)
|
||||
Reference in New Issue
Block a user