app.py

import gevent.monkey
gevent.monkey.patch_all(asyncio=True) # Keep this at the very top

import asyncio
from flask import Flask, request, jsonify
from proxy_lite import Runner, RunnerConfig
import os
import logging
from datetime import datetime
from playwright.async_api import async_playwright, TimeoutError as PlaywrightTimeoutError

logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

app = Flask(__name__)

_runner = None 

async def perform_hardcoded_salesforce_login_and_get_cookies(username, password, login_url, target_url):
    logger.info("Attempting hardcoded Salesforce login with Playwright to obtain cookies...")
    async with async_playwright() as p:
        browser = await p.chromium.launch(headless=True, args=["--no-sandbox", "--disable-setuid-sandbox"])
        context = await browser.new_context()
        page = await context.new_page()

        try:
            await page.goto(login_url, wait_until="domcontentloaded", timeout=60000)
            logger.info(f"Playwright: Navigated to Salesforce login page: {page.url}")

            await page.fill("#username", username)
            await page.fill("#password", password)
            await page.click("#Login")
            logger.info("Playwright: Filled credentials and clicked Login. Waiting for post-login state...")

            try:
                await page.wait_for_url(lambda url: "login.salesforce.com" not in url and "unauthorized" not in url.lower(), timeout=60000)
                logger.info(f"Playwright: Successfully redirected from login page. Current URL: {page.url}")
                await page.wait_for_selector('button[title="App Launcher"]', timeout=30000)
                logger.info("Playwright: Main Salesforce Lightning UI (e.g., App Launcher) detected after login.")

            except PlaywrightTimeoutError:
                logger.error(f"Playwright: Did not detect main UI or expected URL change within timeout after login. Current URL: {page.url}. Login might have failed or stuck on a redirect loop.")
                raise Exception("Salesforce login redirection failed or main UI not detected.")
            
            logger.info(f"Playwright: Navigating to target URL: {target_url} to ensure all relevant cookies are captured.")
            await page.goto(target_url, wait_until="domcontentloaded", timeout=60000)

            try:
                # Wait for generic Salesforce setup page elements to load
                await page.wait_for_selector('.setupPage, .slds-page-header, .slds-card, [data-aura-class*="setup"], .forcePageBlockSectionView', timeout=30000)
                logger.info("Playwright: Detected Salesforce setup page elements loaded successfully.")
            except PlaywrightTimeoutError:
                logger.warning("Playwright: Specific setup page elements not found. Trying generic page load check...")
                try:
                    # Fallback: wait for page to reach network idle state
                    await page.wait_for_load_state("networkidle", timeout=10000)
                    logger.info("Playwright: Page reached network idle state - proceeding with task.")
                except PlaywrightTimeoutError:
                    logger.info("Playwright: Page load validation timed out, but continuing as page may still be functional.")

            await asyncio.sleep(2)
            logger.info(f"Playwright: Successfully navigated to and confirmed content on {page.url}")

            cookies = await context.cookies()
            logger.info(f"Playwright: Extracted {len(cookies)} cookies after successful login and navigation.")
            return cookies

        except PlaywrightTimeoutError as e:
            logger.error(f"Playwright login/navigation failed (Timeout): {e}. Current URL: {page.url}")
            raise 
        except Exception as e:
            logger.error(f"Playwright login/navigation failed (General Error): {e}. Current URL: {page.url}")
            raise
        finally:
            if browser:
                await browser.close()


async def initialize_runner_with_cookies(cookies: list, target_url: str):
    global _runner
    logger.info("Initializing Proxy-lite Runner with provided cookies...")
    
    gemini_api_key = os.environ.get("GEMINI_API_KEY")
    if not gemini_api_key:
        logger.error("GEMINI_API_KEY environment variable not set. Cannot initialize Runner.")
        raise ValueError("GEMINI_API_KEY environment variable not set. Please set it as a Space secret.")

    config_dict = {
        "environment": {
            "name": "webbrowser",
            "homepage": "about:blank",  # Safe startup, we'll open new tab programmatically
            "headless": True, 
            "launch_args": ["--no-sandbox", "--disable-setuid-sandbox"],
            "screenshot_delay": 0.5, 
            "include_html": True,
            "include_poi_text": True,
            "record_pois": True, 
            "viewport_width": 1280,
            "viewport_height": 720,
            "browserbase_timeout": 7200,
            "keep_original_image": False,
            "no_pois_in_image": False,
            "initial_cookies": cookies 
        },
        "solver": {
            "name": "simple",
            "agent": {
                "name": "proxy_lite",
                "client": {
                    "name": "gemini",
                    "model_id": "gemini-2.0-flash-001",
                    "api_key": gemini_api_key,
                    "http_timeout": 50.0,
                    "http_concurrent_connections": 50,
                },
                "history_messages_limit": {
                    "screenshot": 1
                },
                "history_messages_include": None,
            }
        },
        "environment_timeout": 1800.0,
        "action_timeout": 1800.0,
        "task_timeout": 18000.0,
        "max_steps": 150,
        "logger_level": "DEBUG", 
        "save_every_step": True,
        "detailed_logger_name": False
    }
    config = RunnerConfig.from_dict(config_dict)

    logger.info(f"DEBUG: app.py - Initializing Proxy-lite Runner with Gemini Flash 2.0 configuration.")
    _runner = Runner(config=config)
    logger.info("Proxy-lite Runner initialized successfully with Gemini Flash 2.0 and injected cookies.")
    return _runner


