diff_server.txt

��diff --git a/frontend/interactive_server.py b/frontend/interactive_server.py
index ac03d50..346c12d 100644
--- a/frontend/interactive_server.py
+++ b/frontend/interactive_server.py
@@ -2,7 +2,9 @@ import sys
 import os
 import json
 import asyncio
-from llm_panda.metrics_logger import LANGSMITH_AVAILABLE, log_system_metrics, log_interaction
+import uuid
+import psutil
+from llm_panda.metrics_logger import LANGSMITH_AVAILABLE, log_system_metrics, log_interaction, log_episode_metrics
 
 import math
 import time
@@ -33,7 +35,7 @@ from llm_panda import config
 from llm_panda.execution.cancellation import ExecutionInterrupted
 from llm_panda import logger as _logger_setup
 _logger_setup.setup()
-from frontend.ctx.web_runtime import WebRuntime
+from frontend.web_runtime import WebRuntime
 import pybullet_data
 from pybullet_object_models import ycb_objects
 from fastapi.middleware.cors import CORSMiddleware
@@ -115,12 +117,13 @@ class SessionContext:
     def __init__(self, session_id: str):
         self.session_id = session_id
         self.sim = None
-        self.ctx.executor_lock = asyncio.Lock()
-        self.ctx.pybullet_lock = threading.RLock()
-        self.ctx.pipeline_executing = False
-        self.ctx.web_runtime = WebRuntime()
-        self.ctx.command_worker_task = None
-        self.ctx.physics_settle_task = None
+        self.executor_lock = asyncio.Lock()
+        self.pybullet_lock = threading.RLock()
+        self.pipeline_executing = False
+        self.web_runtime = WebRuntime()
+        self.command_worker_task = None
+        self.physics_settle_task = None
+        self.disconnected = False
         self.drag_state = {"constraint_id": None, "active": False}
         self.last_active = time.time()
 
@@ -345,12 +348,17 @@ def get_simulation_state_payload(ctx: SessionContext) -> dict:
         # Read objects pos/ori
         objects_state = {}
         for oid in sim.scene.object_ids:
