Update app.py

app.py

@@ -143,14 +143,18 @@ def _sync_assets_globals_from_manager():
 
 def _any_jam_running() -> bool:
     with jam_lock:
-        return any(w.is_alive() for w in jam_registry.values())
+        return any(info['worker'].is_alive() for info in jam_registry.values())
 
 def _stop_all_jams(timeout: float = 5.0):
     with jam_lock:
-        for sid, w in list(jam_registry.items()):
+        for sid, info in list(jam_registry.items()):
+            w = info['worker']
             if w.is_alive():
                 w.stop()
                 w.join(timeout=timeout)
+                # Release MRT slot
+                if info.get('mrt_index') is not None:
+                    release_mrt(info['mrt_index'])
                 jam_registry.pop(sid, None)
 
 
@@ -202,7 +206,7 @@ def _patch_t5x_for_gpu_coords():
 # Call the patch immediately at import time (before MagentaRT init)
 _patch_t5x_for_gpu_coords()
 
-jam_registry: dict[str, JamWorker] = {}
+jam_registry: dict[str, dict] = {}  # Now stores {'worker': JamWorker, 'mrt_index': int}
 jam_lock = threading.Lock()
 
 # ============================================================================
@@ -1157,13 +1161,16 @@ def jam_stop(session_id: str = Body(..., embed=True)):
     mrt_index = session_info['mrt_index']
 
     worker.stop()
-    worker.join(timeout=5.0)
+    worker.join(timeout=10.0)  # Increased from 5s to 10s to allow chunk generation to finish
     if worker.is_alive():
-        # It's daemon=True, so it won't block process exit, but report it
-        print(f"⚠️ JamWorker {session_id} did not stop within timeout")
+        # Worker still running - don't release MRT slot to avoid corruption
+        print(f"⚠️ JamWorker {session_id} did not stop within timeout - keeping MRT slot reserved")
+        # Keep in registry so we can try to stop it again later
+        return {"stopped": False, "timeout": True, "message": "Worker did not stop in time, retry /jam/stop"}
 
-    # Release MRT back to pool
-    release_mrt(mrt_index)
+    # Only release MRT if worker actually stopped
+    if mrt_index is not None:
+        release_mrt(mrt_index)
 
     with jam_lock:
         jam_registry.pop(session_id, None)
@@ -1187,9 +1194,63 @@ def jam_stop_all():
             # Release MRT back to pool
             release_mrt(mrt_index)
             jam_registry.pop(session_id, None)
-    
+
     return {"stopped_sessions": stopped_sessions, "count": len(stopped_sessions)}
 
+@app.post("/jam/cleanup")
+def jam_cleanup(force: bool = Form(False), idle_threshold_seconds: float = Form(300.0)):
+    """
+    Enhanced cleanup endpoint for stopping stale/orphaned sessions.
+
+    - force=False: Only stops sessions idle for > idle_threshold_seconds (default 5 min)
+    - force=True: Stops ALL sessions regardless of activity (nuclear option)
+    """
+    stopped_sessions = []
+    kept_sessions = []
+    current_time = time.time()
+
+    with jam_lock:
+        for session_id, session_info in list(jam_registry.items()):
+            worker = session_info['worker']
+            mrt_index = session_info['mrt_index']
+
+            # Determine if session should be stopped
+            should_stop = force
+            idle_time = 0
+
+            if not force and hasattr(worker, 'last_activity_at'):
+                idle_time = current_time - worker.last_activity_at
+                should_stop = idle_time > idle_threshold_seconds
+
+            if should_stop:
+                if worker.is_alive():
+                    worker.stop()
+                    worker.join(timeout=10.0)
+
+                # Release MRT regardless of whether worker stopped
+                release_mrt(mrt_index)
+                jam_registry.pop(session_id, None)
+                stopped_sessions.append({
+                    "session_id": session_id,
+                    "idle_seconds": round(idle_time, 1),
+                    "slot": mrt_index
+                })
+            else:
+                kept_sessions.append({
+                    "session_id": session_id,
+                    "idle_seconds": round(idle_time, 1),
+                    "slot": mrt_index
+                })
+
+    return {
+        "stopped": stopped_sessions,
+        "kept": kept_sessions,
+        "stopped_count": len(stopped_sessions),
+        "kept_count": len(kept_sessions),
+        "force": force,
+        "idle_threshold_seconds": idle_threshold_seconds
+    }
+
 @app.post("/jam/update")
 def jam_update(
     session_id: str = Form(...),
@@ -1394,6 +1455,91 @@ def jam_status(session_id: str):
         "last_chunk_completed_at": worker.last_chunk_completed_at,
     }
 
