Use session lifespan context instead of global state

app.py

@@ -1,7 +1,10 @@
+from collections.abc import AsyncIterator
+from contextlib import asynccontextmanager
 from dataclasses import dataclass, field
 from typing import Any, Union
 
-from mcp.server.fastmcp import FastMCP
+from mcp.server.fastmcp import Context, FastMCP
+from mcp.server.session import ServerSession
 
 
 Observation = Union[str, dict[str, Any]]
@@ -60,14 +63,28 @@ class WordleEnv:
         return self._obs
 
 
-wordle = WordleEnv(secret="word")
+@dataclass
+class SessionContext:
+    """Session context with typed dependencies."""
+
+    wordle: WordleEnv
+
+
+@asynccontextmanager
+async def session_lifespan(server: FastMCP) -> AsyncIterator[SessionContext]:
+    """Manage session lifecycle with type-safe context."""
+    # Initialize on session initialization
+    wordle = WordleEnv(secret="word")
+    # try-finally if you need cleanup on session termination
+    yield SessionContext(wordle=wordle)
 
 
 # Stateful server (maintains session state)
-mcp = FastMCP("StatefulServer")
+mcp = FastMCP("StatefulServer", lifespan=session_lifespan)
+
 
 @mcp.tool()
-def step_fn(guess: str) -> tuple[str, float, bool, dict]:
+def step_fn(guess: str, ctx: Context[ServerSession, SessionContext]) -> tuple[str, float, bool, dict]:
     """
     Perform a step in the Wordle environment.
 
@@ -81,6 +98,7 @@ def step_fn(guess: str) -> tuple[str, float, bool, dict]:
             - done: Whether the game is done.
             - info: Additional info.
     """
+    wordle = ctx.request_context.lifespan_context.wordle
     result = wordle.step(guess)
     return result.observation, result.reward, result.done, result.info