diff --git a/agents/optimizer.md b/agents/optimizer.md index 02806dc..fb2178a 100644 --- a/agents/optimizer.md +++ b/agents/optimizer.md @@ -60,6 +60,52 @@ You are a thin-wrapper agent that runs the codeflash CLI to optimize Python and Follow these steps in order: +### 0. Check API Key + +Before anything else, check if a Codeflash API key is available: + +```bash +[ -n "${CODEFLASH_API_KEY:-}" ] && [[ "${CODEFLASH_API_KEY}" == cf-* ]] && printf 'env:ok\n' || printf 'env:missing\n'; grep -l 'CODEFLASH_API_KEY.*cf-' ~/.zshrc ~/.bashrc ~/.profile ~/.kshrc ~/.cshrc ~/codeflash_env.ps1 ~/codeflash_env.bat 2>/dev/null || true +``` + +If the output contains `env:ok`, proceed to Step 1. + +If the output contains `env:missing` but a shell RC file path was listed, source that file to load the key: + +```bash +source ~/.zshrc # or whichever file had the key +``` + +Then proceed to Step 1. + +If **no API key is found anywhere**, run the OAuth login script: + +```bash +bash "$(dirname "$0")/../scripts/oauth-login.sh" +``` + +The script has three possible outcomes: + +1. **Exit 0** — login succeeded, API key saved to shell RC. Source the RC file to load it, then proceed to Step 1. + +2. **Exit 2** — headless environment detected (SSH, CI, no display). The script outputs a JSON line like: + ```json + {"headless":true,"url":"https://app.codeflash.ai/...","state_file":"/tmp/codeflash-oauth-state-XXXXXX.json"} + ``` + In this case: + - Parse the `url` and `state_file` from the JSON output. + - **Ask the user** to visit the URL in their browser, complete authentication, and paste the authorization code they receive. + - Once the user provides the code, run: + ```bash + bash "$(dirname "$0")/../scripts/oauth-login.sh" --exchange-code + ``` + - If that succeeds (exit 0), source the shell RC file and proceed to Step 1. + +3. **Exit 1** — login failed. Stop and inform the user that a Codeflash API key is required. They can get one manually at https://app.codeflash.ai/app/apikeys and set it with: + ``` + export CODEFLASH_API_KEY="cf-your-key-here" + ``` + ### 1. Locate Project Configuration Walk upward from the current working directory to the git repository root (`git rev-parse --show-toplevel`) looking for a project configuration file. Check for both `pyproject.toml` (Python) and `package.json` (JavaScript/TypeScript) at each directory level. Use the **first** (closest to CWD) file found. diff --git a/scripts/oauth-login.sh b/scripts/oauth-login.sh new file mode 100755 index 0000000..68e3aed --- /dev/null +++ b/scripts/oauth-login.sh @@ -0,0 +1,405 @@ +#!/usr/bin/env bash +# OAuth PKCE login flow for Codeflash. +# Opens browser for authentication, exchanges code for API key, +# and saves it to the user's shell RC file. +# +# Usage: +# ./oauth-login.sh # Full browser flow +# ./oauth-login.sh --exchange-code STATE_FILE CODE # Complete headless flow +# +# Exit codes: +# 0 = success (API key saved) +# 1 = error +# 2 = headless mode (remote URL printed to stdout, state saved to temp file) + +set -euo pipefail + +CFWEBAPP_BASE_URL="https://app.codeflash.ai" +TOKEN_URL="${CFWEBAPP_BASE_URL}/codeflash/auth/oauth/token" +CLIENT_ID="cf-cli-app" +TIMEOUT=180 + +# --- Detect if a browser can be launched (matches codeflash's should_attempt_browser_launch) --- +can_open_browser() { + # CI/CD or non-interactive environments + if [ -n "${CI:-}" ] || [ "${DEBIAN_FRONTEND:-}" = "noninteractive" ]; then + return 1 + fi + + # Text-only browsers + local browser_env="${BROWSER:-}" + case "$browser_env" in + www-browser|lynx|links|w3m|elinks|links2) return 1 ;; + esac + + local is_ssh="false" + if [ -n "${SSH_CONNECTION:-}" ]; then + is_ssh="true" + fi + + # Linux: require a display server + if [ "$(uname -s)" = "Linux" ]; then + if [ -z "${DISPLAY:-}" ] && [ -z "${WAYLAND_DISPLAY:-}" ] && [ -z "${MIR_SOCKET:-}" ]; then + return 1 + fi + fi + + # SSH on non-Linux (e.g. macOS remote) — no browser + if [ "$is_ssh" = "true" ] && [ "$(uname -s)" != "Linux" ]; then + return 1 + fi + + return 0 +} + +# --- Save API key to shell RC (matches codeflash's shell_utils.py logic) --- +save_api_key() { + local api_key="$1" + + if [ "${OS:-}" = "Windows_NT" ] || [[ "$(uname -s 2>/dev/null)" == MINGW* ]]; then + # Windows: use dedicated codeflash env files (same as codeflash CLI) + if [ -n "${PSMODULEPATH:-}" ]; then + RC_FILE="$HOME/codeflash_env.ps1" + EXPORT_LINE="\$env:CODEFLASH_API_KEY = \"${api_key}\"" + REMOVE_PATTERN='^\$env:CODEFLASH_API_KEY' + else + RC_FILE="$HOME/codeflash_env.bat" + EXPORT_LINE="set CODEFLASH_API_KEY=\"${api_key}\"" + REMOVE_PATTERN='^set CODEFLASH_API_KEY=' + fi + else + # Unix: use shell RC file (same mapping as codeflash CLI) + SHELL_NAME=$(basename "${SHELL:-/bin/bash}") + case "$SHELL_NAME" in + zsh) RC_FILE="$HOME/.zshrc" ;; + ksh) RC_FILE="$HOME/.kshrc" ;; + csh|tcsh) RC_FILE="$HOME/.cshrc" ;; + dash) RC_FILE="$HOME/.profile" ;; + *) RC_FILE="$HOME/.bashrc" ;; + esac + EXPORT_LINE="export CODEFLASH_API_KEY=\"${api_key}\"" + REMOVE_PATTERN='^export CODEFLASH_API_KEY=' + fi + + # Remove any existing CODEFLASH_API_KEY lines and append the new one + if [ -f "$RC_FILE" ]; then + CLEANED=$(grep -v "$REMOVE_PATTERN" "$RC_FILE" || true) + printf '%s\n' "$CLEANED" > "$RC_FILE" + fi + printf '%s\n' "$EXPORT_LINE" >> "$RC_FILE" + + # Also export for the current session + export CODEFLASH_API_KEY="$api_key" +} + +# --- Exchange code for token and save --- +exchange_and_save() { + local auth_code="$1" + local code_verifier="$2" + local redirect_uri="$3" + + TOKEN_RESPONSE=$(curl -s -X POST "$TOKEN_URL" \ + -H "Content-Type: application/json" \ + -d "{ + \"grant_type\": \"authorization_code\", + \"code\": \"${auth_code}\", + \"code_verifier\": \"${code_verifier}\", + \"redirect_uri\": \"${redirect_uri}\", + \"client_id\": \"${CLIENT_ID}\" + }") + + API_KEY=$(printf '%s' "$TOKEN_RESPONSE" | python3 -c "import sys,json; print(json.load(sys.stdin).get('access_token',''))" 2>/dev/null || true) + + if [ -z "$API_KEY" ] || [[ ! "$API_KEY" == cf-* ]]; then + exit 1 + fi + + save_api_key "$API_KEY" +} + +# =================================================================== +# Mode: --exchange-code STATE_FILE CODE +# Complete a headless flow using a previously saved PKCE state file. +# =================================================================== +if [ "${1:-}" = "--exchange-code" ]; then + STATE_FILE="${2:-}" + MANUAL_CODE="${3:-}" + + if [ -z "$STATE_FILE" ] || [ -z "$MANUAL_CODE" ] || [ ! -f "$STATE_FILE" ]; then + exit 1 + fi + + # Read saved state + CODE_VERIFIER=$(python3 -c "import json; print(json.load(open('${STATE_FILE}')).get('code_verifier',''))" 2>/dev/null || true) + REMOTE_REDIRECT=$(python3 -c "import json; print(json.load(open('${STATE_FILE}')).get('remote_redirect_uri',''))" 2>/dev/null || true) + + rm -f "$STATE_FILE" + + if [ -z "$CODE_VERIFIER" ] || [ -z "$REMOTE_REDIRECT" ]; then + exit 1 + fi + + exchange_and_save "$MANUAL_CODE" "$CODE_VERIFIER" "$REMOTE_REDIRECT" + exit 0 +fi + +# =================================================================== +# Mode: Full OAuth flow (default) +# =================================================================== + +# --- PKCE pair --- +CODE_VERIFIER=$(openssl rand -base64 48 | tr -d '=/+\n' | head -c 64) +CODE_CHALLENGE=$(printf '%s' "$CODE_VERIFIER" | openssl dgst -sha256 -binary | openssl base64 -A | tr '+/' '-_' | tr -d '=') + +# --- State --- +STATE=$(openssl rand -hex 16) + +# --- Find a free port --- +PORT=$(python3 -c "import socket; s=socket.socket(); s.bind(('',0)); print(s.getsockname()[1]); s.close()") + +LOCAL_REDIRECT_URI="http://localhost:${PORT}/callback" +REMOTE_REDIRECT_URI="${CFWEBAPP_BASE_URL}/codeflash/auth/callback" +ENCODED_LOCAL_REDIRECT=$(python3 -c "import urllib.parse; print(urllib.parse.quote('${LOCAL_REDIRECT_URI}'))") +ENCODED_REMOTE_REDIRECT=$(python3 -c "import urllib.parse; print(urllib.parse.quote('${REMOTE_REDIRECT_URI}'))") + +AUTH_PARAMS="response_type=code&client_id=${CLIENT_ID}&code_challenge=${CODE_CHALLENGE}&code_challenge_method=sha256&state=${STATE}" +LOCAL_AUTH_URL="${CFWEBAPP_BASE_URL}/codeflash/auth?${AUTH_PARAMS}&redirect_uri=${ENCODED_LOCAL_REDIRECT}" +REMOTE_AUTH_URL="${CFWEBAPP_BASE_URL}/codeflash/auth?${AUTH_PARAMS}&redirect_uri=${ENCODED_REMOTE_REDIRECT}" + +# --- Headless detection --- +if ! can_open_browser; then + # Save PKCE state so --exchange-code can complete the flow later + HEADLESS_STATE_FILE=$(mktemp /tmp/codeflash-oauth-state-XXXXXX.json) + python3 -c " +import json +json.dump({ + 'code_verifier': '${CODE_VERIFIER}', + 'remote_redirect_uri': '${REMOTE_REDIRECT_URI}', + 'state': '${STATE}' +}, open('${HEADLESS_STATE_FILE}', 'w')) +" + # Output JSON for Claude to parse — this is the ONLY stdout in headless mode + printf '{"headless":true,"url":"%s","state_file":"%s"}\n' "$REMOTE_AUTH_URL" "$HEADLESS_STATE_FILE" + exit 2 +fi + +# --- Temp file for callback result --- +RESULT_FILE=$(mktemp /tmp/codeflash-oauth-XXXXXX.json) +trap 'rm -f "$RESULT_FILE"' EXIT + +# --- Start local callback server with Codeflash-styled pages --- +export PORT STATE RESULT_FILE TIMEOUT +python3 - "$PORT" "$STATE" "$RESULT_FILE" << 'PYEOF' & +import http.server, urllib.parse, json, sys, threading + +port = int(sys.argv[1]) +state = sys.argv[2] +result_file = sys.argv[3] + +STYLE = ( + ":root{" + "--bg:hsl(0,0%,99%);--fg:hsl(222.2,84%,4.9%);--card:hsl(0,0%,100%);" + "--card-fg:hsl(222.2,84%,4.9%);--primary:hsl(38,100%,63%);" + "--muted-fg:hsl(41,8%,46%);--border:hsl(41,30%,90%);" + "--destructive:hsl(0,84.2%,60.2%);--destructive-fg:#fff;" + "--success:hsl(142,76%,36%)}" + "html.dark{" + "--bg:hsl(0,6%,5%);--fg:#fff;--card:hsl(0,3%,11%);" + "--card-fg:#fff;--primary:hsl(38,100%,63%);" + "--muted-fg:hsl(48,20%,65%);--border:hsl(48,20%,25%);" + "--destructive:hsl(0,62.8%,30.6%)}" + "*{margin:0;padding:0;box-sizing:border-box}" + "body{font-family:-apple-system,BlinkMacSystemFont,'Segoe UI',system-ui,sans-serif;" + "background:var(--bg);color:var(--fg);min-height:100vh;" + "display:flex;align-items:center;justify-content:center;padding:20px;position:relative}" + "body::before{content:'';position:fixed;inset:0;" + "background:linear-gradient(to bottom,hsl(38,100%,63%,.1),hsl(38,100%,63%,.05),transparent);" + "pointer-events:none;z-index:0}" + "body::after{content:'';position:fixed;inset:0;" + "background-image:linear-gradient(to right,rgba(128,128,128,.03) 1px,transparent 1px)," + "linear-gradient(to bottom,rgba(128,128,128,.03) 1px,transparent 1px);" + "background-size:24px 24px;pointer-events:none;z-index:0}" + ".ctr{max-width:420px;width:100%;position:relative;z-index:1}" + ".logo-ctr{display:flex;justify-content:center;margin-bottom:48px}" + ".logo{height:40px;width:auto}" + ".ll{display:block}.ld{display:none}" + "html.dark .ll{display:none}html.dark .ld{display:block}" + ".card{background:var(--card);color:var(--card-fg);border:1px solid var(--border);" + "border-radius:16px;box-shadow:0 10px 15px -3px rgba(0,0,0,.1),0 4px 6px -2px rgba(0,0,0,.05);" + "padding:48px;animation:fadeIn .3s ease-out forwards}" + "@keyframes fadeIn{from{opacity:0;transform:translateY(10px)}to{opacity:1;transform:translateY(0)}}" + ".ic{width:48px;height:48px;background:hsl(38,100%,63%,.1);border-radius:12px;" + "display:flex;align-items:center;justify-content:center;margin:0 auto 24px}" + ".spinner{width:24px;height:24px;border:2px solid var(--border);" + "border-top-color:var(--primary);border-radius:50%;animation:spin .8s linear infinite}" + "@keyframes spin{to{transform:rotate(360deg)}}" + ".si{width:64px;height:64px;background:hsl(142,76%,36%,.1);border-radius:12px;" + "display:flex;align-items:center;justify-content:center;margin:0 auto 24px}" + ".sc{width:32px;height:32px;stroke:hsl(142,76%,36%)}" + "h1{font-size:24px;font-weight:600;margin:0 0 12px;color:var(--card-fg);text-align:center}" + "p{color:var(--muted-fg);margin:0;font-size:14px;line-height:1.5;text-align:center}" + ".eb{background:var(--destructive);color:var(--destructive-fg);" + "padding:14px 18px;border-radius:8px;margin-top:24px;font-size:14px;line-height:1.5;text-align:center}" + "@media(max-width:480px){.card{padding:32px 24px}h1{font-size:20px}.logo{height:32px}}" +) + +LOGO = ( + '
' + '' + '' + '
' +) + +# Pre-built static HTML fragments (no user data) +SUCCESS_FRAG = ( + '
' + '' + '

