From f9297bf953f7a07037157d5249b153d348ddb106 Mon Sep 17 00:00:00 2001
From: Son Tran <sontran@google.com>
Date: Thu, 17 Sep 2026 14:38:50 -0700
Subject: [PATCH] Decouple Trajectory Store protocols and store backends from
 Tunix-specific abstractions (TunixTrajectoryMetadata, TunixTrajectory,
 TunixAgentStep, and TunixEnvStep): - Parameterize TrajectoryReader,
 TrajectoryWriter, and TrajectoryStore using TypeVar('T',
 bound=TrajectoryMetadata) and TypeVar('TrajT', bound=TrajectoryMetadata). -
 Add a polymorphic `create_trajectory` factory method on TrajectoryMetadata
 and TunixTrajectoryMetadata to construct full Trajectory instances
 polymorphically. - Remove Tunix-specific imports, `isinstance` branches, and
 static type ignores from InMemoryTrajectoryStore and FileTrajectoryStore. -
 In FileTrajectoryStore, load step data directly as dictionaries and pass them
 to the trajectory factory, allowing Pydantic's recursive field validation to
 automatically instantiate the concrete step models (e.g., TunixAgentStep vs.
 TunixEnvStep) without requiring manual step type configurations. - Add test
 cases in in_memory_store_test.py and file_store_test.py covering step_id=0
 Tunix trajectories and custom TrajectoryMetadata/Trajectory subclasses.

PiperOrigin-RevId: 983429942
---
 .../trajectory/file_store_test.py             | 82 ++++++++++++++++++-
 .../trajectory/in_memory_store_test.py        | 49 ++++++++++-
 tests/experimental/trajectory/store_test.py   | 62 ++++++++++++++
 tunix/experimental/trajectory/file_store.py   | 65 +++++++++------
 .../trajectory/in_memory_store.py             | 45 +++++-----
 tunix/experimental/trajectory/store.py        | 31 ++++---
 tunix/experimental/trajectory/trajectory.py   | 26 +++++-
 7 files changed, 297 insertions(+), 63 deletions(-)

diff --git a/tests/experimental/trajectory/file_store_test.py b/tests/experimental/trajectory/file_store_test.py
index 552ebbec15..095384af80 100644
--- a/tests/experimental/trajectory/file_store_test.py
+++ b/tests/experimental/trajectory/file_store_test.py
@@ -5,6 +5,7 @@
 from absl.testing import absltest
 from absl.testing import parameterized
 from etils import epath
+import pydantic
 from tunix.experimental.trajectory import file_store
 from tunix.experimental.trajectory import store
 from tunix.experimental.trajectory import store_testing
@@ -12,6 +13,14 @@
 from tunix.experimental.trajectory import trajectory_testing
 
 
+class _CustomMetadata(trajectory_lib.TrajectoryMetadata):
+  custom_tag: str = ""
+
+
+class _CustomTrajectory(_CustomMetadata):
+  steps: list[trajectory_lib.Step] = pydantic.Field(default_factory=list)
+
+
 class FileTrajectoryReaderTest(store_testing.TrajectoryReaderTestCase):
   """Contract tests for FileTrajectoryStore's TrajectoryReader implementation."""
 
@@ -199,7 +208,8 @@ def blocking_process_task(task):
     with mock.patch.object(
         self.file_s._writer, "_process_task", side_effect=blocking_process_task
     ):
-      # add_step should enqueue task and return immediately while worker loop is blocked.
+      # add_step should enqueue task and return immediately while worker loop
+      # is blocked.
       self.file_s.add_step(
           trajectory_testing.STEP_1_1, trajectory_testing.METADATA_1
       )