@app.route('/run_proxy_task', methods=['POST'])
async def run_proxy_task_endpoint():
    data = request.json
    request_task_instruction = data.get('task')
    target_url = data.get('url')

    if not request_task_instruction:
        logger.warning("Received request without 'task' field. Returning 400.")
        return jsonify({"error": "No 'task' provided in request body"}), 400
        
    if not target_url:
        logger.warning("Received request without 'url' field. Returning 400.")
        return jsonify({"error": "No 'url' provided in request body"}), 400

    logger.info(f"Received user request task: '{request_task_instruction}'")
    logger.info(f"Target URL: '{target_url}'")

    # Check if this is a Salesforce URL
    is_salesforce_url = "salesforce.com" in target_url or "force.com" in target_url
    
    try:
        if is_salesforce_url:
            # Salesforce automation - requires login
            salesforce_username = os.environ.get("SALESFORCE_USERNAME")
            salesforce_password = os.environ.get("SALESFORCE_PASSWORD")

            if not salesforce_username or not salesforce_password:
                logger.error("Salesforce credentials (SALESFORCE_USERNAME, SALESFORCE_PASSWORD) environment variables not set.")
                return jsonify({"error": "Salesforce credentials not configured. Please set SALESFORCE_USERNAME and SALESFORCE_PASSWORD as Space secrets."}), 500

            salesforce_login_url = "https://login.salesforce.com/"
            logger.info("Executing hardcoded login via Playwright to get session cookies...")
            session_cookies = await perform_hardcoded_salesforce_login_and_get_cookies(
                salesforce_username, salesforce_password, salesforce_login_url, target_url
            )
            logger.info(f"Successfully obtained {len(session_cookies)} cookies. These will be injected into the agent's browser.")
        else:
            # General web browsing - no login required
            logger.info("Non-Salesforce URL detected. Skipping Salesforce login.")
            session_cookies = []

        runner = await initialize_runner_with_cookies(session_cookies, target_url)
        logger.info("Proxy-lite Runner initialized with cookies." if session_cookies else "Proxy-lite Runner initialized for general web browsing.")

        logger.info("Agent will use mandatory new tab tool to bypass loading issues.")
        
        # MANDATORY new tab navigation task - this is critical to avoid loading issues
        agent_task = f"""
CRITICAL FIRST STEP - MANDATORY: 
Your VERY FIRST action must be to use the open_new_tab_and_go_to tool to navigate to {target_url}

DO NOT skip this step. DO NOT use goto. You MUST use: open_new_tab_and_go_to(url='{target_url}')

This is necessary because direct navigation to this URL gets stuck loading. The new tab approach bypasses this issue.

STEP 1: Use open_new_tab_and_go_to(url='{target_url}')
STEP 2: Wait for the page to be fully loaded (no loading spinners visible)
STEP 3: {request_task_instruction}

CRITICAL WORKFLOW - FOLLOW THESE EXACT STEPS IN SEQUENCE:

STEP A: Select Permission Set
- Use select_option_by_text tool to find and select the target permission set from Available list
- Wait for "[ACTION COMPLETED]" response before proceeding

STEP B: Click Add Button  
- After successful selection, immediately click the "Add" button to move permission set to Enabled list
- Do NOT repeat the selection - proceed directly to Add button

STEP C: Click Save Button
- After clicking Add, immediately click "Save" to persist the changes
- After Save, Salesforce redirects to User page indicating SUCCESS

CRITICAL: Do NOT repeat actions. Each step should happen exactly once in sequence.

GENERAL INSTRUCTIONS:
- You must EXECUTE all actions immediately - do NOT just describe what you plan to do
- Do NOT wait for user input or ask "what should I do next?"
- Complete the entire task autonomously using the available tools
- After completing all steps, use the return_value tool to provide your final response
- If you make a plan, IMMEDIATELY execute it step by step using the appropriate tools
        """
        
        logger.info("Executing agent task with mandatory new tab navigation...")
        result = await runner.run(task=agent_task)

        # Extract the actual result value from the Run object
        if hasattr(result, 'value') and result.value:
            task_result = str(result.value)
        elif hasattr(result, 'result') and result.result:
            task_result = str(result.result)
        else:
            task_result = str(result)
            
        logger.info(f"Proxy-lite task completed. Output (truncated for log): {task_result[:500]}...")
        
        # Structure response for LWC integration
        response = {
            "status": "success",
            "message": "Task completed successfully",
            "data": {
                "task_result": task_result,
                "steps_completed": [
                    "Hardcoded Salesforce login completed",
                    "Browser session initialized with cookies",
                    "New tab navigation executed",
                    "Target Salesforce setup page accessed",
                    "Task execution completed successfully"
                ],
                "environment": {
                    "target_url": target_url,
                    "cookies_count": len(session_cookies),
                    "navigation_method": "new_tab_bypass"
                }
            },
            "timestamp": datetime.now().isoformat(),
            "task_request": request_task_instruction
        }
        
        return jsonify(response)
        
    except PlaywrightTimeoutError as e:
        logger.exception(f"Playwright timeout during login/navigation: {e}")
        error_response = {
            "status": "error",
            "error_type": "navigation_timeout",
            "message": "Page loading timed out during login or navigation",
            "data": {
                "error_details": str(e),
                "suggested_action": "Retry the request - network issues may be temporary",
                "steps_completed": ["Login attempted", "Navigation failed due to timeout"]
            },
            "timestamp": datetime.now().isoformat(),
            "task_request": request_task_instruction
        }
        return jsonify(error_response), 500
        
    except ValueError as e:
        logger.exception(f"Configuration error: {e}")
        error_response = {
            "status": "error", 
            "error_type": "configuration_error",
            "message": "System configuration issue",
            "data": {
                "error_details": str(e),
                "suggested_action": "Check environment variables and system configuration",
                "steps_completed": ["Configuration validation failed"]
            },
            "timestamp": datetime.now().isoformat(),
            "task_request": request_task_instruction
        }
        return jsonify(error_response), 500
        
    except Exception as e:
        logger.exception(f"Unexpected error processing Salesforce task: {e}")
        error_response = {
            "status": "error",
            "error_type": "unexpected_error", 
            "message": "An unexpected error occurred during task execution",
            "data": {
                "error_details": str(e),
                "error_class": type(e).__name__,
                "suggested_action": "Check logs for detailed error information and retry",
                "steps_completed": ["Login attempted", "Error occurred during execution"]
            },
            "timestamp": datetime.now().isoformat(),
            "task_request": request_task_instruction
        }
        return jsonify(error_response), 500

