From a347ac4bcbdacc637b839be7eb962304603c6815 Mon Sep 17 00:00:00 2001 From: Chris Guidry Date: Fri, 27 Mar 2026 17:46:23 -0400 Subject: [PATCH] Add entity path index and write-through cache for faster tool calls MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Tool-heavy turns were doing 30+ filesystem existence checks per turn — every entity lookup tried 6 entity type directories sequentially. Now the MCP server builds a name→path index by globbing the world directory once at startup. All entity lookups are dict lookups. The index also has a write-through cache for parsed entity data. establish() and mark() update both the index and cache after writing, so subsequent reads in the same turn never hit disk. Everything is channeled through the same set of tools so the cache is always current. Co-Authored-By: Claude Opus 4.6 (1M context) --- src/storied/engine.py | 15 +++- src/storied/mcp_server.py | 13 +++- src/storied/tools.py | 140 +++++++++++++++++++++++++++++++------- tests/test_entities.py | 119 +++++++++++++++++++++++++++++++- 4 files changed, 258 insertions(+), 29 deletions(-) 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"]) -- 2.51.2