diff --git a/src/storied/engine.py b/src/storied/engine.py index 3493c20..89a2958 100644 --- a/src/storied/engine.py +++ b/src/storied/engine.py @@ -297,10 +297,23 @@ class DMEngine: return resolver.load(name, content_type=content_type) def _find_entity(self, name: str) -> dict | None: - """Find an entity by name, searching entity directories in priority order.""" + """Find an entity by name, using the entity index for O(1) lookups.""" if not self.world_id: return None + # Fast path: use the entity index from the MCP server + entity_index = self._mcp.entity_index + if entity_index: + path = entity_index.resolve(name) + if path and path.exists(): + return { + "name": name, + "body": path.read_text(), + "entity_type": path.parent.name, + } + return None + + # Fallback: filesystem scan entity = load_entity_content(name, self.world_id, self.base_path) if entity: return { diff --git a/src/storied/mcp_server.py b/src/storied/mcp_server.py index dcc27b6..0872282 100644 --- a/src/storied/mcp_server.py +++ b/src/storied/mcp_server.py @@ -19,6 +19,7 @@ from mcp.types import TextContent, Tool from storied.log import CampaignLog from storied.tools import ( + EntityIndex, PLANNER_TOOL_DEFINITIONS, SEEDER_TOOL_DEFINITIONS, TOOL_DEFINITIONS, @@ -59,9 +60,13 @@ def _find_free_port() -> int: class MCPServerHandle: """Handle to a running in-process MCP server.""" - def __init__(self, port: int, thread: threading.Thread): + def __init__( + self, port: int, thread: threading.Thread, + entity_index: EntityIndex | None = None, + ): self.port = port self.url = f"http://127.0.0.1:{port}/sse" + self.entity_index = entity_index self._thread = thread def stop(self) -> None: @@ -85,6 +90,9 @@ def start_server( if campaign_log is None: campaign_log = CampaignLog(world_id, base_path) + world_dir = base_path / "worlds" / world_id + entity_index = EntityIndex(world_dir) + definitions = TOOL_SETS.get(tool_set, TOOL_DEFINITIONS) executor = EXECUTORS.get(tool_set, execute_tool) mcp_tools = [_to_mcp_tool(d) for d in definitions] @@ -104,6 +112,7 @@ def start_server( player_id=player_id, base_path=base_path, campaign_log=campaign_log, + entity_index=entity_index, ) return [TextContent(type="text", text=str(result))] @@ -143,4 +152,4 @@ def start_server( except OSError: time.sleep(0.1) - return MCPServerHandle(port, thread) + return MCPServerHandle(port, thread, entity_index=entity_index) diff --git a/src/storied/tools.py b/src/storied/tools.py index 3b37de4..5156ad2 100644 --- a/src/storied/tools.py +++ b/src/storied/tools.py @@ -10,6 +10,14 @@ from pathlib import Path import yaml +from storied.character import create_character as char_create +from storied.character import update_character as char_update +from storied.content import ContentResolver +from storied.dice import roll as dice_roll +from storied.log import CampaignLog, load_log +from storied.session import name_to_slug +from storied.session import update_session as session_update + # Per-file locks for thread-safe entity writes (establish, mark) _file_locks: dict[Path, threading.Lock] = {} _file_locks_lock = threading.Lock() @@ -22,13 +30,38 @@ def _get_file_lock(path: Path) -> threading.Lock: _file_locks[path] = threading.Lock() return _file_locks[path] -from storied.character import create_character as char_create -from storied.character import update_character as char_update -from storied.content import ContentResolver -from storied.dice import roll as dice_roll -from storied.log import CampaignLog, load_log -from storied.session import name_to_slug -from storied.session import update_session as session_update + +class EntityIndex: + """Name→path index with write-through entity cache. + + Built once at startup by globbing the world directory. Lookups are + dict lookups instead of 6 sequential Path.exists() calls. The cache + stores parsed entity dicts so repeated loads within a turn skip disk I/O. + establish() and mark() update both the index and cache on write. + """ + + def __init__(self, world_dir: Path | None = None): + self._paths: dict[str, Path] = {} + self._cache: dict[Path, dict] = {} + if world_dir and world_dir.exists(): + for md in world_dir.rglob("*.md"): + self._paths[md.stem] = md + + def resolve(self, name: str) -> Path | None: + """Look up an entity's file path by name.""" + return self._paths.get(name) + + def register(self, name: str, path: Path) -> None: + """Register or update an entity's path in the index.""" + self._paths[name] = path + + def cache_get(self, path: Path) -> dict | None: + """Get a cached parsed entity, or None if not cached.""" + return self._cache.get(path) + + def cache_put(self, path: Path, data: dict) -> None: + """Store a parsed entity in the cache.""" + self._cache[path] = data def roll(notation: str, reason: str | None = None) -> dict: @@ -240,6 +273,7 @@ def set_scene( base_path: Path | None = None, campaign_log: CampaignLog | None = None, world_id: str = "default", + entity_index: EntityIndex | None = None, ) -> str: """Call this after every response. Logs what happened, advances the clock, and updates the scene state. @@ -297,6 +331,7 @@ def set_scene( if event and present and world_id: marked = _auto_mark_present( present, event, world_id, base_path, campaign_log, + entity_index=entity_index, ) if marked: parts.append(f"Auto-marked: {', '.join(marked)}") @@ -314,6 +349,7 @@ def establish( will: list[str] | None = None, world_id: str | None = None, base_path: Path | None = None, + entity_index: EntityIndex | None = None, ) -> str: """Establish or update an entity in the world. @@ -361,7 +397,7 @@ def establish( lock = _get_file_lock(file_path) with lock: # Load existing content if file exists (for partial updates) - existing = _load_entity(file_path) + existing = _load_entity(file_path, entity_index) # Merge with existing content (new values override) if description is None: @@ -380,12 +416,28 @@ def establish( file_content = _format_entity(name, description, location, knows, wants, will, was) file_path.write_text(file_content) + # Write-through: update index and cache + updated_data = { + "description": description, "location": location, + "knows": knows, "wants": wants, "will": will, "was": was, + } + if entity_index: + entity_index.register(name, file_path) + entity_index.cache_put(file_path, updated_data) + action = "Updated" if existing else "Established" return f"{action} {entity_type.rstrip('s')} '{name}'" -def _load_entity(file_path: Path) -> dict: +def _load_entity( + file_path: Path, entity_index: EntityIndex | None = None, +) -> dict: """Load an existing entity file and parse its structure.""" + if entity_index: + cached = entity_index.cache_get(file_path) + if cached is not None: + return cached + if not file_path.exists(): return {} @@ -430,6 +482,9 @@ def _load_entity(file_path: Path) -> dict: if was_match: result["was"] = _parse_list_items(was_match.group(1)) + if entity_index: + entity_index.cache_put(file_path, result) + return result @@ -499,6 +554,7 @@ def _auto_mark_present( world_id: str, base_path: Path | None, campaign_log: CampaignLog | None, + entity_index: EntityIndex | None = None, ) -> list[str]: """Auto-mark present entities with the current event. @@ -516,20 +572,27 @@ def _auto_mark_present( continue name = link_match.group(1) - # Try each entity type directory - for etype in ("npcs", "locations", "items", "factions"): - file_path = base_path / "worlds" / world_id / etype / f"{name}.md" - if file_path.exists(): - mark( - entity_type=etype, - name=name, - event=event, - world_id=world_id, - base_path=base_path, - campaign_log=campaign_log, - ) - marked.append(name) - break + # Resolve via index (O(1)) or fall back to directory scan + file_path = entity_index.resolve(name) if entity_index else None + if file_path is None: + for etype in ("npcs", "locations", "items", "factions"): + candidate = base_path / "worlds" / world_id / etype / f"{name}.md" + if candidate.exists(): + file_path = candidate + break + + if file_path and file_path.exists(): + entity_type = file_path.parent.name + mark( + entity_type=entity_type, + name=name, + event=event, + world_id=world_id, + base_path=base_path, + campaign_log=campaign_log, + entity_index=entity_index, + ) + marked.append(name) return marked @@ -542,6 +605,7 @@ def mark( world_id: str | None = None, base_path: Path | None = None, campaign_log: CampaignLog | None = None, + entity_index: EntityIndex | None = None, ) -> str: """Record an event in an entity's history (## Was section). @@ -570,7 +634,14 @@ def mark( if base_path is None: base_path = Path.cwd() - file_path = base_path / "worlds" / world_id / entity_type / f"{name}.md" + # Resolve path via index or filesystem + if entity_index: + file_path = entity_index.resolve(name) + if file_path is None: + file_path = base_path / "worlds" / world_id / entity_type / f"{name}.md" + else: + file_path = base_path / "worlds" / world_id / entity_type / f"{name}.md" + if not file_path.exists(): return f"Error: Entity '{name}' not found in {entity_type}" @@ -584,7 +655,7 @@ def mark( lock = _get_file_lock(file_path) with lock: # Load existing entity - existing = _load_entity(file_path) + existing = _load_entity(file_path, entity_index) # Append to Was section was = existing.get("was", []) @@ -610,6 +681,17 @@ def mark( ) file_path.write_text(file_content) + # Write-through cache + if entity_index: + entity_index.cache_put(file_path, { + "description": existing.get("description", ""), + "location": existing.get("location", ""), + "knows": existing.get("knows", []), + "wants": existing.get("wants", []), + "will": will, + "was": was, + }) + result = f"Marked: {event}" if resolved: if len(resolved) == 1: @@ -985,6 +1067,7 @@ def execute_tool( player_id: str = "default", base_path: Path | None = None, campaign_log: CampaignLog | None = None, + entity_index: EntityIndex | None = None, ) -> str: """Execute a tool by name with the given input. @@ -1060,6 +1143,7 @@ def execute_tool( base_path=base_path, campaign_log=campaign_log, world_id=world_id or "default", + entity_index=entity_index, ) elif tool_name == "establish": @@ -1073,6 +1157,7 @@ def execute_tool( will=tool_input.get("will"), world_id=world_id, base_path=base_path, + entity_index=entity_index, ) elif tool_name == "mark": @@ -1084,6 +1169,7 @@ def execute_tool( world_id=world_id, base_path=base_path, campaign_log=campaign_log, + entity_index=entity_index, ) elif tool_name == "note_discovery": @@ -1120,6 +1206,7 @@ def planner_execute_tool( world_id: str | None = None, base_path: Path | None = None, campaign_log: CampaignLog | None = None, + entity_index: EntityIndex | None = None, ) -> str: """Execute a planner-allowed tool. Rejects anything outside the allowed set.""" if tool_name not in PLANNER_TOOLS: @@ -1130,6 +1217,7 @@ def planner_execute_tool( world_id=world_id, base_path=base_path, campaign_log=campaign_log, + entity_index=entity_index, ) @@ -1145,6 +1233,7 @@ def seeder_execute_tool( player_id: str | None = None, base_path: Path | None = None, campaign_log: CampaignLog | None = None, + entity_index: EntityIndex | None = None, ) -> str: """Execute a seeder-allowed tool. Rejects anything outside the allowed set.""" if tool_name not in SEEDER_TOOLS: @@ -1156,4 +1245,5 @@ def seeder_execute_tool( player_id=player_id or "default", base_path=base_path, campaign_log=campaign_log, + entity_index=entity_index, ) diff --git a/tests/test_entities.py b/tests/test_entities.py index 6b031e6..9078a06 100644 --- a/tests/test_entities.py +++ b/tests/test_entities.py @@ -10,7 +10,7 @@ from storied.session import ( load_entity_content, resolve_wiki_link, ) -from storied.tools import establish, mark +from storied.tools import EntityIndex, establish, mark @pytest.fixture @@ -452,3 +452,120 @@ class TestLoadEntityContent: def test_load_entity_content_not_found(self, world_base: Path): entity = load_entity_content("Nobody", "test-world", world_base) assert entity is None + + +# ── EntityIndex ────────────────────────────────────────────────────────── + + +@pytest.fixture +def indexed_world(world_base: Path) -> tuple[Path, EntityIndex]: + """Create a world with entities and build an index.""" + world_dir = world_base / "worlds" / "test-world" + for etype, name in [("npcs", "Vera"), ("locations", "Tavern"), ("items", "Sword")]: + d = world_dir / etype + d.mkdir(parents=True, exist_ok=True) + (d / f"{name}.md").write_text(f"# {name}\n\n## Is\n\nA {etype[:-1]}.\n") + return world_dir, EntityIndex(world_dir) + + +class TestEntityIndex: + """Tests for the entity path index and write-through cache.""" + + def test_resolve_existing(self, indexed_world: tuple[Path, EntityIndex]): + _, index = indexed_world + path = index.resolve("Vera") + assert path is not None + assert path.name == "Vera.md" + assert "npcs" in str(path) + + def test_resolve_missing(self, indexed_world: tuple[Path, EntityIndex]): + _, index = indexed_world + assert index.resolve("Nobody") is None + + def test_register_new_entity(self, indexed_world: tuple[Path, EntityIndex]): + world_dir, index = indexed_world + new_path = world_dir / "npcs" / "Henrik.md" + index.register("Henrik", new_path) + assert index.resolve("Henrik") == new_path + + def test_cache_miss_returns_none(self, indexed_world: tuple[Path, EntityIndex]): + world_dir, index = indexed_world + assert index.cache_get(world_dir / "npcs" / "Vera.md") is None + + def test_cache_put_and_get(self, indexed_world: tuple[Path, EntityIndex]): + world_dir, index = indexed_world + path = world_dir / "npcs" / "Vera.md" + data = {"description": "A tavern keeper.", "knows": ["a secret"]} + index.cache_put(path, data) + assert index.cache_get(path) == data + + def test_cache_survives_register(self, indexed_world: tuple[Path, EntityIndex]): + world_dir, index = indexed_world + path = world_dir / "npcs" / "Vera.md" + data = {"description": "Vera"} + index.cache_put(path, data) + index.register("Vera", path) + assert index.cache_get(path) == data + + def test_empty_world_dir(self, tmp_path: Path): + world_dir = tmp_path / "worlds" / "empty" + world_dir.mkdir(parents=True) + index = EntityIndex(world_dir) + assert index.resolve("anything") is None + + def test_none_world_dir(self): + index = EntityIndex(None) + assert index.resolve("anything") is None + + +class TestEstablishWithIndex: + """Tests that establish writes through to the index and cache.""" + + def test_establish_registers_in_index(self, world_base: Path): + world_dir = world_base / "worlds" / "test-world" + index = EntityIndex(world_dir) + establish( + entity_type="npcs", name="New NPC", + description="A stranger.", + world_id="test-world", base_path=world_base, + entity_index=index, + ) + assert index.resolve("New NPC") is not None + + def test_establish_caches_entity(self, world_base: Path): + world_dir = world_base / "worlds" / "test-world" + index = EntityIndex(world_dir) + establish( + entity_type="npcs", name="Cached NPC", + description="Cached.", + knows=["a secret"], + world_id="test-world", base_path=world_base, + entity_index=index, + ) + path = index.resolve("Cached NPC") + cached = index.cache_get(path) + assert cached is not None + assert cached["description"] == "Cached." + assert cached["knows"] == ["a secret"] + + def test_mark_updates_cache( + self, world_base: Path, campaign_log: CampaignLog, + ): + world_dir = world_base / "worlds" / "test-world" + index = EntityIndex(world_dir) + establish( + entity_type="npcs", name="Markable", + description="An NPC.", + world_id="test-world", base_path=world_base, + entity_index=index, + ) + mark( + entity_type="npcs", name="Markable", + event="Something happened", + world_id="test-world", base_path=world_base, + campaign_log=campaign_log, + entity_index=index, + ) + path = index.resolve("Markable") + cached = index.cache_get(path) + assert any("Something happened" in w for w in cached["was"])