import fix

client.py

@@ -3,7 +3,7 @@ Second Brain Environment — WebSocket Client
 Connects to the running server and exposes reset()/step()/state().
 """
 from openenv.core.env_client import EnvClient
-from models import SecondBrainAction, SecondBrainObservation
+from models import SecondBrainAction, SecondBrainObservation , SecondBrainState
 
 
 class SecondBrainEnv(EnvClient):
@@ -25,6 +25,7 @@ class SecondBrainEnv(EnvClient):
 
     action_type      = SecondBrainAction
     observation_type = SecondBrainObservation
+    state_type       = SecondBrainState  
 
     def __init__(self, base_url: str = "http://localhost:8000", task_name: str = "note_categorization"):
         super().__init__(base_url=base_url)

server/app.py

@@ -3,8 +3,7 @@ FastAPI server for Second Brain OpenEnv environment.
 Exposes reset(), step(), state() over HTTP + WebSocket.
 """
 import os
-from typing import Optional
-from fastapi import FastAPI, Query
+from fastapi import FastAPI
 from openenv.core.env_server import create_app
 
 from second_brain_env_environment import SecondBrainEnvironment