feat(sdk): allow custom FileStore injection

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,

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)

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()