From 9a1e1523b09018243ef6dadd24fffc7079bb4836 Mon Sep 17 00:00:00 2001
From: HeshamHM28 <HeshamMohamedFathy@outlook.com>
Date: Tue, 17 Mar 2026 08:07:41 +0200
Subject: [PATCH 1/3] feat: add OAuth PKCE login flow for automatic API key
 authentication

---
 .claude-plugin/plugin.json  |   2 +-
 agents/optimizer.md         |  95 ++++++++++++
 scripts/oauth-login.sh      | 300 ++++++++++++++++++++++++++++++++++++
 scripts/suggest-optimize.sh |  43 +++++-
 4 files changed, 436 insertions(+), 4 deletions(-)
 create mode 100755 scripts/oauth-login.sh

diff --git a/.claude-plugin/plugin.json b/.claude-plugin/plugin.json
index 50b7c18..dae8b5c 100644
--- a/.claude-plugin/plugin.json
+++ b/.claude-plugin/plugin.json
@@ -1,7 +1,7 @@
 {
   "name": "codeflash",
   "description": "Run codeflash as a background agent to optimize Python code for performance",
-  "version": "0.1.10",
+  "version": "0.1.11",
   "author": {
     "name": "Codeflash",
     "url": "https://codeflash.ai"
diff --git a/agents/optimizer.md b/agents/optimizer.md
index 3c70002..572ea7c 100644
--- a/agents/optimizer.md
+++ b/agents/optimizer.md
@@ -42,6 +42,101 @@ You are a thin-wrapper agent that runs the codeflash CLI to optimize Python code
 
 Follow these steps in order:
 
+### 0. Check API Key
+
+Before anything else, check if a Codeflash API key is available:
+
+```bash
+echo "${CODEFLASH_API_KEY:-}"; grep '^export CODEFLASH_API_KEY="cf-' ~/.zshrc ~/.bashrc ~/.profile 2>/dev/null || true
+```
+
+If the environment variable is set and starts with `cf-`, proceed to Step 1.
+
+If the env var is empty but a key was found in a shell RC file, source that file to load it:
+
+```bash
+source ~/.zshrc  # or whichever file had the key
+```
+
+Then proceed to Step 1.
+
+If **no API key is found anywhere**, perform an OAuth PKCE login flow to authenticate the user. Run this script exactly:
+
+```bash
+set -euo pipefail
+CFWEBAPP_BASE_URL="https://app.codeflash.ai"
+TOKEN_URL="${CFWEBAPP_BASE_URL}/codeflash/auth/oauth/token"
+CLIENT_ID="cf-cli-app"
+
+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=$(openssl rand -hex 16)
+PORT=$(python3 -c "import socket; s=socket.socket(); s.bind(('',0)); print(s.getsockname()[1]); s.close()")
+REDIRECT_URI="http://localhost:${PORT}/callback"
+ENCODED_REDIRECT=$(python3 -c "import urllib.parse; print(urllib.parse.quote('${REDIRECT_URI}'))")
+AUTH_URL="${CFWEBAPP_BASE_URL}/codeflash/auth?response_type=code&client_id=${CLIENT_ID}&code_challenge=${CODE_CHALLENGE}&code_challenge_method=sha256&state=${STATE}&redirect_uri=${ENCODED_REDIRECT}"
+
+RESULT_FILE=$(mktemp /tmp/codeflash-oauth-XXXXXX.json)
+
+PORT=$PORT STATE=$STATE RESULT_FILE=$RESULT_FILE python3 -c "
+import http.server, urllib.parse, json, os, threading
+port, state, rf = int(os.environ['PORT']), os.environ['STATE'], os.environ['RESULT_FILE']
+class H(http.server.BaseHTTPRequestHandler):
+    def do_GET(self):
+        p = urllib.parse.urlparse(self.path)
+        if p.path != '/callback':
+            self.send_response(404); self.end_headers(); return
+        params = urllib.parse.parse_qs(p.query)
+        code, st, err = params.get('code',[None])[0], params.get('state',[None])[0], params.get('error',[None])[0]
+        self.send_response(200); self.send_header('Content-type','text/html'); self.end_headers()
+        if err or not code or st != state:
+            self.wfile.write(b'<h2>Authentication failed.</h2>')
+            with open(rf,'w') as f: json.dump({'error': err or 'state_mismatch'}, f)
+        else:
+            self.wfile.write(b'<h2>Success! You can close this window.</h2>')
+            with open(rf,'w') as f: json.dump({'code': code}, f)
+        threading.Thread(target=self.server.shutdown, daemon=True).start()
+    def log_message(self, *a): pass
+http.server.HTTPServer(('localhost', port), H).serve_forever()
+" &
+SERVER_PID=$!
+
+if [[ "$(uname)" == "Darwin" ]]; then open "$AUTH_URL" 2>/dev/null || true
+elif command -v xdg-open >/dev/null 2>&1; then xdg-open "$AUTH_URL" 2>/dev/null || true; fi
+
+echo "Opening browser for Codeflash login..."
+echo "If the browser didn't open, visit: $AUTH_URL"
+
+WAITED=0
+while [ ! -s "$RESULT_FILE" ] && [ "$WAITED" -lt 180 ]; do sleep 1; WAITED=$((WAITED + 1)); done
+kill "$SERVER_PID" 2>/dev/null || true; wait "$SERVER_PID" 2>/dev/null || true
+
+AUTH_CODE=$(python3 -c "import json; print(json.load(open('${RESULT_FILE}')).get('code',''))" 2>/dev/null || true)
+rm -f "$RESULT_FILE"
+
+if [ -z "$AUTH_CODE" ]; then echo "Login failed"; exit 1; fi
+
+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=$(echo "$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 echo "Token exchange failed"; exit 1; fi
+
+SHELL_NAME=$(basename "${SHELL:-/bin/bash}")
+case "$SHELL_NAME" in zsh) RC="$HOME/.zshrc";; *) RC="$HOME/.bashrc";; esac
+[ -f "$RC" ] && grep -v '^export CODEFLASH_API_KEY=' "$RC" > "${RC}.tmp" && mv "${RC}.tmp" "$RC"
+echo "export CODEFLASH_API_KEY=\"${API_KEY}\"" >> "$RC"
+export CODEFLASH_API_KEY="$API_KEY"
+echo "Login successful! API key saved to $RC"
+```
+
+After the login script completes, **source the shell RC file** to load the key, then proceed to Step 1.
+
+If the login fails or times out, **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 `pyproject.toml`. Use the **first** (closest to CWD) file found. Record:
diff --git a/scripts/oauth-login.sh b/scripts/oauth-login.sh
new file mode 100755
index 0000000..74c5a31
--- /dev/null
+++ b/scripts/oauth-login.sh
@@ -0,0 +1,300 @@
+#!/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
+# Outputs the API key on success, exits non-zero on failure.
+
+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
+
+# --- 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()")
+
+REDIRECT_URI="http://localhost:${PORT}/callback"
+ENCODED_REDIRECT=$(python3 -c "import urllib.parse; print(urllib.parse.quote('${REDIRECT_URI}'))")
+
+AUTH_URL="${CFWEBAPP_BASE_URL}/codeflash/auth?response_type=code&client_id=${CLIENT_ID}&code_challenge=${CODE_CHALLENGE}&code_challenge_method=sha256&state=${STATE}&redirect_uri=${ENCODED_REDIRECT}"
+
+# --- 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 = (
+    '<div class="logo-ctr">'
+    '<img src="https://app.codeflash.ai/images/codeflash_light.svg" alt="CodeFlash" class="logo ll"/>'
+    '<img src="https://app.codeflash.ai/images/codeflash_darkmode.svg" alt="CodeFlash" class="logo ld"/>'
+    '</div>'
+)
+
+# Pre-built static HTML fragments (no user data)
+SUCCESS_FRAG = (
+    '<div class="si"><svg class="sc" fill="none" stroke="currentColor" viewBox="0 0 24 24">'
+    '<path stroke-linecap="round" stroke-linejoin="round" stroke-width="2" d="M5 13l4 4L19 7"/>'
+    '</svg></div><h1>Success!</h1>'
+    '<p>Authentication completed. You can close this window and return to your terminal.</p>'
+)
+
+ERR_ICON_FRAG = (
+    '<div class="ic" style="background:hsl(0,84.2%,60.2%,.1)">'
+    '<svg width="24" height="24" fill="none" stroke="hsl(0,84.2%,60.2%)" viewBox="0 0 24 24">'
+    '<circle cx="12" cy="12" r="10" stroke-width="2"/>'
+    '<line x1="12" y1="8" x2="12" y2="12" stroke-width="2" stroke-linecap="round"/>'
+    '<line x1="12" y1="16" x2="12.01" y2="16" stroke-width="2" stroke-linecap="round"/>'
+    '</svg></div><h1>Authentication Failed</h1>'
+)
+
+
+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 (
+        '<!DOCTYPE html><html><head><meta charset="UTF-8">'
+        '<meta name="viewport" content="width=device-width,initial-scale=1.0">'
+        '<title>CodeFlash Authentication</title>'
+        f'<style>{STYLE}</style>'
+        '<script>if(window.matchMedia&&window.matchMedia("(prefers-color-scheme:dark)").matches)'
+        '{document.documentElement.classList.add("dark")}</script>'
+        '</head><body>'
+        f'<div class="ctr">{LOGO}'
+        '<div class="card" id="c">'
+        '<div class="ic"><div class="spinner"></div></div>'
+        '<h1>Authenticating</h1>'
+        '<p>Please wait while we verify your credentials...</p>'
+        '</div></div>'
+        '<script>'
+        f'var SF={success_js},EI={err_icon_js};'
+        'var n=0,mx=60;'
+        'function ck(){'
+        'fetch("/status").then(function(r){return r.json()}).then(function(d){'
+        'var e=document.getElementById("c");'
+        'if(d.success){while(e.firstChild)e.removeChild(e.firstChild);'
+        'e.insertAdjacentHTML("beforeend",SF)}'
+        'else if(d.error){while(e.firstChild)e.removeChild(e.firstChild);'
+        'e.insertAdjacentHTML("beforeend",EI);'
+        'var b=document.createElement("div");b.className="eb";b.textContent=d.error;e.appendChild(b)}'
+        'else if(n<mx){n++;setTimeout(ck,500)}'
+        'else{while(e.firstChild)e.removeChild(e.firstChild);'
+        'e.insertAdjacentHTML("beforeend",EI);'
+        'var b2=document.createElement("div");b2.className="eb";'
+        'b2.textContent="Authentication timed out. Please try again.";e.appendChild(b2)}'
+        '}).catch(function(){if(n<mx){n++;setTimeout(ck,500)}})}'
+        'setTimeout(ck,1000);'
+        '</script></body></html>'
+    )
+
+
+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 ---
+if [[ "$(uname)" == "Darwin" ]]; then
+  open "$AUTH_URL" 2>/dev/null || true
+elif command -v xdg-open >/dev/null 2>&1; then
+  xdg-open "$AUTH_URL" 2>/dev/null || true
+fi
+
+echo "Opening browser for Codeflash login..."
+echo ""
+echo "If the browser didn't open, visit this URL:"
+echo "$AUTH_URL"
+echo ""
+
+# --- 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
+  echo "Authentication timed out after ${TIMEOUT}s." >&2
+  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
+  echo "No authorization code received." >&2
+  exit 1
+fi
+
+# --- Exchange code for token ---
+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=$(echo "$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
+  kill "$SERVER_PID" 2>/dev/null || true
+  wait "$SERVER_PID" 2>/dev/null || true
+  echo "Failed to obtain API key from token exchange." >&2
+  echo "Response: $TOKEN_RESPONSE" >&2
+  exit 1
+fi
+
+# 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
+
+# --- Save to shell RC ---
+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 any existing CODEFLASH_API_KEY lines and append the new one
+if [ -f "$RC_FILE" ]; then
+  CLEANED=$(grep -v '^export CODEFLASH_API_KEY=' "$RC_FILE" || true)
+  echo "$CLEANED" > "$RC_FILE"
+fi
+echo "$EXPORT_LINE" >> "$RC_FILE"
+
+# Also export for the current session
+export CODEFLASH_API_KEY="$API_KEY"
+
+echo "Login successful! API key saved to ${RC_FILE}"
+echo "Restart your shell or run: source ${RC_FILE}"
+
+# Output the key so callers can capture it
+echo "$API_KEY"
\ No newline at end of file
diff --git a/scripts/suggest-optimize.sh b/scripts/suggest-optimize.sh
index b3cf829..d2ebfd3 100755
--- a/scripts/suggest-optimize.sh
+++ b/scripts/suggest-optimize.sh
@@ -81,6 +81,29 @@ echo "$COMMIT_HASH" >> "$SEEN_MARKER"
 
 # --- From here on, we know there are new Python 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 shell RC files
+  for rc in "$HOME/.zshrc" "$HOME/.bashrc" "$HOME/.profile"; do
+    if [ -f "$rc" ] && grep -q '^export 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 (this will open a browser for authentication and save the API key)."
+fi
+
 # Walk from $PWD upward to $REPO_ROOT looking for pyproject.toml.
 # Sets: PYPROJECT_DIR, PYPROJECT_PATH, PYPROJECT_CONFIGURED
 find_pyproject() {
@@ -170,7 +193,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}\`.
@@ -207,7 +230,9 @@ if [ "$PYPROJECT_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\`.
 
@@ -238,7 +263,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}
@@ -256,6 +281,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. This will open a browser for authentication and save the API key to your shell config.
+
+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 "$PYPROJECT_DIR" ] && [ "$PYPROJECT_DIR" != "$PWD" ]; then
   RUN_CMD="cd $PYPROJECT_DIR && $CODEFLASH_BIN --subagent"

From 750e9770e2b135f2b78455708d6d1bfcf6a6c5b1 Mon Sep 17 00:00:00 2001
From: HeshamHM28 <HeshamMohamedFathy@outlook.com>
Date: Thu, 19 Mar 2026 01:27:04 +0200
Subject: [PATCH 2/3] fix headless server

---
 agents/optimizer.md         | 101 ++++------------
 scripts/oauth-login.sh      | 235 ++++++++++++++++++++++++++----------
 scripts/suggest-optimize.sh |  14 ++-
 3 files changed, 206 insertions(+), 144 deletions(-)

diff --git a/agents/optimizer.md b/agents/optimizer.md
index 3c64f11..fb2178a 100644
--- a/agents/optimizer.md
+++ b/agents/optimizer.md
@@ -65,12 +65,12 @@ Follow these steps in order:
 Before anything else, check if a Codeflash API key is available:
 
 ```bash
-echo "${CODEFLASH_API_KEY:-}"; grep '^export CODEFLASH_API_KEY="cf-' ~/.zshrc ~/.bashrc ~/.profile 2>/dev/null || true
+[ -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 environment variable is set and starts with `cf-`, proceed to Step 1.
+If the output contains `env:ok`, proceed to Step 1.
 
-If the env var is empty but a key was found in a shell RC file, source that file to load it:
+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
@@ -78,82 +78,33 @@ source ~/.zshrc  # or whichever file had the key
 
 Then proceed to Step 1.
 
-If **no API key is found anywhere**, perform an OAuth PKCE login flow to authenticate the user. Run this script exactly:
+If **no API key is found anywhere**, run the OAuth login script:
 
 ```bash
-set -euo pipefail
-CFWEBAPP_BASE_URL="https://app.codeflash.ai"
-TOKEN_URL="${CFWEBAPP_BASE_URL}/codeflash/auth/oauth/token"
-CLIENT_ID="cf-cli-app"
-
-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=$(openssl rand -hex 16)
-PORT=$(python3 -c "import socket; s=socket.socket(); s.bind(('',0)); print(s.getsockname()[1]); s.close()")
-REDIRECT_URI="http://localhost:${PORT}/callback"
-ENCODED_REDIRECT=$(python3 -c "import urllib.parse; print(urllib.parse.quote('${REDIRECT_URI}'))")
-AUTH_URL="${CFWEBAPP_BASE_URL}/codeflash/auth?response_type=code&client_id=${CLIENT_ID}&code_challenge=${CODE_CHALLENGE}&code_challenge_method=sha256&state=${STATE}&redirect_uri=${ENCODED_REDIRECT}"
-
-RESULT_FILE=$(mktemp /tmp/codeflash-oauth-XXXXXX.json)
-
-PORT=$PORT STATE=$STATE RESULT_FILE=$RESULT_FILE python3 -c "
-import http.server, urllib.parse, json, os, threading
-port, state, rf = int(os.environ['PORT']), os.environ['STATE'], os.environ['RESULT_FILE']
-class H(http.server.BaseHTTPRequestHandler):
-    def do_GET(self):
-        p = urllib.parse.urlparse(self.path)
-        if p.path != '/callback':
-            self.send_response(404); self.end_headers(); return
-        params = urllib.parse.parse_qs(p.query)
-        code, st, err = params.get('code',[None])[0], params.get('state',[None])[0], params.get('error',[None])[0]
-        self.send_response(200); self.send_header('Content-type','text/html'); self.end_headers()
-        if err or not code or st != state:
-            self.wfile.write(b'<h2>Authentication failed.</h2>')
-            with open(rf,'w') as f: json.dump({'error': err or 'state_mismatch'}, f)
-        else:
-            self.wfile.write(b'<h2>Success! You can close this window.</h2>')
-            with open(rf,'w') as f: json.dump({'code': code}, f)
-        threading.Thread(target=self.server.shutdown, daemon=True).start()
-    def log_message(self, *a): pass
-http.server.HTTPServer(('localhost', port), H).serve_forever()
-" &
-SERVER_PID=$!
-
-if [[ "$(uname)" == "Darwin" ]]; then open "$AUTH_URL" 2>/dev/null || true
-elif command -v xdg-open >/dev/null 2>&1; then xdg-open "$AUTH_URL" 2>/dev/null || true; fi
-
-echo "Opening browser for Codeflash login..."
-echo "If the browser didn't open, visit: $AUTH_URL"
-
-WAITED=0
-while [ ! -s "$RESULT_FILE" ] && [ "$WAITED" -lt 180 ]; do sleep 1; WAITED=$((WAITED + 1)); done
-kill "$SERVER_PID" 2>/dev/null || true; wait "$SERVER_PID" 2>/dev/null || true
-
-AUTH_CODE=$(python3 -c "import json; print(json.load(open('${RESULT_FILE}')).get('code',''))" 2>/dev/null || true)
-rm -f "$RESULT_FILE"
-
-if [ -z "$AUTH_CODE" ]; then echo "Login failed"; exit 1; fi
-
-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=$(echo "$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 echo "Token exchange failed"; exit 1; fi
-
-SHELL_NAME=$(basename "${SHELL:-/bin/bash}")
-case "$SHELL_NAME" in zsh) RC="$HOME/.zshrc";; *) RC="$HOME/.bashrc";; esac
-[ -f "$RC" ] && grep -v '^export CODEFLASH_API_KEY=' "$RC" > "${RC}.tmp" && mv "${RC}.tmp" "$RC"
-echo "export CODEFLASH_API_KEY=\"${API_KEY}\"" >> "$RC"
-export CODEFLASH_API_KEY="$API_KEY"
-echo "Login successful! API key saved to $RC"
+bash "$(dirname "$0")/../scripts/oauth-login.sh"
 ```
 
-After the login script completes, **source the shell RC file** to load the key, then proceed to Step 1.
-
-If the login fails or times out, **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"
-```
+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 <state_file> <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
 
diff --git a/scripts/oauth-login.sh b/scripts/oauth-login.sh
index 74c5a31..68e3aed 100755
--- a/scripts/oauth-login.sh
+++ b/scripts/oauth-login.sh
@@ -3,8 +3,14 @@
 # Opens browser for authentication, exchanges code for API key,
 # and saves it to the user's shell RC file.
 #
-# Usage: ./oauth-login.sh
-# Outputs the API key on success, exits non-zero on failure.
+# 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
 
@@ -13,6 +19,134 @@ 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 '=')
@@ -23,10 +157,31 @@ 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()")
 
-REDIRECT_URI="http://localhost:${PORT}/callback"
-ENCODED_REDIRECT=$(python3 -c "import urllib.parse; print(urllib.parse.quote('${REDIRECT_URI}'))")
-
-AUTH_URL="${CFWEBAPP_BASE_URL}/codeflash/auth?response_type=code&client_id=${CLIENT_ID}&code_challenge=${CODE_CHALLENGE}&code_challenge_method=sha256&state=${STATE}&redirect_uri=${ENCODED_REDIRECT}"
+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)
@@ -205,19 +360,17 @@ httpd.serve_forever()
 PYEOF
 SERVER_PID=$!
 
-# --- Open browser ---
+# --- Open browser (macOS, Linux, WSL) ---
 if [[ "$(uname)" == "Darwin" ]]; then
-  open "$AUTH_URL" 2>/dev/null || true
+  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 "$AUTH_URL" 2>/dev/null || true
+  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
 
-echo "Opening browser for Codeflash login..."
-echo ""
-echo "If the browser didn't open, visit this URL:"
-echo "$AUTH_URL"
-echo ""
-
 # --- Wait for callback ---
 WAITED=0
 while [ ! -s "$RESULT_FILE" ] && [ "$WAITED" -lt "$TIMEOUT" ]; do
@@ -231,7 +384,6 @@ done
 if [ ! -s "$RESULT_FILE" ]; then
   kill "$SERVER_PID" 2>/dev/null || true
   wait "$SERVER_PID" 2>/dev/null || true
-  echo "Authentication timed out after ${TIMEOUT}s." >&2
   exit 1
 fi
 
@@ -241,60 +393,13 @@ AUTH_CODE=$(python3 -c "import json; print(json.load(open('${RESULT_FILE}')).get
 if [ -z "$AUTH_CODE" ]; then
   kill "$SERVER_PID" 2>/dev/null || true
   wait "$SERVER_PID" 2>/dev/null || true
-  echo "No authorization code received." >&2
   exit 1
 fi
 
 # --- Exchange code for token ---
-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=$(echo "$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
-  kill "$SERVER_PID" 2>/dev/null || true
-  wait "$SERVER_PID" 2>/dev/null || true
-  echo "Failed to obtain API key from token exchange." >&2
-  echo "Response: $TOKEN_RESPONSE" >&2
-  exit 1
-fi
+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
-
-# --- Save to shell RC ---
-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 any existing CODEFLASH_API_KEY lines and append the new one
-if [ -f "$RC_FILE" ]; then
-  CLEANED=$(grep -v '^export CODEFLASH_API_KEY=' "$RC_FILE" || true)
-  echo "$CLEANED" > "$RC_FILE"
-fi
-echo "$EXPORT_LINE" >> "$RC_FILE"
-
-# Also export for the current session
-export CODEFLASH_API_KEY="$API_KEY"
-
-echo "Login successful! API key saved to ${RC_FILE}"
-echo "Restart your shell or run: source ${RC_FILE}"
-
-# Output the key so callers can capture it
-echo "$API_KEY"
\ No newline at end of file
+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 cb59c87..52fbbc5 100755
--- a/scripts/suggest-optimize.sh
+++ b/scripts/suggest-optimize.sh
@@ -99,19 +99,25 @@ has_api_key() {
   if [ -n "${CODEFLASH_API_KEY:-}" ] && [[ "${CODEFLASH_API_KEY}" == cf-* ]]; then
     return 0
   fi
-  # Check shell RC files
-  for rc in "$HOME/.zshrc" "$HOME/.bashrc" "$HOME/.profile"; do
+  # 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 (this will open a browser for authentication and save the API key)."
+- 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 <state_file> <code>\` to complete the login."
 fi
 
 # Walk from $PWD upward to $REPO_ROOT looking for project config.
@@ -406,7 +412,7 @@ fi
 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. This will open a browser for authentication and save the API key to your shell config.
+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 <state_file> <code>\` to complete the login.
 
 After login, run \`codeflash --subagent\` as a background Bash task (run_in_background: true) to optimize the code."
 

From 5f02b07910543a0880e3f80fa3f4ddd151e11bd8 Mon Sep 17 00:00:00 2001
From: Ubuntu <ubuntu@ip-172-31-39-200.ec2.internal>
Date: Wed, 18 Mar 2026 23:32:07 +0000
Subject: [PATCH 3/3] fix: set mock API key in auto-allow test to bypass OAuth
 check

Test 27 was failing because the new has_api_key check exits before
reaching the permissions.allow message. Providing a mock API key
lets the test exercise the auto-allow logic as intended.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
---
 tests/test_suggest_optimize.bats | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

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*)'