Fix vLLM token parameter and improve streaming error handling

- Remove 'token' parameter from vLLM LLM() call (uses HF_TOKEN env var)
- Add better error handling for generation thread
- Add debug logging for streamer token consumption
- Add timeout handling for generation thread

app.py

@@ -138,10 +138,10 @@ def load_vllm_model(model_name: str):
     try:
         # vLLM configuration optimized for ZeroGPU H200 slice
         # vLLM natively supports AWQ via llm-compressor (replaces deprecated AutoAWQ)
+        # Note: HF_TOKEN is passed via environment variable, not as a parameter
         llm_kwargs = {
             "model": repo,
             "trust_remote_code": True,
-            "token": HF_TOKEN,
             "dtype": "bfloat16",  # Prefer bf16 over int8 for speed
             "gpu_memory_utilization": 0.90,  # Leave headroom for KV cache
             "max_model_len": 16384,  # Adjust based on GPU memory
@@ -675,45 +675,79 @@ def _generate_router_plan_streaming_internal(
                 "pad_token_id": tokenizer.pad_token_id or tokenizer.eos_token_id,
             }
 
+            generation_error = None
+            
             def _generate():
-                with torch.inference_mode():
-                    model.generate(**generation_kwargs)
+                nonlocal generation_error
+                try:
+                    with torch.inference_mode():
+                        model.generate(**generation_kwargs)
+                except Exception as e:
+                    generation_error = e
+                    print(f"[DEBUG] Generation thread error: {e}")
+                    import traceback
+                    traceback.print_exc()
 
             thread = Thread(target=_generate)
             thread.start()
 
-        # Stream tokens
+            # Stream tokens
             completion = ""
             parsed_plan: Dict[str, Any] | None = None
             validation_msg = "🔄 Generating..."
-
-            for new_text in streamer:
-                completion += new_text
-                chunk = completion
-                finished = False
-                display_plan = parsed_plan or {}
-
-                chunk, finished = trim_at_stop_sequences(chunk)
-
-                try:
-                    json_block = extract_json_from_text(chunk)
-                    candidate_plan = json.loads(json_block)
-                    ok, issues = validate_router_plan(candidate_plan)
-                    validation_msg = format_validation_message(ok, issues)
-                    parsed_plan = candidate_plan if ok else parsed_plan
-                    display_plan = candidate_plan
-                except Exception:
-                    # Ignore until JSON is complete
-                    pass
-
-                yield chunk, display_plan, validation_msg, prompt
-
-                if finished:
-                    completion = chunk
-                    break
+            
+            print(f"[DEBUG] Starting to consume streamer...")
+            token_count = 0
+            
+            try:
+                for new_text in streamer:
+                    if generation_error:
+                        raise generation_error
+                    
+                    if new_text:
+                        token_count += 1
+                        completion += new_text
+                        chunk = completion
+                        finished = False
+                        display_plan = parsed_plan or {}
+
+                        chunk, finished = trim_at_stop_sequences(chunk)
+
+                        try:
+                            json_block = extract_json_from_text(chunk)
+                            candidate_plan = json.loads(json_block)
+                            ok, issues = validate_router_plan(candidate_plan)
+                            validation_msg = format_validation_message(ok, issues)
+                            parsed_plan = candidate_plan if ok else parsed_plan
+                            display_plan = candidate_plan
+                        except Exception:
+                            # Ignore until JSON is complete
+                            pass
+
+                        yield chunk, display_plan, validation_msg, prompt
+
+                        if finished:
+                            completion = chunk
+                            break
+                
+                print(f"[DEBUG] Streamer finished. Received {token_count} tokens.")
+            except Exception as stream_error:
+                print(f"[DEBUG] Streamer error: {stream_error}")
+                import traceback
+                traceback.print_exc()
+                # Wait for thread to finish
+                thread.join(timeout=5.0)
+                if generation_error:
+                    raise generation_error
+                raise stream_error
             
             # Final processing after streaming completes
-            thread.join()
+            thread.join(timeout=30.0)
+            if thread.is_alive():
+                print("[DEBUG] WARNING: Generation thread still running after timeout")
+            
+            if generation_error:
+                raise generation_error
 
         completion = trim_at_stop_sequences(completion.strip())[0]
         print(f"[DEBUG] Final completion length: {len(completion)}")