@@ -471,6 +481,76 @@ def test_get_trajectories_metadata_nonexistent_root_dir_returns_empty(
     store_instance = file_store.FileTrajectoryStore(root_dir=nonexistent_root)
     self.assertEmpty(store_instance.get_trajectories_metadata())
 
+  def test_tunix_trajectory_with_step_zero(self) -> None:
+    """Verifies storing and retrieving TunixTrajectoryMetadata and TunixTrajectory with step_id=0."""
+    tunix_store: file_store.FileTrajectoryStore[
+        trajectory_lib.TunixTrajectoryMetadata, trajectory_lib.TunixTrajectory
+    ] = file_store.FileTrajectoryStore(
+        root_dir=self.tmp_dir,
+        run_id="tunix_run",
+        metadata_cls=trajectory_lib.TunixTrajectoryMetadata,
+    )
+    meta = trajectory_lib.TunixTrajectoryMetadata(
+        trajectory_id="tunix_1",
+        agent=trajectory_lib.Agent(name="a1", version="1.0"),
+        status="RUNNING",
+    )
+    step0 = trajectory_lib.TunixEnvStep(
+        step_id=0, source=trajectory_lib.Source.USER, message="prompt"
+    )
+    step1 = trajectory_lib.TunixAgentStep(
+        step_id=1, source=trajectory_lib.Source.AGENT, message="response"
+    )
+    tunix_store.add_step(step0, meta)
+    tunix_store.add_step(step1, meta)
+    tunix_store.flush()
+
+    metas = tunix_store.get_trajectories_metadata(["tunix_1"])
+    self.assertLen(metas, 1)
+    self.assertIsInstance(metas[0], trajectory_lib.TunixTrajectoryMetadata)
+    self.assertEqual(metas[0].status, "RUNNING")
+
+    trajs = tunix_store.get_trajectories(["tunix_1"])
+    self.assertLen(trajs, 1)
+    self.assertIsInstance(trajs[0], trajectory_lib.TunixTrajectory)
+    self.assertEqual(trajs[0].steps[0].step_id, 0)
+    self.assertEqual(trajs[0].steps[1].step_id, 1)
+    self.assertIsInstance(trajs[0].steps[0], trajectory_lib.TunixEnvStep)
+    self.assertIsInstance(trajs[0].steps[1], trajectory_lib.TunixAgentStep)
+
+  def test_custom_metadata_and_trajectory_subclass(self) -> None:
+    """Verifies FileTrajectoryStore supports custom metadata and trajectory types."""
+    custom_store: file_store.FileTrajectoryStore[
+        _CustomMetadata, _CustomTrajectory
+    ] = file_store.FileTrajectoryStore(
+        root_dir=self.tmp_dir,
+        run_id="custom_run",
+        metadata_cls=_CustomMetadata,
+        trajectory_cls=_CustomTrajectory,
+    )
+    meta = _CustomMetadata(
+        trajectory_id="custom_1",
+        agent=trajectory_lib.Agent(name="custom_agent", version="1.0"),
+        custom_tag="experiment_42",
+    )
+    step = trajectory_lib.Step(
+        step_id=1, source=trajectory_lib.Source.AGENT, message="custom step"
+    )
+    custom_store.add_step(step, meta)
+    custom_store.flush()
+
+    metas = custom_store.get_trajectories_metadata(["custom_1"])
+    self.assertLen(metas, 1)
+    self.assertIsInstance(metas[0], _CustomMetadata)
+    self.assertEqual(metas[0].custom_tag, "experiment_42")
+
+    trajs = custom_store.get_trajectories(["custom_1"])
+    self.assertLen(trajs, 1)
+    self.assertIsInstance(trajs[0], _CustomTrajectory)
+    self.assertEqual(trajs[0].custom_tag, "experiment_42")
+    self.assertLen(trajs[0].steps, 1)
+    self.assertEqual(trajs[0].steps[0].message, "custom step")
+
 
 if __name__ == "__main__":
   absltest.main()
diff --git a/tests/experimental/trajectory/in_memory_store_test.py b/tests/experimental/trajectory/in_memory_store_test.py
index 65b4a2ef67..8346024bc8 100644
--- a/tests/experimental/trajectory/in_memory_store_test.py
+++ b/tests/experimental/trajectory/in_memory_store_test.py
@@ -1,10 +1,19 @@
 from absl.testing import absltest
+import pydantic
 from tunix.experimental.trajectory import in_memory_store
 from tunix.experimental.trajectory import store
 from tunix.experimental.trajectory import store_testing
 from tunix.experimental.trajectory import trajectory as trajectory_lib
 
 
+class _CustomMetadata(trajectory_lib.TrajectoryMetadata):
+  custom_tag: str = ""
+
+
+class _CustomTrajectory(_CustomMetadata):
+  steps: list[trajectory_lib.Step] = pydantic.Field(default_factory=list)
+
+
 class InMemoryTrajectoryReaderTest(store_testing.TrajectoryReaderTestCase):
   """Contract tests for InMemoryTrajectoryStore's TrajectoryReader implementation."""
 
@@ -37,7 +46,8 @@ def _create_reader_and_writer(
     mem_store = in_memory_store.InMemoryTrajectoryStore()
     return mem_store, mem_store
 
-  def test_update_metadata(self):
+  def test_update_metadata(self) -> None:
+    """Verifies that updating metadata in-memory updates the stored metadata."""
     mem_store = in_memory_store.InMemoryTrajectoryStore()
     meta = trajectory_lib.TrajectoryMetadata(
         trajectory_id="t1",
@@ -53,7 +63,8 @@ def test_update_metadata(self):
     read_meta = mem_store.get_trajectories_metadata()[0]
     self.assertEqual(read_meta.extra["status"], "SUCCEEDED")
 
-  def test_tunix_trajectory_with_step_zero(self):
+  def test_tunix_trajectory_with_step_zero(self) -> None:
+    """Verifies storing and retrieving TunixTrajectoryMetadata and TunixTrajectory with step_id=0."""
     mem_store = in_memory_store.InMemoryTrajectoryStore()
     meta = trajectory_lib.TunixTrajectoryMetadata(
         trajectory_id="tunix_1",
@@ -73,8 +84,40 @@ def test_tunix_trajectory_with_step_zero(self):
     self.assertIsInstance(trajs[0], trajectory_lib.TunixTrajectory)
     self.assertEqual(trajs[0].steps[0].step_id, 0)
     self.assertEqual(trajs[0].steps[1].step_id, 1)
+    self.assertIsInstance(trajs[0].steps[0], trajectory_lib.TunixEnvStep)
+    self.assertIsInstance(trajs[0].steps[1], trajectory_lib.TunixAgentStep)
+
+  def test_custom_metadata_and_trajectory_subclass(self) -> None:
+    """Verifies InMemoryTrajectoryStore supports custom metadata and trajectory types."""
+    custom_store: in_memory_store.InMemoryTrajectoryStore[
+        _CustomMetadata, _CustomTrajectory
+    ] = in_memory_store.InMemoryTrajectoryStore(
+        trajectory_cls=_CustomTrajectory
+    )
+    meta = _CustomMetadata(
+        trajectory_id="custom_1",
+        agent=trajectory_lib.Agent(name="custom_agent", version="1.0"),
+        custom_tag="experiment_42",
+    )
+    step = trajectory_lib.Step(
+        step_id=1, source=trajectory_lib.Source.AGENT, message="custom step"
+    )
+    custom_store.add_step(step, meta)
+
+    metas = custom_store.get_trajectories_metadata(["custom_1"])
+    self.assertLen(metas, 1)
+    self.assertIsInstance(metas[0], _CustomMetadata)
+    self.assertEqual(metas[0].custom_tag, "experiment_42")
+
+    trajs = custom_store.get_trajectories(["custom_1"])
+    self.assertLen(trajs, 1)
+    self.assertIsInstance(trajs[0], _CustomTrajectory)
+    self.assertEqual(trajs[0].custom_tag, "experiment_42")
+    self.assertLen(trajs[0].steps, 1)
+    self.assertEqual(trajs[0].steps[0].message, "custom step")
 
-  def test_metadata_mutation_isolation(self):
+  def test_metadata_mutation_isolation(self) -> None:
+    """Verifies that mutating returned metadata does not alter internal store state."""
     mem_store = in_memory_store.InMemoryTrajectoryStore()
     meta = trajectory_lib.TrajectoryMetadata(
         trajectory_id="iso_1",
diff --git a/tests/experimental/trajectory/store_test.py b/tests/experimental/trajectory/store_test.py
index fe67aa56be..50c8a87e11 100644
--- a/tests/experimental/trajectory/store_test.py
+++ b/tests/experimental/trajectory/store_test.py
@@ -21,6 +21,7 @@
 from tunix.experimental.trajectory import file_store
 from tunix.experimental.trajectory import in_memory_store
 from tunix.experimental.trajectory import store as store_lib
+from tunix.experimental.trajectory import trajectory as trajectory_lib
 from tunix.experimental.trajectory import trajectory_testing
 
 
@@ -192,5 +193,66 @@ def test_file_store_without_run_id_raises_on_to_config(self):
     store.close()
 
 
+class GenericTypeParametersTest(absltest.TestCase):
+  """Tests that TrajectoryStore and backends preserve generic type parameters without Generic."""
+
+  def test_classes_inherit_parameters_without_explicit_generic(self) -> None:
+    """Verifies that Python typing automatically discovers (T, TrajT)."""
+    self.assertLen(store_lib.TrajectoryStore.__parameters__, 2)
+    self.assertLen(in_memory_store.InMemoryTrajectoryStore.__parameters__, 2)
+    self.assertLen(file_store.FileTrajectoryStore.__parameters__, 2)
+
+    self.assertEqual(
+        store_lib.TrajectoryStore.__parameters__,
+        (store_lib.T, store_lib.TrajT),
+    )
+
+  def test_classes_are_runtime_subscriptable(self) -> None:
+    """Verifies that classes are subscriptable with concrete models at runtime."""
+    subscripted_store = store_lib.TrajectoryStore[
+        trajectory_lib.TunixTrajectoryMetadata, trajectory_lib.TunixTrajectory
+    ]
+    subscripted_mem = in_memory_store.InMemoryTrajectoryStore[
+        trajectory_lib.TunixTrajectoryMetadata, trajectory_lib.TunixTrajectory
+    ]
+    subscripted_file = file_store.FileTrajectoryStore[
+        trajectory_lib.TunixTrajectoryMetadata, trajectory_lib.TunixTrajectory
+    ]
+    self.assertIsNotNone(subscripted_store)
+    self.assertIsNotNone(subscripted_mem)
+    self.assertIsNotNone(subscripted_file)
+
+  def test_instances_satisfy_protocols_and_abc(self) -> None:
+    """Verifies protocol and ABC conformance on instantiated instances."""
+    mem = in_memory_store.InMemoryTrajectoryStore()
+    self.assertIsInstance(mem, store_lib.TrajectoryStore)
+    self.assertIsInstance(mem, store_lib.TrajectoryReader)
+    self.assertIsInstance(mem, store_lib.TrajectoryWriter)
+
+    tmp_dir = self.create_tempdir().full_path
+    f_store = file_store.FileTrajectoryStore(root_dir=tmp_dir)
+    self.assertIsInstance(f_store, store_lib.TrajectoryStore)
+    self.assertIsInstance(f_store, store_lib.TrajectoryReader)
+    self.assertIsInstance(f_store, store_lib.TrajectoryWriter)
+    f_store.close()
+
+  def test_subclassing_closes_type_parameters(self) -> None:
+    """Verifies that concrete subclassing binds and closes type parameters."""
+
+    class ConcreteFileStore(
+        file_store.FileTrajectoryStore[
+            trajectory_lib.TunixTrajectoryMetadata,
+            trajectory_lib.TunixTrajectory,
+        ]
+    ):
+      pass
+
+    self.assertEmpty(ConcreteFileStore.__parameters__)
+    tmp_dir = self.create_tempdir().full_path
+    c_store = ConcreteFileStore(root_dir=tmp_dir)
+    self.assertIsInstance(c_store, store_lib.TrajectoryStore)
+    c_store.close()
+
+
 if __name__ == "__main__":
   absltest.main()
diff --git a/tunix/experimental/trajectory/file_store.py b/tunix/experimental/trajectory/file_store.py
index 2fb8b89456..74736395a5 100644
--- a/tunix/experimental/trajectory/file_store.py
+++ b/tunix/experimental/trajectory/file_store.py
@@ -1,15 +1,19 @@
 """File-based implementation for Trajectory Store."""
 
 import functools
+import json
 import re
 import types
-from typing import Any, ClassVar, Final, Mapping
+from typing import Any, ClassVar, Final, Mapping, TypeVar, cast
 
 from etils import epath
 from tunix.experimental.trajectory import async_writer
 from tunix.experimental.trajectory import store
 from tunix.experimental.trajectory import trajectory as trajectory_lib
 
+T = TypeVar("T", bound=trajectory_lib.TrajectoryMetadata)
+TrajT = TypeVar("TrajT", bound=trajectory_lib.TrajectoryMetadata)
+
 _METADATA_FILENAME: Final[str] = "metadata.json"
 _TRAJECTORY_DIR_PREFIX: Final[str] = "traj_"
 # Characters allowed in a trajectory_id: ASCII letters, digits, underscores, and
@@ -50,9 +54,7 @@ def _validate_trajectory_id(trajectory_id: str | None) -> str:
   return trajectory_id
 
 
-class FileTrajectoryStore(
-    store.TrajectoryStore, store.TrajectoryReader, store.TrajectoryWriter
-):
+class FileTrajectoryStore(store.TrajectoryStore[T, TrajT]):
   """File-based implementation satisfying TrajectoryReader and TrajectoryWriter.
 
   Architectural Separation of Responsibilities:
@@ -79,7 +81,12 @@ class FileTrajectoryStore(
   BACKEND: ClassVar[str] = "file"
 
   def __init__(
-      self, root_dir: epath.PathLike, run_id: str | None = None
+      self,
+      root_dir: epath.PathLike,
+      run_id: str | None = None,
+      *,
+      metadata_cls: type[T] = trajectory_lib.TrajectoryMetadata,
+      trajectory_cls: type[TrajT] | None = None,
   ) -> None:
     """Initializes FileTrajectoryStore.
 
@@ -90,6 +97,10 @@ def __init__(
         scoped under root_dir / run_id. This ID MUST stay the same when
         recovering from failures or process restarts as long as the same RL
         process is being continued.
+      metadata_cls: Type of TrajectoryMetadata to deserialize. Defaults to
+        TrajectoryMetadata.
+      trajectory_cls: Optional explicit Trajectory type to instantiate. If None,
+        calls `meta.create_trajectory(steps=steps)`.
 
     Raises:
       ValueError: If root_dir is empty, or run_id is given but cannot be used
@@ -109,10 +120,14 @@ def __init__(
       )
     self._raw_root_dir = epath.Path(root_dir)
     self._run_id = run_id
+    self._metadata_cls = metadata_cls
+    self._trajectory_cls = trajectory_cls
     self._writer = async_writer.AsyncFileWriter()
 
   @classmethod
-  def _from_config(cls, config: Mapping[str, Any]) -> "FileTrajectoryStore":
+  def _from_config(
+      cls, config: Mapping[str, Any]
+  ) -> "FileTrajectoryStore[Any, Any]":
     """Builds a file-backed store from `config`.
 
     Args:
@@ -186,7 +201,7 @@ def get_step_path(self, trajectory_id: str, step_id: int) -> epath.Path:
 
   def get_trajectories_metadata(
       self, trajectory_ids: list[str] | None = None
-  ) -> list[trajectory_lib.TrajectoryMetadata]:
+  ) -> list[T]:
     """Retrieves metadata for trajectories in the run.
 
     Args:
@@ -201,7 +216,7 @@ def get_trajectories_metadata(
       store.TrajectoryMetadataNotFoundError: If any requested trajectory ID does
         not exist.
     """
-    metas: list[trajectory_lib.TrajectoryMetadata] = []
+    metas: list[T] = []
     if trajectory_ids is None:
       if not self.root_dir.exists():
         return metas
@@ -217,16 +232,12 @@ def get_trajectories_metadata(
       meta_path = self.get_trajectory_metadata_path(traj_id)
       if not meta_path.exists():
         raise store.TrajectoryMetadataNotFoundError(traj_id)
-      meta = trajectory_lib.TrajectoryMetadata.model_validate_json(
-          meta_path.read_text()
-      )
+      meta = self._metadata_cls.model_validate_json(meta_path.read_text())
       metas.append(meta)
 
     return metas
 
-  def get_trajectories(
-      self, trajectory_ids: list[str]
-  ) -> list[trajectory_lib.Trajectory]:
+  def get_trajectories(self, trajectory_ids: list[str]) -> list[TrajT]:
     """Retrieves full trajectories for a list of trajectory IDs.
 
     Args:
@@ -239,7 +250,7 @@ def get_trajectories(
       store.TrajectoryNotFoundError: If any requested trajectory ID does not
       exist.
     """
-    trajs: list[trajectory_lib.Trajectory] = []
+    trajs: list[TrajT] = []
 
     for traj_id in trajectory_ids:
       traj_dir = self.get_trajectory_dir(traj_id)
@@ -247,27 +258,27 @@ def get_trajectories(
       if not meta_path.exists():
         raise store.TrajectoryNotFoundError(traj_id)
 
-      meta = trajectory_lib.TrajectoryMetadata.model_validate_json(
-          meta_path.read_text()
-      )
-      steps: list[trajectory_lib.Step] = []
+      meta = self._metadata_cls.model_validate_json(meta_path.read_text())
+      steps: list[Any] = []
 
       for file_entry in traj_dir.iterdir():
         if not _STEP_FILENAME_REGEX.match(file_entry.name):
           continue
-        step = trajectory_lib.Step.model_validate_json(file_entry.read_text())
-        steps.append(step)
+        steps.append(json.loads(file_entry.read_text()))
 
-      traj_data = meta.model_dump()
-      traj_data["steps"] = steps
-      trajs.append(trajectory_lib.Trajectory(**traj_data))
+      if self._trajectory_cls is not None:
+        traj_data = meta.model_dump()
+        traj_data["steps"] = steps
+        trajs.append(self._trajectory_cls(**traj_data))
+      else:
+        trajs.append(cast(TrajT, meta.create_trajectory(steps=steps)))
 
     return trajs
 
   def add_step(
       self,
       step: trajectory_lib.Step,
-      metadata: trajectory_lib.TrajectoryMetadata,
+      metadata: T,
   ) -> None:
     """Asynchronously logs a turn step and its trajectory metadata.
 
@@ -288,7 +299,7 @@ def add_step(
 
   def update_metadata(
       self,
-      metadata: trajectory_lib.TrajectoryMetadata,
+      metadata: T,
       step: trajectory_lib.Step | None = None,
   ) -> None:
     """Updates (or creates) trajectory metadata asynchronously, optionally writing a step.
@@ -340,7 +351,7 @@ def close(self) -> None:
     """
     self._writer.close()
 
-  def __enter__(self) -> "FileTrajectoryStore":
+  def __enter__(self) -> "FileTrajectoryStore[T, TrajT]":
     """Returns this store, for use as a context manager."""
     return self
 
diff --git a/tunix/experimental/trajectory/in_memory_store.py b/tunix/experimental/trajectory/in_memory_store.py
index a8317ca70b..5c3a063737 100644
--- a/tunix/experimental/trajectory/in_memory_store.py
+++ b/tunix/experimental/trajectory/in_memory_store.py
@@ -1,11 +1,14 @@
 """In-memory implementation for Trajectory Store."""
 
 import collections
-from typing import Any, ClassVar, Mapping
+from typing import Any, ClassVar, Mapping, TypeVar, cast
 
 from tunix.experimental.trajectory import store
 from tunix.experimental.trajectory import trajectory as trajectory_lib
 
+T = TypeVar("T", bound=trajectory_lib.TrajectoryMetadata)
+TrajT = TypeVar("TrajT", bound=trajectory_lib.TrajectoryMetadata)
+
 
 def _validate_trajectory_id(trajectory_id: str | None) -> str:
   """Validates that trajectory_id is non-empty.
@@ -24,9 +27,7 @@ def _validate_trajectory_id(trajectory_id: str | None) -> str:
   return trajectory_id
 
 
-class InMemoryTrajectoryStore(
-    store.TrajectoryStore, store.TrajectoryReader, store.TrajectoryWriter
-):
+class InMemoryTrajectoryStore(store.TrajectoryStore[T, TrajT]):
   """In-memory implementation satisfying TrajectoryReader and TrajectoryWriter.
 
   Process-local: the steps written here are visible only to the process that
@@ -36,17 +37,21 @@ class InMemoryTrajectoryStore(
 
   BACKEND: ClassVar[str] = "memory"
 
-  def __init__(self) -> None:
+  def __init__(
+      self,
+      trajectory_cls: type[TrajT] | None = None,
+  ) -> None:
     """Initializes the InMemoryTrajectoryStore."""
-    self._metadata_by_trajectory_id: dict[
-        str, trajectory_lib.TrajectoryMetadata
-    ] = {}
+    self._metadata_by_trajectory_id: dict[str, T] = {}
     self._steps_by_trajectory_id: dict[str, list[trajectory_lib.Step]] = (
         collections.defaultdict(list)
     )
+    self._trajectory_cls = trajectory_cls
 
   @classmethod
-  def _from_config(cls, config: Mapping[str, Any]) -> "InMemoryTrajectoryStore":
+  def _from_config(
+      cls, config: Mapping[str, Any]
+  ) -> "InMemoryTrajectoryStore[Any, Any]":
     """Builds an in-memory store; this backend takes no configuration.
 
     Args:
@@ -65,7 +70,7 @@ def to_config(self) -> dict[str, Any]:
 
   def get_trajectories_metadata(
       self, trajectory_ids: list[str] | None = None
-  ) -> list[trajectory_lib.TrajectoryMetadata]:
+  ) -> list[T]:
     """Retrieves metadata for trajectories in the run.
 
     Args:
@@ -82,7 +87,7 @@ def get_trajectories_metadata(
     """
     if trajectory_ids is None:
       trajectory_ids = list(self._metadata_by_trajectory_id.keys())
-    metas: list[trajectory_lib.TrajectoryMetadata] = []
+    metas: list[T] = []
     for traj_id in trajectory_ids:
       if traj_id not in self._metadata_by_trajectory_id:
         raise store.TrajectoryMetadataNotFoundError(traj_id)
@@ -92,7 +97,7 @@ def get_trajectories_metadata(
 
   def get_trajectories(
       self, trajectory_ids: list[str]
-  ) -> list[trajectory_lib.Trajectory]:
+  ) -> list[TrajT]:
     """Retrieves full trajectories for a list of trajectory IDs.
 
     Args:
@@ -105,7 +110,7 @@ def get_trajectories(
       store.TrajectoryNotFoundError: If any requested trajectory ID does not
       exist.
     """
-    result: list[trajectory_lib.Trajectory] = []
+    result: list[TrajT] = []
     for traj_id in trajectory_ids:
       if traj_id not in self._metadata_by_trajectory_id:
         raise store.TrajectoryNotFoundError(traj_id)
@@ -114,18 +119,18 @@ def get_trajectories(
           s.model_copy(deep=True)
           for s in self._steps_by_trajectory_id.get(traj_id, [])
       ]
-      traj_data = meta.model_dump()
-      traj_data["steps"] = steps
-      if isinstance(meta, trajectory_lib.TunixTrajectoryMetadata):
-        result.append(trajectory_lib.TunixTrajectory(**traj_data))  # pyrefly: ignore[bad-argument-type]
+      if self._trajectory_cls is not None:
+        traj_data = meta.model_dump()
+        traj_data["steps"] = steps
+        result.append(self._trajectory_cls(**traj_data))
       else:
-        result.append(trajectory_lib.Trajectory(**traj_data))
+        result.append(cast(TrajT, meta.create_trajectory(steps=steps)))
     return result
 
   def add_step(
       self,
       step: trajectory_lib.Step,
-      metadata: trajectory_lib.TrajectoryMetadata,
+      metadata: T,
   ) -> None:
     """Atomically logs a turn step and its trajectory metadata.
 
@@ -154,7 +159,7 @@ def add_step(
 
   def update_metadata(
       self,
-      metadata: trajectory_lib.TrajectoryMetadata,
+      metadata: T,
   ) -> None:
     """Updates (or creates) trajectory metadata.
 
diff --git a/tunix/experimental/trajectory/store.py b/tunix/experimental/trajectory/store.py
index 9d0cf0f772..b35dd9362a 100644
--- a/tunix/experimental/trajectory/store.py
+++ b/tunix/experimental/trajectory/store.py
@@ -1,10 +1,13 @@
 """Protocols defining Trajectory Store interfaces."""
 
 import abc
-from typing import Any, ClassVar, Mapping, Protocol, runtime_checkable
+from typing import Any, ClassVar, Mapping, Protocol, TypeVar, runtime_checkable
 
 from tunix.experimental.trajectory import trajectory as trajectory_lib
 
+T = TypeVar("T", bound=trajectory_lib.TrajectoryMetadata)
+TrajT = TypeVar("TrajT", bound=trajectory_lib.TrajectoryMetadata)
+
 # ==============================================================================
 # Custom Exceptions
 # ==============================================================================
@@ -32,12 +35,12 @@ def __init__(self, trajectory_id: str) -> None:
 
 
 @runtime_checkable
-class TrajectoryReader(Protocol):
+class TrajectoryReader(Protocol[T, TrajT]):
   """Structural protocol defining read-only Trajectory Store operations."""
 
   def get_trajectories_metadata(
       self, trajectory_ids: list[str] | None = None
-  ) -> list[trajectory_lib.TrajectoryMetadata]:
+  ) -> list[T]:
     """Retrieves metadata for trajectories in the run.
 
     Args:
@@ -56,7 +59,7 @@ def get_trajectories_metadata(
 
   def get_trajectories(
       self, trajectory_ids: list[str]
-  ) -> list[trajectory_lib.Trajectory]:
+  ) -> list[TrajT]:
     """Retrieves full trajectories for a list of trajectory IDs.
 
     Args:
@@ -72,13 +75,13 @@ def get_trajectories(
 
 
 @runtime_checkable
-class TrajectoryWriter(Protocol):
+class TrajectoryWriter(Protocol[T]):
   """Structural protocol defining write Trajectory Store operations."""
 
   def add_step(
       self,
       step: trajectory_lib.Step,
-      metadata: trajectory_lib.TrajectoryMetadata,
+      metadata: T,
   ) -> None:
     """Logs a turn step and its trajectory metadata.
 
@@ -97,7 +100,7 @@ def add_step(
 
   def update_metadata(
       self,
-      metadata: trajectory_lib.TrajectoryMetadata,
+      metadata: T,
   ) -> None:
     """Updates (or creates) trajectory metadata.
 
@@ -133,7 +136,11 @@ def close(self) -> None:
 # ==============================================================================
 
 
-class TrajectoryStore(TrajectoryReader, TrajectoryWriter, abc.ABC):
+class TrajectoryStore(
+    TrajectoryReader[T, TrajT],
+    TrajectoryWriter[T],
+    abc.ABC,
+):
   """Base class pairing a store implementation with its own configuration.
 
   Every backend owns both directions of its configuration: `_from_config`
@@ -151,7 +158,7 @@ class TrajectoryStore(TrajectoryReader, TrajectoryWriter, abc.ABC):
 
   # The value of the config's "backend" key that selects this class.
   BACKEND: ClassVar[str]
-  _REGISTRY: ClassVar[dict[str, type["TrajectoryStore"]]] = {}
+  _REGISTRY: ClassVar[dict[str, type["TrajectoryStore[Any, Any]"]]] = {}
 
   def __init_subclass__(cls, **kwargs: Any) -> None:
     super().__init_subclass__(**kwargs)
@@ -161,7 +168,9 @@ def __init_subclass__(cls, **kwargs: Any) -> None:
 
   @classmethod
   @abc.abstractmethod
-  def _from_config(cls, config: Mapping[str, Any]) -> "TrajectoryStore":
+  def _from_config(
+      cls, config: Mapping[str, Any]
+  ) -> "TrajectoryStore[Any, Any]":
     """Builds an instance of this backend from `config`.
 
     Implementations read the keys they care about and raise ValueError for a
@@ -192,7 +201,7 @@ def to_config(self) -> dict[str, Any]:
   @classmethod
   def from_config(
       cls, config: Mapping[str, Any] | None
-  ) -> "TrajectoryStore | None":
+  ) -> "TrajectoryStore[Any, Any] | None":
     """Builds the store described by `config`, or None when it is disabled.
 
     Call once per process and hold onto the result: the process that built a
diff --git a/tunix/experimental/trajectory/trajectory.py b/tunix/experimental/trajectory/trajectory.py
index dbdc6b71ce..9bffafca09 100644
--- a/tunix/experimental/trajectory/trajectory.py
+++ b/tunix/experimental/trajectory/trajectory.py
@@ -9,7 +9,7 @@
 import dataclasses
 import datetime
 import enum
-from typing import Annotated, Any, Final, Literal, get_args
+from typing import Annotated, Any, Final, Literal, Sequence, get_args
 
 import numpy as np
 import pydantic
@@ -391,6 +391,18 @@ class TrajectoryMetadata(pydantic.BaseModel):
       description="Custom root-level metadata.",
   )
 
+  def create_trajectory(
+      self,
+      steps: Sequence[Any] | None = None,
+      subagent_trajectories: Sequence[Any] | None = None,
+  ) -> Any:
+    """Creates a full Trajectory from this metadata and given steps."""
+    data = self.model_dump()
+    data["steps"] = list(steps) if steps is not None else []
+    if subagent_trajectories is not None:
+      data["subagent_trajectories"] = list(subagent_trajectories)
+    return Trajectory(**data)
+
 
 class Trajectory(TrajectoryMetadata):
   """Root trajectory object containing the interaction history."""
@@ -618,6 +630,18 @@ class TunixTrajectoryMetadata(TrajectoryMetadata):
       description="Timing information for reward operations.",
   )
 
+  def create_trajectory(
+      self,
+      steps: Sequence[Any] | None = None,
+      subagent_trajectories: Sequence[Any] | None = None,
+  ) -> Any:
+    """Creates a full TunixTrajectory from this metadata and given steps."""
+    data = self.model_dump()
+    data["steps"] = list(steps) if steps is not None else []
+    if subagent_trajectories is not None:
+      data["subagent_trajectories"] = list(subagent_trajectories)
+    return TunixTrajectory(**data)
+
 
 class TunixTrajectory(TunixTrajectoryMetadata):
   """Tunix-specific trajectory object containing the interaction history."""