@app.route('/')
def root():
    logger.info("Root endpoint accessed.")
    return "Proxy-lite API is running. Send POST requests to /run_proxy_task with a 'task' in JSON body."

@app.route('/health', methods=['GET'])
def health_check():
    """Health check endpoint for monitoring and debugging"""
    logger.info("Health check endpoint accessed.")
    
    # Check environment variables
    env_status = {
        "GEMINI_API_KEY": "✓" if os.environ.get("GEMINI_API_KEY") else "✗",
        "SALESFORCE_USERNAME": "✓" if os.environ.get("SALESFORCE_USERNAME") else "✗", 
        "SALESFORCE_PASSWORD": "✓" if os.environ.get("SALESFORCE_PASSWORD") else "✗"
    }
    
    health_response = {
        "status": "healthy",
        "message": "Proxy-lite API is running",
        "environment_variables": env_status,
        "endpoints": {
            "POST /run_proxy_task": "Execute Salesforce automation tasks (requires 'task' and 'url' parameters)",
            "GET /health": "Health check and status",
            "GET /": "API information"
        },
        "supported_pages": [
            "Warranty Lifecycle Management",
            "Account Forecasting Settings", 
            "Sales Agreements",
            "Account Manager Targets",
            "Any Salesforce Setup page"
        ],
        "timestamp": datetime.now().isoformat()
    }
    
    return jsonify(health_response)

if __name__ == '__main__':
    if not os.environ.get("GEMINI_API_KEY"):
        logger.error("GEMINI_API_KEY environment variable is not set. Please set it for local testing.")
    logger.info("Starting Flask development server on 0.0.0.0:6101...")
    app.run(host='0.0.0.0', port=6101, debug=True)