-            name = sim.scene.get_name(oid)
-            pos, ori = pb.getBasePositionAndOrientation(oid, physicsClientId=sim.client)
-            objects_state[name] = {
-                "pos": list(pos),
-                "ori": list(ori)
-            }
+            if oid < 0:
+                continue
+            try:
+                name = sim.scene.get_name(oid)
+                pos, ori = pb.getBasePositionAndOrientation(oid, physicsClientId=sim.client)
+                objects_state[name] = {
+                    "pos": list(pos),
+                    "ori": list(ori)
+                }
+            except Exception as e:
+                print(f"Failed to get pose for object {oid}: {e}", flush=True)
         
         return {
             "type": "state",
@@ -477,8 +485,10 @@ def run_pipeline_threadsafe(
 
     def publish(payload: dict, journal: bool = True):
         payload = {**payload, "task_id": task_id, "generation": generation}
+        if getattr(ctx, "episode_id", None) is not None:
+            payload["run_id"] = ctx.episode_id
         if getattr(ctx, "episode_run_tree", None) is not None:
-            payload["run_id"] = str(ctx.episode_run_tree.id)
+            payload["langsmith_run_id"] = str(ctx.episode_run_tree.id)
         ctx.web_runtime.publish_threadsafe(payload, loop, journal=journal)
     
     def log(message: str, step: Optional[int] = None):
@@ -530,7 +540,7 @@ def run_pipeline_threadsafe(
         curr_time = time.time()
         if curr_time - last_stream_time[0] >= throttle_interval:
             last_stream_time[0] = curr_time
-            state_payload = get_simulation_state_payload(current_sim)
+            state_payload = get_simulation_state_payload(ctx)
             if state_payload:
                 publish(state_payload, journal=False)
 
@@ -560,7 +570,9 @@ def run_pipeline_threadsafe(
             state_text = json.dumps(scene_state, indent=2)
 
             if attempt == 0:
+                ctx.episode_id = str(uuid.uuid4())
                 ctx.episode_run_tree = None
+                ctx.llm_latency_total = 0
                 if LANGSMITH_AVAILABLE:
                     from langsmith.run_trees import RunTree
                     ctx.episode_run_tree = RunTree(
@@ -608,6 +620,7 @@ def run_pipeline_threadsafe(
                     # LangSmith automatically handles children if wrapped, but we can also store explicitly
                     pass # Handled by LLMAgent wrapping if we did that, or we'll just end the tree at the end.
                 llm_time = time.time() - llm_start_time
+                ctx.llm_latency_total += llm_time
                 cancellation_token.checkpoint()
                 
             # Discard an LLM result that arrived after the user pressed stop.
@@ -683,7 +696,7 @@ def run_pipeline_threadsafe(
                 "score_value": 0,
                 "success": False,
             })
-            state_payload = get_simulation_state_payload(current_sim)
+            state_payload = get_simulation_state_payload(ctx)
             if state_payload:
                 publish(state_payload, journal=False)
             ctx.pipeline_executing = False
@@ -692,13 +705,13 @@ def run_pipeline_threadsafe(
         if not plan:
             log("������ [1] K�Q%% ho�Q%�ch r�W%�ng.")
             log("��% [5/5] %�ang t%�nh to%�n %�i�W%�m s�Q%�p x�Q%%p...")
-            score_val, score_md = compute_and_format_score(current_sim, [], True)
+            score_val, score_md = compute_and_format_score(ctx, [], True)
             publish({
                 "type": "result",
                 "score_display": score_md,
                 "score_value": score_val,
             })
-            state_payload = get_simulation_state_payload(current_sim)
+            state_payload = get_simulation_state_payload(ctx)
             if state_payload:
                 publish(state_payload, journal=False)
             ctx.pipeline_executing = False
@@ -740,7 +753,7 @@ def run_pipeline_threadsafe(
         
         try:
             # Trigger an initial simulation state transmission
-            state_payload = get_simulation_state_payload(current_sim)
+            state_payload = get_simulation_state_payload(ctx)
             if state_payload:
                 publish(state_payload, journal=False)
                 
@@ -821,7 +834,7 @@ def run_pipeline_threadsafe(
                 pb.stepSimulation = original_step
                 
             # Send a final state to ensure visual sync
-            state_payload = get_simulation_state_payload(current_sim)
+            state_payload = get_simulation_state_payload(ctx)
             if state_payload:
                 publish(state_payload, journal=False)
                 
@@ -886,7 +899,7 @@ def run_pipeline_threadsafe(
             
     # Step 5: Scoring
     cancellation_token.checkpoint()
-    score_val, score_md = compute_and_format_score(current_sim, plan, executed_successfully and goal_completed)
+    score_val, score_md = compute_and_format_score(ctx, plan, executed_successfully and goal_completed)
     
     error_cause = None
     if not plan:
@@ -898,6 +911,43 @@ def run_pipeline_threadsafe(
         
     is_success = executed_successfully and goal_completed and score_val >= 50
     
+    # Calculate PostHog metrics
+    token_usage = {"prompt_tokens": 0, "completion_tokens": 0}
+    if hasattr(ctx.sim.pipeline.agent, "planner") and hasattr(ctx.sim.pipeline.agent.planner, "last_token_usage"):
+        token_usage = getattr(ctx.sim.pipeline.agent.planner, "last_token_usage")
+    
+    system_ram_mb = psutil.virtual_memory().used / (1024 * 1024)
+    timeout_network_error = False
+    plan_validity_first_try = False
+    collision_free = True
+    kinematic_error = False
+    
+    if last_failed_sig:
+        if "Rate limit" in last_failed_sig or "Timeout" in last_failed_sig:
+            timeout_network_error = True
+        if "Kinematic" in last_failed_sig or "IK" in last_failed_sig or "outside workspace" in last_failed_sig.lower():
+            kinematic_error = True
+        if "collision" in last_failed_sig.lower():
+            collision_free = False
+            
+    if attempt == 0 and not last_failed_sig:
+        plan_validity_first_try = True
+
+    try:
+        log_episode_metrics(
+            session_id=getattr(ctx, "episode_id", "unknown"),
+            llm_latency_sec=getattr(ctx, "llm_latency_total", 0),
+            token_usage=token_usage,
+            system_ram_mb=system_ram_mb,
+            timeout_network_error=timeout_network_error,
+            plan_validity_first_try=plan_validity_first_try,
+            collision_free=collision_free,
+            kinematic_error=kinematic_error,
+            replan_count=attempt
+        )
+    except Exception as e:
+        log(f"Failed to send PostHog metrics: {e}")
+
     publish({
         "type": "result",
         "score_display": score_md,
@@ -914,19 +964,19 @@ async def command_worker_loop(loop, ctx: SessionContext):
     
     try:
         while True:
-            command = ctx.web_ctx.web_runtime.take_pending()
+            command = ctx.web_runtime.take_pending()
             if command is None:
                 return
             generation = command["generation"]
             task_id = command["task_id"]
             interrupted = False
-            await ctx.web_ctx.web_runtime.publish({
+            await ctx.web_runtime.publish({
                 "type": "command_accepted",
                 "prompt": command["prompt"],
                 "task_id": task_id,
                 "generation": generation,
             })
-            await ctx.web_ctx.web_runtime.publish({
+            await ctx.web_runtime.publish({
                 "type": "execution_status",
                 "state": "planning",
                 "task_id": task_id,
@@ -938,7 +988,7 @@ async def command_worker_loop(loop, ctx: SessionContext):
                     command["prompt"],
                     loop,
                     command["use_mock"],
-                    ctx.web_runtime,
+                    ctx,
                     command["token"],
                     generation,
                     task_id,
@@ -950,19 +1000,19 @@ async def command_worker_loop(loop, ctx: SessionContext):
                 interrupted = True
                 ctx.pipeline_executing = False
                 with ctx.pybullet_lock:
-                    if current_sim is not None and ctx.sim.pipeline is not None:
+                    if ctx.sim is not None and ctx.sim.pipeline is not None:
                         ctx.sim.pipeline.interrupt_hold()
-                ctx.web_ctx.web_runtime.set_status("interrupted", generation=generation)
+                ctx.web_runtime.set_status("interrupted", generation=generation)
                 state_payload = await asyncio.to_thread(
                     get_simulation_state_payload,
-                    current_sim,
+                    ctx.sim,
                 )
                 if state_payload:
-                    await ctx.web_ctx.web_runtime.publish(state_payload, journal=False)
-                ctx.web_ctx.web_runtime.conversation_manager.record_interruption(
+                    await ctx.web_runtime.publish(state_payload, journal=False)
+                ctx.web_runtime.conversation_manager.record_interruption(
                     command["prompt"]
                 )
-                await ctx.web_ctx.web_runtime.publish({
+                await ctx.web_runtime.publish({
                     "type": "interrupted",
                     "reason": str(exc),
                     "task_id": task_id,
@@ -970,7 +1020,7 @@ async def command_worker_loop(loop, ctx: SessionContext):
                 })
             except Exception as exc:
                 ctx.pipeline_executing = False
-                await ctx.web_ctx.web_runtime.publish({
+                await ctx.web_runtime.publish({
                     "type": "result",
                     "success": False,
                     "error_cause": "runtime",
@@ -979,24 +1029,24 @@ async def command_worker_loop(loop, ctx: SessionContext):
                     "generation": generation,
                 })
             finally:
-                if current_sim is not None and ctx.sim.pipeline is not None:
+                if ctx.sim is not None and ctx.sim.pipeline is not None:
                     ctx.sim.pipeline.set_cancellation_token(None)
-                ctx.web_ctx.web_runtime.complete(
+                ctx.web_runtime.complete(
                     generation,
                     state="interrupted" if interrupted else "completed",
                 )
     finally:
         ctx.command_worker_task = None
-        if ctx.web_ctx.web_runtime.has_pending:
+        if ctx.web_runtime.has_pending:
             ctx.command_worker_task = asyncio.create_task(command_worker_loop(loop))
 
 
 async def wait_for_active_task_stop(ctx: SessionContext, timeout=5.0):
-    ctx.web_ctx.web_runtime.cancel_all("runtime_mutation")
+    ctx.web_runtime.cancel_all("runtime_mutation")
     deadline = time.monotonic() + timeout
-    while ctx.web_ctx.web_runtime.active_token is not None and time.monotonic() < deadline:
+    while ctx.web_runtime.active_token is not None and time.monotonic() < deadline:
         await asyncio.sleep(0.02)
-    return ctx.web_ctx.web_runtime.active_token is None
+    return ctx.web_runtime.active_token is None
 
 
 
@@ -1093,6 +1143,8 @@ async def websocket_endpoint(websocket: WebSocket):
     async with ctx.executor_lock:
         def init_sim():
             with ctx.pybullet_lock:
+                if ctx.disconnected:
+                    return
                 if ctx.sim is None:
                     print(f"[{session_id}] Initializing PandaApp lazily...", flush=True)
                     env_path = os.path.join(PROJECT_SRC, "envs", "env1.json")
@@ -1104,6 +1156,8 @@ async def websocket_endpoint(websocket: WebSocket):
                             agent=None,
                             viz=False
                         )
+                        if ctx.disconnected:
+                            pb.disconnect(physicsClientId=ctx.sim.client)
                     except Exception as e:
                         print(f"Failed to initialize lazy PandaApp: {e}", flush=True)
                         raise e
@@ -1603,6 +1657,7 @@ async def websocket_endpoint(websocket: WebSocket):
         print(f"[{session_id}] WS Exception: {e}")
     finally:
         if 'ctx' in locals():
+            ctx.disconnected = True
             ctx.drag_state["active"] = False
             if ctx.drag_state.get("constraint_id") is not None:
                 try:
@@ -1632,4 +1687,5 @@ async def websocket_endpoint(websocket: WebSocket):
 
 if __name__ == "__main__":
     port = int(os.environ.get("PORT", 8000))
+    import uvicorn
     uvicorn.run("interactive_server:app", host="0.0.0.0", port=port, reload=False)