From 55949498e34c94db3af98cf29bc0e3b3c62dabb7 Mon Sep 17 00:00:00 2001 From: StatPan Date: Thu, 14 May 2026 14:19:36 +0900 Subject: [PATCH] feat(sdk): allow custom FileStore injection Co-authored-by: openhands --- .../conversation/impl/local_conversation.py | 10 +- .../openhands/sdk/conversation/state.py | 30 ++-- .../local/test_custom_file_store.py | 129 ++++++++++++++++++ 3 files changed, 156 insertions(+), 13 deletions(-) create mode 100644 tests/sdk/conversation/local/test_custom_file_store.py diff --git a/openhands-sdk/openhands/sdk/conversation/impl/local_conversation.py b/openhands-sdk/openhands/sdk/conversation/impl/local_conversation.py index af6ee4b2b9..a0f7d770b3 100644 --- a/openhands-sdk/openhands/sdk/conversation/impl/local_conversation.py +++ b/openhands-sdk/openhands/sdk/conversation/impl/local_conversation.py @@ -38,7 +38,7 @@ ) from openhands.sdk.event.conversation_error import ConversationErrorEvent from openhands.sdk.hooks import HookConfig, HookEventProcessor, create_hook_callback -from openhands.sdk.io import LocalFileStore +from openhands.sdk.io import FileStore, LocalFileStore from openhands.sdk.llm import LLM, Message, TextContent from openhands.sdk.llm.llm_profile_store import LLMProfileStore from openhands.sdk.llm.llm_registry import LLMRegistry @@ -109,6 +109,7 @@ def __init__( delete_on_close: bool = True, cipher: Cipher | None = None, tags: dict[str, str] | None = None, + file_store: FileStore | None = None, **_: object, ): """Initialize the conversation. @@ -124,7 +125,8 @@ def __init__( semantics: skills override by name (last wins), MCP config override by key (last wins), hooks concatenate (all run). persistence_dir: Directory for persisting conversation state and events. - Can be a string path or Path object. + Can be a string path or Path object. When file_store is provided, + this value is still used to derive environment observation paths. conversation_id: Optional ID for the conversation. If provided, will be used to identify the conversation. The user might want to suffix their persistent filestore with this ID. @@ -150,6 +152,9 @@ def __init__( (lost) on serialization. tags: Optional key-value tags for the conversation. Keys must be lowercase alphanumeric, values up to 256 characters. + file_store: Optional FileStore to use for conversation state and EventLog + persistence. If provided, this takes precedence over persistence_dir + for state and EventLog storage. """ super().__init__() # Initialize with span tracking # Mark cleanup as initiated as early as possible to avoid races or partially @@ -185,6 +190,7 @@ def __init__( persistence_dir=self.get_persistence_dir(persistence_dir, desired_id) if persistence_dir else None, + file_store=file_store, max_iterations=max_iteration_per_run, stuck_detection=stuck_detection, cipher=cipher, diff --git a/openhands-sdk/openhands/sdk/conversation/state.py b/openhands-sdk/openhands/sdk/conversation/state.py index 7e137929f2..37377da8e9 100644 --- a/openhands-sdk/openhands/sdk/conversation/state.py +++ b/openhands-sdk/openhands/sdk/conversation/state.py @@ -284,6 +284,7 @@ def create( stuck_detection: bool = True, cipher: Cipher | None = None, tags: dict[str, str] | None = None, + file_store: FileStore | None = None, ) -> "ConversationState": """Create a new conversation state or resume from persistence. @@ -304,7 +305,10 @@ def create( id: Unique conversation identifier agent: The Agent to use (tools must match persisted on restore) workspace: Working directory for agent operations - persistence_dir: Directory for persisting state and events + persistence_dir: Directory for persisting state and events when + file_store is not provided. When file_store is provided, this + value is still stored on the state and used for environment + observation paths. max_iterations: Maximum iterations per run stuck_detection: Whether to enable stuck detection cipher: Optional cipher for encrypting/decrypting secrets in @@ -313,6 +317,9 @@ def create( are redacted (lost) on serialization. tags: Optional key-value tags for the conversation. Keys must be lowercase alphanumeric, values up to 256 characters. + file_store: Optional FileStore to use for state and EventLog + persistence. If provided, this takes precedence over + persistence_dir for state and EventLog storage. Returns: ConversationState ready for use @@ -321,16 +328,17 @@ def create( ValueError: If conversation ID or tools mismatch on restore ValidationError: If agent or other fields fail Pydantic validation """ - if persistence_dir: - file_store = LocalFileStore( - persistence_dir, cache_limit_size=max_iterations - ) - else: - logger.warning( - "No persistence_dir provided; falling back to InMemoryFileStore. " - "EventLog data will not persist across requests." - ) - file_store = InMemoryFileStore() + if file_store is None: + if persistence_dir: + file_store = LocalFileStore( + persistence_dir, cache_limit_size=max_iterations + ) + else: + logger.warning( + "No persistence_dir provided; falling back to InMemoryFileStore. " + "EventLog data will not persist across requests." + ) + file_store = InMemoryFileStore() try: base_text = file_store.read(BASE_STATE) diff --git a/tests/sdk/conversation/local/test_custom_file_store.py b/tests/sdk/conversation/local/test_custom_file_store.py new file mode 100644 index 0000000000..c8810cb41a --- /dev/null +++ b/tests/sdk/conversation/local/test_custom_file_store.py @@ -0,0 +1,129 @@ +"""Tests for custom FileStore injection in local conversations.""" + +import logging +import uuid +from pathlib import Path + +from pydantic import SecretStr + +from openhands.sdk.agent import Agent +from openhands.sdk.conversation.impl.local_conversation import LocalConversation +from openhands.sdk.conversation.state import ConversationState +from openhands.sdk.io import InMemoryFileStore +from openhands.sdk.llm import LLM +from openhands.sdk.workspace import LocalWorkspace + + +def create_test_agent() -> Agent: + """Create a test agent.""" + llm = LLM(model="gpt-4o-mini", api_key=SecretStr("test-key"), usage_id="test-llm") + return Agent(llm=llm, tools=[]) + + +def test_conversation_state_uses_injected_file_store(tmp_path, caplog): + """ConversationState.create uses an injected FileStore without warning.""" + file_store = InMemoryFileStore() + conversation_id = uuid.uuid4() + workspace = LocalWorkspace(working_dir=tmp_path / "workspace") + + with caplog.at_level(logging.WARNING): + state = ConversationState.create( + id=conversation_id, + agent=create_test_agent(), + workspace=workspace, + file_store=file_store, + ) + + assert state.id == conversation_id + assert state.persistence_dir is None + assert state._fs is file_store + assert file_store.exists("base_state.json") + assert not any( + "No persistence_dir provided; falling back to InMemoryFileStore" + in record.message + for record in caplog.records + ) + + resumed_state = ConversationState.create( + id=conversation_id, + agent=create_test_agent(), + workspace=workspace, + file_store=file_store, + ) + + assert resumed_state.id == conversation_id + assert resumed_state._fs is file_store + + +def test_conversation_state_file_store_takes_precedence_over_persistence_dir(tmp_path): + """Injected FileStore stores state while persistence_dir remains metadata.""" + file_store = InMemoryFileStore() + persistence_dir = tmp_path / "persistence" + + state = ConversationState.create( + id=uuid.uuid4(), + agent=create_test_agent(), + workspace=LocalWorkspace(working_dir=tmp_path / "workspace"), + persistence_dir=str(persistence_dir), + file_store=file_store, + ) + + assert state.persistence_dir == str(persistence_dir) + assert state._fs is file_store + assert file_store.exists("base_state.json") + assert not (persistence_dir / "base_state.json").exists() + + +def test_local_conversation_uses_injected_file_store(tmp_path): + """LocalConversation forwards the injected FileStore to ConversationState.""" + file_store = InMemoryFileStore() + conversation_id = uuid.uuid4() + + conversation = LocalConversation( + agent=create_test_agent(), + workspace=tmp_path / "workspace", + conversation_id=conversation_id, + file_store=file_store, + visualizer=None, + ) + + assert conversation.id == conversation_id + assert conversation.state.persistence_dir is None + assert conversation.state._fs is file_store + assert file_store.exists("base_state.json") + + resumed_conversation = LocalConversation( + agent=create_test_agent(), + workspace=tmp_path / "workspace", + conversation_id=conversation_id, + file_store=file_store, + visualizer=None, + ) + + assert resumed_conversation.id == conversation_id + assert resumed_conversation.state._fs is file_store + + +def test_local_conversation_keeps_persistence_dir_with_injected_file_store(tmp_path): + """persistence_dir still sets observation paths when FileStore is injected.""" + file_store = InMemoryFileStore() + persistence_root = tmp_path / "persistence" + conversation_id = uuid.uuid4() + + conversation = LocalConversation( + agent=create_test_agent(), + workspace=tmp_path / "workspace", + persistence_dir=persistence_root, + conversation_id=conversation_id, + file_store=file_store, + visualizer=None, + ) + + expected_persistence_dir = Path(persistence_root) / conversation_id.hex + assert conversation.state.persistence_dir == str(expected_persistence_dir) + assert conversation.state.env_observation_persistence_dir == str( + expected_persistence_dir / "observations" + ) + assert conversation.state._fs is file_store + assert file_store.exists("base_state.json") + assert not (expected_persistence_dir / "base_state.json").exists()