+@app.get("/jam/sessions")
+def jam_sessions():
+    """List all active JAM sessions with metadata for monitoring"""
+    sessions = []
+    current_time = time.time()
+
+    with jam_lock:
+        for session_id, session_info in jam_registry.items():
+            worker = session_info['worker']
+            mrt_index = session_info['mrt_index']
+
+            # Calculate uptime and idle time
+            uptime = current_time - worker.created_at if hasattr(worker, 'created_at') else 0
+            last_activity = worker.last_activity_at if hasattr(worker, 'last_activity_at') else worker.created_at if hasattr(worker, 'created_at') else current_time
+            idle_time = current_time - last_activity
+
+            # Get generation stats
+            with worker._lock:
+                last_generated = int(worker.idx)
+                last_delivered = int(worker._last_delivered_index)
+                queued = len(worker.outbox)
+
+            sessions.append({
+                "session_id": session_id,
+                "mrt_slot": mrt_index,
+                "running": worker.is_alive(),
+                "uptime_seconds": round(uptime, 1),
+                "idle_seconds": round(idle_time, 1),
+                "chunks_generated": last_generated,
+                "chunks_delivered": last_delivered,
+                "chunks_queued": queued,
+                "bpm": worker.params.bpm,
+                "bars_per_chunk": worker.params.bars_per_chunk,
+                "created_at": worker.created_at if hasattr(worker, 'created_at') else None,
+                "last_activity_at": last_activity,
+            })
+
+    return {
+        "sessions": sessions,
+        "total_active": len(sessions),
+        "mrt_pool_size": len(_MRT_POOL),
+    }
+
+@app.get("/jam/sessions")
+def jam_sessions():
+    """List all active JAM sessions with metadata for monitoring"""
+    sessions = []
+    current_time = time.time()
+
+    with jam_lock:
+        for session_id, session_info in jam_registry.items():
+            worker = session_info['worker']
+            mrt_index = session_info['mrt_index']
+
+            # Calculate uptime and idle time
+            uptime = current_time - worker.created_at if hasattr(worker, 'created_at') else 0
+            last_activity = worker.last_activity_at if hasattr(worker, 'last_activity_at') else (worker.created_at if hasattr(worker, 'created_at') else current_time)
+            idle_time = current_time - last_activity
+
+            # Get generation stats
+            with worker._lock:
+                last_generated = int(worker.idx)
+                last_delivered = int(worker._next_to_deliver) - 1
+                queued = len(worker._outbox)
+
+            sessions.append({
+                "session_id": session_id,
+                "mrt_slot": mrt_index,
+                "running": worker.is_alive(),
+                "uptime_seconds": round(uptime, 1),
+                "idle_seconds": round(idle_time, 1),
+                "chunks_generated": last_generated,
+                "chunks_delivered": last_delivered,
+                "chunks_queued": queued,
+                "bpm": worker.params.bpm,
+                "bars_per_chunk": worker.params.bars_per_chunk,
+                "created_at": worker.created_at if hasattr(worker, 'created_at') else None,
+                "last_activity_at": last_activity,
+            })
+
+    return {
+        "sessions": sessions,
+        "total_active": len(sessions),
+        "mrt_pool_size": len(_MRT_POOL),
+    }
 
 @app.get("/health")
 def health():
@@ -1426,7 +1572,7 @@ def health():
     # 3) Ready; include operational hints
     warmed = bool(_WARMED)
     with jam_lock:
-        active_jams = sum(1 for w in jam_registry.values() if w.is_alive())
+        active_jams = sum(1 for info in jam_registry.values() if info['worker'].is_alive())
     return {
         "ok": True,
         "status": "ready" if warmed else "initializing",
@@ -1550,66 +1696,28 @@ async def ws_jam(websocket: WebSocket):
                         sid = str(uuid.uuid4())
                         with jam_lock:
                             # single active jam per GPU, mirroring /jam/start
-                            for _sid, w in list(jam_registry.items()):
-                                if w.is_alive():
+                            for _sid, info in list(jam_registry.items()):
+                                if info['worker'].is_alive():
                                     await send_json({"type":"error","error":"A jam is already running"})
                                     worker = None; sid = None
                                     break
                             if worker is not None:
-                                jam_registry[sid] = worker
+                                jam_registry[sid] = {'worker': worker, 'mrt_index': None}  # WebSocket mode doesn't use MRT pool
                                 worker.start()
 
                     else:
-                        # mode == "rt" (with optional loop context)
+                        # mode == "rt" (Colab-style, no loop context)
                         mrt = get_mrt()
                         state = mrt.init_state()
 
-                        # Build context tokens (silent or from loop)
+                        # Build silent context (10s) tokens
                         codec_fps   = float(mrt.codec.frame_rate)
                         ctx_seconds = float(mrt.config.context_length_frames) / codec_fps
                         sr = int(mrt.sample_rate)
-
-                        # Check for optional loop audio
-                        loop_b64 = msg.get("loop_audio_b64")
-                        loop_embed = None
-
-                        if loop_b64:
-                            try:
-                                # Decode and load loop (similar to bar-mode)
-                                loop_bytes = base64.b64decode(loop_b64)
-                                with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as tmp:
-                                    tmp.write(loop_bytes)
-                                    tmp_path = tmp.name
-
-                                loop = au.Waveform.from_file(tmp_path).resample(mrt.sample_rate).as_stereo()
-
-                                # Extract bar-aligned tail for embedding
-                                bpm = float(params.get("bpm", 120.0))
-                                bpb = int(params.get("beats_per_bar", 4))
-                                loop_tail = take_bar_aligned_tail(loop, bpm, bpb, ctx_seconds)
-
-                                # Embed the loop audio
-                                loop_embed = mrt.embed_style(loop_tail)
-
-                                # Use loop audio for context tokens instead of silent
-                                tokens = mrt.codec.encode(loop_tail).astype(np.int32)[:, :mrt.config.decoder_codec_rvq_depth]
-                                state.context_tokens = tokens
-
-                            except Exception as e:
-                                # Log error but continue with silent context
-                                print(f"Loop audio processing failed: {e}")
-                                loop_embed = None
-                                # Fall back to silent context
-                                samples = int(max(1, round(ctx_seconds * sr)))
-                                silent = au.Waveform(np.zeros((samples, 2), np.float32), sr)
-                                tokens = mrt.codec.encode(silent).astype(np.int32)[:, :mrt.config.decoder_codec_rvq_depth]
-                                state.context_tokens = tokens
-                        else:
-                            # No loop provided - use silent context (original behavior)
-                            samples = int(max(1, round(ctx_seconds * sr)))
-                            silent = au.Waveform(np.zeros((samples, 2), np.float32), sr)
-                            tokens = mrt.codec.encode(silent).astype(np.int32)[:, :mrt.config.decoder_codec_rvq_depth]
-                            state.context_tokens = tokens
+                        samples = int(max(1, round(ctx_seconds * sr)))
+                        silent = au.Waveform(np.zeros((samples, 2), np.float32), sr)
+                        tokens = mrt.codec.encode(silent).astype(np.int32)[:, :mrt.config.decoder_codec_rvq_depth]
+                        state.context_tokens = tokens
 
                         # Parse params (including steering)
                         asset_manager.ensure_assets_loaded(get_mrt())
@@ -1617,7 +1725,6 @@ async def ws_jam(websocket: WebSocket):
                         style_weights_str = params.get("style_weights", "") or ""
                         mean_w            = float(params.get("mean", 0.0) or 0.0)
                         cw_str            = str(params.get("centroid_weights", "") or "")
-                        loop_weight       = float(params.get("loop_weight", 1.0) or 1.0)
 
                         text_list = [s.strip() for s in styles_str.split(",") if s.strip()]
                         try:
@@ -1633,13 +1740,13 @@ async def ws_jam(websocket: WebSocket):
                         if _CENTROIDS is not None and len(cw) > int(_CENTROIDS.shape[0]):
                             cw = cw[: int(_CENTROIDS.shape[0])]
 
-                        # Build initial style vector WITH optional loop_embed
+                        # Build initial style vector (no loop_embed in rt mode)
                         style_vec = build_style_vector(
                             mrt,
                             text_styles=text_list,
                             text_weights=text_w,
-                            loop_embed=loop_embed,
-                            loop_weight=loop_weight,
+                            loop_embed=None,
+                            loop_weight=None,
                             mean_weight=mean_w,
                             centroid_weights=cw,
                         )
@@ -1798,15 +1905,16 @@ async def ws_jam(websocket: WebSocket):
 
             elif mtype == "consume" and mode == "bar":
                 with jam_lock:
-                    worker = jam_registry.get(msg.get("session_id"))
-                if worker is not None:
-                    worker.mark_chunk_consumed(int(msg.get("chunk_index", -1)))
+                    session_info = jam_registry.get(msg.get("session_id"))
+                if session_info is not None:
+                    session_info['worker'].mark_chunk_consumed(int(msg.get("chunk_index", -1)))
 
             elif mtype == "reseed" and mode == "bar":
                 with jam_lock:
-                    worker = jam_registry.get(msg.get("session_id"))
-                if worker is None or not worker.is_alive():
+                    session_info = jam_registry.get(msg.get("session_id"))
+                if session_info is None or not session_info['worker'].is_alive():
                     await send_json({"type":"error","error":"Session not found"}); continue
+                worker = session_info['worker']
                 loop_b64 = msg.get("loop_audio_b64")
                 if not loop_b64:
                     await send_json({"type":"error","error":"loop_audio_b64 required"}); continue
@@ -1819,9 +1927,10 @@ async def ws_jam(websocket: WebSocket):
 
             elif mtype == "reseed_splice" and mode == "bar":
                 with jam_lock:
-                    worker = jam_registry.get(msg.get("session_id"))
-                if worker is None or not worker.is_alive():
+                    session_info = jam_registry.get(msg.get("session_id"))
+                if session_info is None or not session_info['worker'].is_alive():
                     await send_json({"type":"error","error":"Session not found"}); continue
+                worker = session_info['worker']
                 anchor = float(msg.get("anchor_bars", 2.0))
                 b64 = msg.get("combined_audio_b64")
                 if b64: