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)