Skip to content

Commit d9f3a8a

Browse files
xzrderekxzrderek
authored
Merge pull request #10 from reward-protocol/derekx/refactor-session
Session Management Consolidation
2 parents 661f950 + 3c62324 commit d9f3a8a

1 file changed

Lines changed: 178 additions & 40 deletions

File tree

‎reward_kit/mcp/mcpgym.py‎

Lines changed: 178 additions & 40 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,8 @@
1212
- Session-aware control plane endpoints via @control_plane_endpoint decorator
1313
"""
1414

15+
import hashlib
16+
import threading
1517
import inspect
1618
import json
1719
import logging
@@ -84,6 +86,12 @@ def __init__(
8486
"""
8587
super().__init__(server_name, adapter)
8688

89+
# Multi-session support
90+
self.sessions = (
91+
{}
92+
) # session_id -> {"env": env, "obs": obs, "session_data": data}
93+
self.session_lock = threading.Lock()
94+
8795
# Control plane endpoints dictionary
8896
self._control_plane_endpoints: Dict[str, Callable] = {}
8997

@@ -103,6 +111,136 @@ def __init__(
103111
# Discover and register control plane endpoints
104112
self._discover_and_register_control_plane_endpoints()
105113

114+
def _get_session_id(self, ctx: Context) -> str:
115+
"""
116+
Extract session ID from MCP context using proper FastMCP pattern.
117+
118+
Creates stable session IDs based on client info (seed + config + client details)
119+
for consistent session management across reconnections.
120+
"""
121+
print(f"🔍 _get_session_id: Starting session ID extraction")
122+
print(f"🔍 _get_session_id: ctx type: {type(ctx)}")
123+
print(f"🔍 _get_session_id: hasattr(ctx, 'session'): {hasattr(ctx, 'session')}")
124+
125+
# Use stable session ID based on client info (following simulation_server.py pattern)
126+
if hasattr(ctx, "session") and hasattr(ctx.session, "client_params"):
127+
client_params = ctx.session.client_params
128+
print(f"🔍 _get_session_id: client_params type: {type(client_params)}")
129+
print(
130+
f"🔍 _get_session_id: hasattr(client_params, 'clientInfo'): {hasattr(client_params, 'clientInfo')}"
131+
)
132+
133+
if hasattr(client_params, "clientInfo"):
134+
client_info = client_params.clientInfo
135+
print(f"🔍 _get_session_id: client_info: {client_info}")
136+
print(
137+
f"🔍 _get_session_id: hasattr(client_info, '_extra'): {hasattr(client_info, '_extra')}"
138+
)
139+
140+
if client_info and hasattr(client_info, "_extra"):
141+
extra_data = client_info._extra
142+
print(f"🔍 _get_session_id: extra_data: {extra_data}")
143+
print(f"🔍 _get_session_id: extra_data type: {type(extra_data)}")
144+
145+
if extra_data and isinstance(extra_data, dict):
146+
# Create a stable session ID based on seed and other config
147+
seed_value = extra_data.get("seed")
148+
config_value = extra_data.get("config", {})
149+
150+
print(
151+
f"🔍 _get_session_id: seed_value: {seed_value} (type: {type(seed_value)})"
152+
)
153+
print(f"🔍 _get_session_id: config_value: {config_value}")
154+
155+
stable_data = {
156+
"seed": seed_value,
157+
"config": config_value,
158+
"name": client_info.name,
159+
"version": client_info.version,
160+
}
161+
162+
print(f"🔍 _get_session_id: stable_data: {stable_data}")
163+
stable_str = json.dumps(stable_data, sort_keys=True)
164+
session_id = hashlib.md5(stable_str.encode()).hexdigest()
165+
print(
166+
f"🎯 Generated stable session_id: {session_id} for seed: {seed_value}"
167+
)
168+
return session_id
169+
170+
# Fallback for testing or other scenarios
171+
session_id = f"gym_{id(ctx)}"
172+
print(f"🎯 Generated fallback session_id: {session_id}")
173+
return session_id
174+
175+
def _get_or_create_session(self, ctx: Context) -> Dict[str, Any]:
176+
"""
177+
Get or create session data for the given context.
178+
179+
This method handles comprehensive session creation with seed extraction
180+
from MCP context and proper environment initialization.
181+
"""
182+
session_id = self._get_session_id(ctx)
183+
print(f"🔍 _get_or_create_session: session_id: {session_id}")
184+
185+
with self.session_lock:
186+
if session_id not in self.sessions:
187+
print(
188+
f"🔍 _get_or_create_session: Creating new session for {session_id}"
189+
)
190+
# Extract seed from context using proper FastMCP pattern
191+
seed = None
192+
config = self._get_default_config()
193+
print(f"🔍 _get_or_create_session: default_config: {config}")
194+
195+
if hasattr(ctx, "session") and hasattr(ctx.session, "client_params"):
196+
client_params = ctx.session.client_params
197+
if hasattr(client_params, "clientInfo"):
198+
client_info = client_params.clientInfo
199+
if client_info and hasattr(client_info, "_extra"):
200+
extra_data = client_info._extra
201+
print(
202+
f"🔍 _get_or_create_session: extra_data in session creation: {extra_data}"
203+
)
204+
if extra_data and isinstance(extra_data, dict):
205+
# Extract seed from client info
206+
seed = extra_data.get("seed")
207+
print(
208+
f"🌱 Extracted seed from client_info: {seed} (type: {type(seed)})"
209+
)
210+
# Update config with any additional options
211+
if "config" in extra_data:
212+
config.update(extra_data["config"])
213+
print(
214+
f"🔍 _get_or_create_session: updated config: {config}"
215+
)
216+
217+
print(
218+
f"🔍 _get_or_create_session: About to create environment with seed: {seed}"
219+
)
220+
221+
env, obs, info = self._new_env(seed=seed)
222+
print(
223+
f"🔍 _get_or_create_session: environment created with obs: {obs}, info: {info}"
224+
)
225+
226+
# Initialize session state
227+
self.sessions[session_id] = {
228+
"env": env,
229+
"obs": obs,
230+
"session_data": {}, # Subclasses can store additional data here
231+
"session_id": session_id,
232+
}
233+
234+
print(
235+
f"🎮 Created new session {session_id[:16]}... with seed {seed}, initial obs: {obs}"
236+
)
237+
else:
238+
print(
239+
f"🔍 _get_or_create_session: Returning existing session {session_id}"
240+
)
241+
242+
return self.sessions[session_id]
243+
106244
def _discover_and_register_control_plane_endpoints(self):
107245
"""
108246
Discover and register control plane endpoints on the subclass instance.
@@ -140,10 +278,16 @@ async def endpoint_handler(request: Request) -> JSONResponse:
140278
# For initial state endpoint, we need to create the session
141279
# based on the session ID and available information
142280
if func.__name__ == "get_initial_state_endpoint":
143-
# Create session with extracted seed from session ID
144-
session_data = self._create_session_from_id(
145-
session_id
146-
)
281+
env, obs, info = self._new_env(seed=None)
282+
# Initialize session state with extracted seed from session ID
283+
session_data = {
284+
"env": env,
285+
"obs": obs,
286+
"session_data": {}, # Subclasses can store additional data here
287+
"session_id": session_id,
288+
}
289+
# Store the session
290+
self.sessions[session_id] = session_data
147291
else:
148292
return JSONResponse(
149293
{"error": f"Session {session_id} not found"},
@@ -176,42 +320,6 @@ async def endpoint_handler(request: Request) -> JSONResponse:
176320
else:
177321
logger.info("⚠️ No session-aware control plane endpoints discovered")
178322

179-
def _create_session_from_id(self, session_id: str) -> Dict[str, Any]:
180-
"""
181-
Create a session based on session ID when the initial state endpoint is called.
182-
183-
The session ID is a hash of seed + config, so we can't extract the original values.
184-
Instead, we'll create a session with default values and let the tool calls handle
185-
the proper seed extraction from the MCP context.
186-
187-
Args:
188-
session_id: Session ID from the client
189-
190-
Returns:
191-
Session data dictionary
192-
"""
193-
# Create environment with default settings
194-
# The proper seed will be applied when the first tool is called
195-
config = self.adapter.get_default_config()
196-
197-
# Create environment without seed initially
198-
env = self.adapter.create_environment(config)
199-
obs, info = self.adapter.reset_environment(env, seed=None)
200-
201-
# Initialize session state
202-
session_data = {
203-
"env": env,
204-
"obs": obs,
205-
"session_data": {}, # Subclasses can store additional data here
206-
"session_id": session_id,
207-
}
208-
209-
# Store the session
210-
self.sessions[session_id] = session_data
211-
212-
print(f"🎮 Created session {session_id[:16]}... for initial state endpoint")
213-
214-
return session_data
215323

216324
def _update_control_plane(
217325
self, reward: float, terminated: bool, truncated: bool, info: Dict[str, Any]
@@ -342,6 +450,36 @@ def _execute_session_environment_step(
342450
# Return ONLY data plane information (no rewards/termination)
343451
return self.format_observation(obs, env)
344452

453+
454+
def _new_env(self, seed: Optional[int] = None) -> Tuple[Any, Any, Dict]:
455+
"""Create new environment and return initial state."""
456+
config = self.adapter.get_default_config()
457+
458+
try:
459+
env, obs, info = self.adapter.create_environment_with_seed(config, seed=seed)
460+
except AttributeError:
461+
env = self.adapter.create_environment(config)
462+
obs, info = self.adapter.reset_environment(env, seed=seed)
463+
464+
return env, obs, info
465+
466+
def _render(self, obs) -> Dict[str, Any]:
467+
"""Format observation using subclass implementation."""
468+
return self.format_observation(obs, self.env)
469+
470+
471+
def _get_default_config(self) -> Dict[str, Any]:
472+
"""
473+
Get default configuration from adapter.
474+
475+
Wrapper method to handle potential adapter interface issues.
476+
"""
477+
try:
478+
return self.adapter.get_default_config()
479+
except AttributeError:
480+
# Fallback for adapters that don't implement get_default_config
481+
return {}
482+
345483
# ===== SESSION-AWARE CONTROL PLANE ENDPOINTS =====
346484
# These provide session-specific control plane data via HTTP endpoints
347485
# instead of global MCP resources, enabling proper multi-session support.

0 commit comments

Comments
 (0)