diff --git a/src/openenv/cli/templates/openenv_env/client.py b/src/openenv/cli/templates/openenv_env/client.py
index 3ee1a0747..720090431 100644
--- a/src/openenv/cli/templates/openenv_env/client.py
+++ b/src/openenv/cli/templates/openenv_env/client.py
@@ -16,7 +16,7 @@
 
 
 class __ENV_CLASS_NAME__Env(
-    EnvClient[__ENV_CLASS_NAME__Action, __ENV_CLASS_NAME__Observation]
+    EnvClient[__ENV_CLASS_NAME__Action, __ENV_CLASS_NAME__Observation, State]
 ):
     """
     Client for the __ENV_TITLE_NAME__ Environment.
diff --git a/src/openenv/cli/templates/openenv_env/server/__ENV_NAME___environment.py b/src/openenv/cli/templates/openenv_env/server/__ENV_NAME___environment.py
index df640f89c..bbde58219 100644
--- a/src/openenv/cli/templates/openenv_env/server/__ENV_NAME___environment.py
+++ b/src/openenv/cli/templates/openenv_env/server/__ENV_NAME___environment.py
@@ -13,10 +13,14 @@
 
 from uuid import uuid4
 
-from models import __ENV_CLASS_NAME__Action, __ENV_CLASS_NAME__Observation
 from openenv.core.env_server.interfaces import Environment
 from openenv.core.env_server.types import State
 
+try:
+    from ..models import __ENV_CLASS_NAME__Action, __ENV_CLASS_NAME__Observation
+except ImportError:
+    from models import __ENV_CLASS_NAME__Action, __ENV_CLASS_NAME__Observation
+
 
 class __ENV_CLASS_NAME__Environment(Environment):
     """
diff --git a/src/openenv/cli/templates/openenv_env/server/app.py b/src/openenv/cli/templates/openenv_env/server/app.py
index 73d26cdec..898911a2a 100644
--- a/src/openenv/cli/templates/openenv_env/server/app.py
+++ b/src/openenv/cli/templates/openenv_env/server/app.py
@@ -35,10 +35,12 @@
         "openenv is required for the web interface. Install dependencies with '\n    uv sync\n'"
     ) from e
 
-# Import from local models.py (PYTHONPATH includes /app/env in Docker)
-from models import __ENV_CLASS_NAME__Action, __ENV_CLASS_NAME__Observation
-
-from .__ENV_NAME___environment import __ENV_CLASS_NAME__Environment
+try:
+    from ..models import __ENV_CLASS_NAME__Action, __ENV_CLASS_NAME__Observation
+    from .__ENV_NAME___environment import __ENV_CLASS_NAME__Environment
+except ModuleNotFoundError:
+    from models import __ENV_CLASS_NAME__Action, __ENV_CLASS_NAME__Observation
+    from server.__ENV_NAME___environment import __ENV_CLASS_NAME__Environment
 
 
 # Create the app with web interface and README integration