Success!

' + '

Authentication completed. You can close this window and return to your terminal.

' +) + +ERR_ICON_FRAG = ( + '
' + '' + '' + '' + '' + '

Authentication Failed

' +) + + +def loading_page(): + # Inline the static fragments as JS string constants so the polling + # script only inserts pre-defined trusted HTML, never user data. + success_js = json.dumps(SUCCESS_FRAG) + err_icon_js = json.dumps(ERR_ICON_FRAG) + return ( + '' + '' + 'CodeFlash Authentication' + f'' + '' + '' + f'
{LOGO}' + '
' + '
' + '

Authenticating

' + '

Please wait while we verify your credentials...

' + '
' + '' + ) + + +class H(http.server.BaseHTTPRequestHandler): + server_version = "CFHTTP" + + def do_GET(self): + p = urllib.parse.urlparse(self.path) + + if p.path == "/status": + self.send_response(200) + self.send_header("Content-type", "application/json") + self.send_header("Access-Control-Allow-Origin", "*") + self.end_headers() + body = json.dumps({ + "success": self.server.token_error is None and self.server.auth_code is not None, + "error": self.server.token_error, + }) + self.wfile.write(body.encode()) + return + + if p.path != "/callback": + self.send_response(404) + self.end_headers() + return + + params = urllib.parse.parse_qs(p.query) + code = params.get("code", [None])[0] + recv_state = params.get("state", [None])[0] + error = params.get("error", [None])[0] + + if error or not code or recv_state != state: + self.server.token_error = error or "state_mismatch" + else: + self.server.auth_code = code + with open(result_file, "w") as f: + json.dump({"code": code}, f) + + self.send_response(200) + self.send_header("Content-type", "text/html") + self.end_headers() + self.wfile.write(loading_page().encode()) + + def log_message(self, *a): + pass + + +httpd = http.server.HTTPServer(("localhost", port), H) +httpd.auth_code = None +httpd.token_error = None +httpd.serve_forever() +PYEOF +SERVER_PID=$! + +# --- Open browser (macOS, Linux, WSL) --- +if [[ "$(uname)" == "Darwin" ]]; then + open "$LOCAL_AUTH_URL" 2>/dev/null || true +elif command -v wslview >/dev/null 2>&1; then + wslview "$LOCAL_AUTH_URL" 2>/dev/null || true +elif command -v xdg-open >/dev/null 2>&1; then + xdg-open "$LOCAL_AUTH_URL" 2>/dev/null || true +elif command -v cmd.exe >/dev/null 2>&1; then + cmd.exe /c start "" "$LOCAL_AUTH_URL" 2>/dev/null || true +fi + +# --- Wait for callback --- +WAITED=0 +while [ ! -s "$RESULT_FILE" ] && [ "$WAITED" -lt "$TIMEOUT" ]; do + sleep 1 + WAITED=$((WAITED + 1)) + if ! kill -0 "$SERVER_PID" 2>/dev/null; then + break + fi +done + +if [ ! -s "$RESULT_FILE" ]; then + kill "$SERVER_PID" 2>/dev/null || true + wait "$SERVER_PID" 2>/dev/null || true + exit 1 +fi + +# --- Parse callback result --- +AUTH_CODE=$(python3 -c "import json; print(json.load(open('${RESULT_FILE}')).get('code',''))" 2>/dev/null || true) + +if [ -z "$AUTH_CODE" ]; then + kill "$SERVER_PID" 2>/dev/null || true + wait "$SERVER_PID" 2>/dev/null || true + exit 1 +fi + +# --- Exchange code for token --- +exchange_and_save "$AUTH_CODE" "$CODE_VERIFIER" "$LOCAL_REDIRECT_URI" + +# Give the browser a moment to poll /status and see success, then shut down +sleep 2 +kill "$SERVER_PID" 2>/dev/null || true +wait "$SERVER_PID" 2>/dev/null || true \ No newline at end of file diff --git a/scripts/suggest-optimize.sh b/scripts/suggest-optimize.sh index bb8bcba..52fbbc5 100755 --- a/scripts/suggest-optimize.sh +++ b/scripts/suggest-optimize.sh @@ -91,6 +91,35 @@ echo "$COMMIT_HASH" >> "$SEEN_MARKER" # --- From here on, we know there are new commits to optimize --- +# --- Check if CODEFLASH_API_KEY is available --- +OAUTH_SCRIPT="$(dirname "$0")/oauth-login.sh" + +has_api_key() { + # Check env var + if [ -n "${CODEFLASH_API_KEY:-}" ] && [[ "${CODEFLASH_API_KEY}" == cf-* ]]; then + return 0 + fi + # Check Unix shell RC files + for rc in "$HOME/.zshrc" "$HOME/.bashrc" "$HOME/.profile" "$HOME/.kshrc" "$HOME/.cshrc"; do + if [ -f "$rc" ] && grep -q '^export CODEFLASH_API_KEY="cf-' "$rc" 2>/dev/null; then + return 0 + fi + done + # Check Windows-specific files (PowerShell / CMD, matching codeflash CLI) + for rc in "$HOME/codeflash_env.ps1" "$HOME/codeflash_env.bat"; do + if [ -f "$rc" ] && grep -q 'CODEFLASH_API_KEY.*cf-' "$rc" 2>/dev/null; then + return 0 + fi + done + return 1 +} + +LOGIN_STEP="" +if ! has_api_key; then + LOGIN_STEP=" +- Run \`${OAUTH_SCRIPT}\` to log in to Codeflash. If it exits with code 0, the key is saved. If it exits with code 2 (headless environment), parse the JSON output for the \`url\` and \`state_file\`, ask the user to visit the URL and provide the authorization code, then run \`${OAUTH_SCRIPT} --exchange-code \` to complete the login." +fi + # Walk from $PWD upward to $REPO_ROOT looking for project config. # Sets: PROJECT_TYPE, PROJECT_DIR, PROJECT_CONFIG_PATH, PROJECT_CONFIGURED detect_project() { @@ -291,7 +320,7 @@ Then, add \`Bash(*codeflash*)\` to the \`permissions.allow\` array in \`$SETTING VENV_MSG="Python files were changed in a recent commit, but no Python virtual environment was found. Set up codeflash to optimize the Python code for performance. Perform the following steps: - +${LOGIN_STEP} 1. Create a virtual environment by running: \`python3 -m venv .venv\` in \`$REPO_ROOT\`, then activate it by running: \`source $REPO_ROOT/.venv/bin/activate\`. 2. Install codeflash by running: \`${INSTALL_CMD}\`. @@ -328,7 +357,9 @@ if [ "$PROJECT_CONFIGURED" != "true" ]; then SETUP_MSG="Python files were changed in a recent commit, but codeflash is not configured yet. -Set up codeflash to optimize the Python code for performance. Automatically determine the following by exploring the project's actual directory structure using Glob and Read (do NOT present options to the user — discover the correct paths yourself): +Set up codeflash to optimize the Python code for performance: +${LOGIN_STEP} +Automatically determine the following by exploring the project's actual directory structure using Glob and Read (do NOT present options to the user — discover the correct paths yourself): 1. The relative path with respect to pyproject.toml to the root of the Python module — the module root is where tests import from. for example, if the module root is abc/ then the tests would be importing code as \`from abc import xyz\`. 2. The relative path with respect to pyproject.toml to the tests folder — look for existing directories named \`tests\` or \`test\`, or folders containing files that start with \`test_\`. If none exist, default to \`tests\`. @@ -359,7 +390,7 @@ fi # Codeflash is configured but not installed in the venv if [ "$CODEFLASH_INSTALLED" != "true" ]; then INSTALL_MSG="Python files were changed in a recent commit, but codeflash is not installed in the active virtual environment (\`$VIRTUAL_ENV\`). - +${LOGIN_STEP} Install codeflash to optimize the Python code for performance by running: ${INSTALL_CMD} @@ -377,6 +408,18 @@ Then, add \`Bash(*codeflash*)\` to the \`permissions.allow\` array in \`$SETTING exit 0 fi +# Check for API key before running codeflash +if ! has_api_key; then + LOGIN_MSG="Python files were changed in a recent commit, but no Codeflash API key was found. + +Run \`${OAUTH_SCRIPT}\` to log in to Codeflash. If it exits with code 0, the key is saved. If it exits with code 2 (headless environment), parse the JSON output for the \`url\` and \`state_file\`, ask the user to visit the URL and provide the authorization code, then run \`${OAUTH_SCRIPT} --exchange-code \` to complete the login. + +After login, run \`codeflash --subagent\` as a background Bash task (run_in_background: true) to optimize the code." + + jq -nc --arg reason "$LOGIN_MSG" '{"decision": "block", "reason": $reason}' + exit 0 +fi + # Instruct Claude to run codeflash as a background subagent if [ -n "$PROJECT_DIR" ] && [ "$PROJECT_DIR" != "$PWD" ]; then RUN_CMD="cd $PROJECT_DIR && $CODEFLASH_BIN --subagent" diff --git a/tests/test_suggest_optimize.bats b/tests/test_suggest_optimize.bats index 9746aca..2e74e80 100755 --- a/tests/test_suggest_optimize.bats +++ b/tests/test_suggest_optimize.bats @@ -367,7 +367,7 @@ setup() { create_pyproject true create_fake_venv "$REPO/.venv" - run run_hook false "VIRTUAL_ENV=$REPO/.venv" + run run_hook false "VIRTUAL_ENV=$REPO/.venv" "CODEFLASH_API_KEY=cf-test-key" assert_block assert_reason_contains "permissions.allow" assert_reason_contains 'Bash(*codeflash*)'