fix: stabilize e2e test suite — session resilience + LegitData API call reduction

Tests 2-6 (page load timeout): add _navigate_and_ensure_logged_in() to
submit_job — waits for tabs OR login form, re-logs in on session expiry,
retries up to 5min if space is rebuilding after a push.

Test 1 (45-min hang): generate_entity_batch was making N/20 API calls per
dimension table (10 calls × 5 tables = 50+ calls for Banking/Marketing standard
preset). Changed to one seed call (≤50 entities) then cycle — cuts API calls
from 50+ down to ~5, well within the 15-min POP_TIMEOUT.

Monitoring loop: now detects pipeline error text in page body and stage stuck
>20min to bail early instead of waiting the full 45min timeout.

Also includes ★ CUSTOM ★ vertical rename (demo_personas.py + chat_interface.py).

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

chat_interface.py

@@ -5411,15 +5411,15 @@ def create_chat_tab(chat_controller_state, settings, current_stage, current_mode
         # App tab: vertical → line + function cascade
         def update_line_on_vertical(vertical):
             lines = VERTICAL_LINES.get(vertical, [])
-            if vertical == "Custom" or not lines:
+            if vertical == "★ CUSTOM ★" or not lines:
                 return (
-                    gr.Dropdown(choices=[], value=None, interactive=False),
-                    gr.Dropdown(choices=[], value=None, interactive=False),
+                    gr.Dropdown(choices=["— not used —"], value="— not used —", interactive=False, label="Line (n/a for custom)"),
+                    gr.Dropdown(choices=["— not used —"], value="— not used —", interactive=False, label="Function (n/a for custom)"),
                     gr.Textbox(label="Context *", placeholder="Describe your use case, industry, and key metrics...", interactive=True),
                 )
             return (
-                gr.Dropdown(choices=lines, value=lines[0], interactive=True),
-                gr.Dropdown(choices=DEMO_FUNCTIONS, value=DEMO_FUNCTIONS[0], interactive=True),
+                gr.Dropdown(choices=lines, value=lines[0], interactive=True, label="Line"),
+                gr.Dropdown(choices=DEMO_FUNCTIONS, value=DEMO_FUNCTIONS[0], interactive=True, label="Function"),
                 gr.Textbox(label="Context", placeholder="Any extra context for the demo...", interactive=True),
             )
 
@@ -5457,7 +5457,7 @@ def create_chat_tab(chat_controller_state, settings, current_stage, current_mode
             else:
                 raw_company = f"{line or vertical} Demo"
 
-            is_custom = (vertical == "Custom")
+            is_custom = (vertical == "★ CUSTOM ★")
             hide_welcome = gr.update(visible=False)
 
             if is_custom:
@@ -5762,7 +5762,7 @@ def create_settings_tab():
     # Build use case choices from VERTICAL_LINES × DEMO_FUNCTIONS
     use_case_choices = []
     for v_name, lines in VERTICAL_LINES.items():
-        if v_name == "Custom":
+        if v_name == "★ CUSTOM ★":
             continue
         for line in lines:
             for f_name in DEMO_FUNCTIONS:

demo_personas.py

@@ -115,7 +115,7 @@ VERTICAL_LINES = {
     "Agriculture": [
         "Forestry", "Fishing", "Hunting", "Cultivation", "Raw Materials",
     ],
-    "Custom": [],
+    "★ CUSTOM ★": [],
 }
 
 DEMO_FUNCTIONS = ["Sales", "Marketing", "Finance", "HR", "IT", "Legal"]

legitdata_project/legitdata/sourcer/ai_generator.py

@@ -221,15 +221,12 @@ Return ONLY a JSON array of strings. Example: ["value1", "value2", "value3"]"""
         columns_desc = "\n".join(column_specs)
         column_names = [col['name'] for col in columns]
         
-        # Generate in batches to avoid token limits
+        # Generate a "seed" batch via a single API call, then cycle for remaining rows.
+        # Avoids making N/20 API calls when num_entities is large (e.g. 200 dim rows).
+        seed_target = min(50, num_entities)
         all_entities = []
-        batch_size = min(20, num_entities)
-        
-        while len(all_entities) < num_entities:
-            remaining = num_entities - len(all_entities)
-            current_batch = min(batch_size, remaining)
-            
-            prompt = f"""Generate {current_batch} realistic, coherent {table_name} records for a database.
+
+        prompt = f"""Generate {seed_target} realistic, coherent {table_name} records for a database.
 
 Company: {context.company_name}
 Industry: {context.industry}
@@ -243,7 +240,7 @@ CRITICAL: The values in each row must be COHERENT with each other. For example:
 - If product_category is "Books", product_description should describe a book
 - If brand_name is "Nike", the product should be sports/athletic related
 
-Generate {current_batch} diverse, realistic records. Each record should tell a coherent story.
+Generate {seed_target} diverse, realistic records. Each record should tell a coherent story.
 
 Return ONLY a JSON array of objects. Example format:
 [
@@ -251,54 +248,55 @@ Return ONLY a JSON array of objects. Example format:
   ...
 ]"""
 
-            try:
-                response = self.anthropic_client.messages.create(
-                    model="claude-sonnet-4-20250514",
-                    max_tokens=4000,
-                    messages=[{"role": "user", "content": prompt}]
-                )
-                
-                json_str = response.content[0].text.strip()
-                # Handle markdown
-                if json_str.startswith('```'):
-                    json_str = re.sub(r'^```\w*\n?', '', json_str)
-                    json_str = re.sub(r'\n?```$', '', json_str)
-                
-                entities = json.loads(json_str)
-                
-                if isinstance(entities, list):
-                    for entity in entities:
-                        if isinstance(entity, dict):
-                            # Normalize keys to match column names
-                            normalized = {}
-                            for col in columns:
-                                col_name = col['name']
-                                # Try exact match first, then case-insensitive
-                                if col_name in entity:
-                                    normalized[col_name] = str(entity[col_name])
-                                else:
-                                    for key in entity:
-                                        if key.lower() == col_name.lower():
-                                            normalized[col_name] = str(entity[key])
-                                            break
-                            if normalized:
-                                all_entities.append(normalized)
-                
-            except Exception as e:
-                print(f"Warning: Entity batch generation failed for {table_name}: {e}")
-                # Fall back to single-column generation
-                break
-            
-            if len(all_entities) >= num_entities:
-                break
-        
-        # If we didn't get enough, fill with fallback
+        try:
+            response = self.anthropic_client.messages.create(
+                model="claude-sonnet-4-20250514",
+                max_tokens=4000,
+                messages=[{"role": "user", "content": prompt}]
+            )
+
+            json_str = response.content[0].text.strip()
+            # Handle markdown
+            if json_str.startswith('```'):
+                json_str = re.sub(r'^```\w*\n?', '', json_str)
+                json_str = re.sub(r'\n?```$', '', json_str)
+
+            entities = json.loads(json_str)
+
+            if isinstance(entities, list):
+                for entity in entities:
+                    if isinstance(entity, dict):
+                        # Normalize keys to match column names
+                        normalized = {}
+                        for col in columns:
+                            col_name = col['name']
+                            # Try exact match first, then case-insensitive
+                            if col_name in entity:
+                                normalized[col_name] = str(entity[col_name])
+                            else:
+                                for key in entity:
+                                    if key.lower() == col_name.lower():
+                                        normalized[col_name] = str(entity[key])
+                                        break
+                        if normalized:
+                            all_entities.append(normalized)
+
+        except Exception as e:
+            print(f"Warning: Entity batch generation failed for {table_name}: {e}")
+
+        # If seed generation produced results, cycle them to fill remaining rows
+        if all_entities and len(all_entities) < num_entities:
+            seed = list(all_entities)
+            while len(all_entities) < num_entities:
+                all_entities.append(seed[len(all_entities) % len(seed)].copy())
+
+        # If we still don't have enough (seed generation failed), fall back
         if len(all_entities) < num_entities:
             fallback = self._generate_entity_fallback(
                 table_name, columns, num_entities - len(all_entities)
             )
             all_entities.extend(fallback)
-        
+
         return all_entities[:num_entities]
     
     def _generate_entity_fallback(

tests/e2e_quality.py

@@ -185,11 +185,55 @@ def select_gradio_dropdown(page: Page, label: str, value: str):
     page.wait_for_timeout(300)
 
 
+def _do_login(page: Page):
+    """Fill and submit the login form, then wait for tabs."""
+    page.fill('input[type=text]', TEST_USER)
+    page.fill('input[type=password]', TEST_PASSWORD)
+    page.click('button:has-text("Login")')
+    page.wait_for_selector('button[role=tab]', timeout=90000)
+    page.wait_for_timeout(3000)
+
+
+def _navigate_and_ensure_logged_in(page: Page, max_wait_secs: int = 300):
+    """
+    Navigate to BASE_URL and ensure we're on the logged-in app.
+    Handles: HF space sleeping/rebuilding after a long test, session expiry.
+    Retries for up to max_wait_secs before raising.
+    """
+    deadline = time.time() + max_wait_secs
+    attempt = 0
+    while True:
+        attempt += 1
+        try:
+            page.goto(BASE_URL, timeout=90000)
+            # Wait for either the logged-in app (tabs) or the login form
+            page.wait_for_selector(
+                'button[role=tab], input[type=password]',
+                timeout=60000,
+            )
+            break
+        except Exception as nav_err:
+            remaining = int(deadline - time.time())
+            if remaining <= 0:
+                raise RuntimeError(
+                    f"Space not reachable after {max_wait_secs}s: {nav_err}"
+                ) from nav_err
+            print(f"  ⚠️ Space not ready (attempt {attempt}) — retrying in 30s ({remaining}s left)...")
+            time.sleep(30)
+
+    # If we landed on the login page (session expired or space rebuilt), re-login
+    try:
+        if page.locator('input[type=password]').is_visible(timeout=2000):
+            print("  🔑 Session expired — re-logging in...")
+            _do_login(page)
+    except Exception:
+        pass  # Already on the app — no login needed
+
+
 def submit_job(page: Page, test_case: dict):
     """Fill the form and click GO."""
-    # Navigate to App tab, then App sub-tab
-    page.goto(BASE_URL, timeout=90000)
-    page.wait_for_selector('button[role=tab]', timeout=60000)
+    # Navigate, re-logging in if the session expired (e.g. after a long prior test)
+    _navigate_and_ensure_logged_in(page)
     page.wait_for_timeout(2000)
 
     page.click('button[role=tab]:has-text("📱 App")', timeout=10000)
@@ -914,8 +958,17 @@ def run_single_test(page: Page, test_case: dict, config: dict) -> dict:
         submit_job(page, test_case)
         print(f"  ⏳ Monitoring pipeline (timeout: {config['grading']['timeout_minutes']}min)...")
 
-        poll_interval = 15
-        last_stages   = {}
+        poll_interval    = 15
+        last_stages      = {}
+        last_change_time = time.time()
+        STUCK_THRESHOLD  = 20 * 60  # 20 min with no stage change → bail early
+        PIPELINE_ERROR_INDICATORS = [
+            "pipeline has been interrupted",
+            "An unexpected error occurred",
+            "Population failed",
+            "Pipeline failed",
+            "Something went wrong during the pipeline",
+        ]
         while time.time() - start < timeout_sec:
             time.sleep(poll_interval)
             stages = read_progress(page)
@@ -923,10 +976,26 @@ def run_single_test(page: Page, test_case: dict, config: dict) -> dict:
                 done    = [k for k, v in stages.items() if v == "complete"]
                 running = [k for k, v in stages.items() if v == "running"]
                 print(f"    ✓ {done}  ▶ {running}")
-                last_stages = stages
+                last_stages      = stages
+                last_change_time = time.time()
             if pipeline_finished(stages):
                 print("  ✅ Pipeline complete")
                 break
+            # Detect hard pipeline failure in the chat output
+            try:
+                page_text = page.inner_text('body')
+                if any(ind in page_text for ind in PIPELINE_ERROR_INDICATORS):
+                    result["timed_out"] = True
+                    print("  ❌ Pipeline error detected in page — stopping early")
+                    break
+            except Exception:
+                pass
+            # Bail if stages have been stuck for too long
+            stuck_secs = time.time() - last_change_time
+            if stuck_secs > STUCK_THRESHOLD and any(v == "running" for v in stages.values()):
+                result["timed_out"] = True
+                print(f"  ⏰ Stage stuck for {int(stuck_secs/60)}min — treating as failure")
+                break
         else:
             result["timed_out"] = True
             print(f"  ⏰ Timed out after {config['grading']['timeout_minutes']} min")
@@ -1039,11 +1108,7 @@ def run_quality_suite(max_tests: int = None):
         page = ctx.new_page()
         page.goto(BASE_URL, timeout=90000)
         page.wait_for_selector('input[placeholder="Type here..."]', timeout=90000)
-        page.fill('input[type=text]', TEST_USER)
-        page.fill('input[type=password]', TEST_PASSWORD)
-        page.click('button:has-text("Login")')
-        page.wait_for_selector('button[role=tab]', timeout=90000)
-        page.wait_for_timeout(3000)
+        _do_login(page)
         print("✅ Logged in\n")
 
         for i, test_case in enumerate(suite, 1):