From 9a36d36c1017d341417e4820f1161e239644bc76 Mon Sep 17 00:00:00 2001
From: Yoav Katz <katz@il.ibm.com>
Date: Tue, 28 Jul 2026 11:08:50 +0300
Subject: [PATCH 1/3] Refactor Keycloak direct-access-grants into shared
 helper; harden eval port-forwards; add cluster/dry flags to ibac comparison

- Extract the duplicated "enable Direct Access Grants for the rossoctl
  client" logic from deploy-agent.sh and deploy-benchmark.sh into a new
  libsh/keycloak-direct-access.sh helper (enable_direct_access_grants).
  Resolves master-realm admin creds from env vars, the
  keycloak-initial-admin secret, then admin/admin defaults.
- delete-all-deployments.sh now calls the helper too: its token requests
  use grant_type=password against the rossoctl client, which requires
  Direct Access Grants to be enabled first.
- evaluate-benchmark.sh: detect and clean up stale kubectl port-forwards
  holding a local port before binding, refusing to touch non-kubectl
  processes; capture port-forward output to surface the real error on
  startup failure.
- run-ibac-comparison.sh: add --kind/--openshift/--in-cluster cluster
  flags and a --dry run mode, forwarded to deploy-and-evaluate.sh.

Signed-off-by: Yoav Katz <katz@il.ibm.com>
---
 exgentic_a2a_runner/delete-all-deployments.sh |  6 ++
 exgentic_a2a_runner/deploy-agent.sh           | 32 +-------
 exgentic_a2a_runner/deploy-benchmark.sh       | 52 +------------
 exgentic_a2a_runner/evaluate-benchmark.sh     | 63 ++++++++++++++-
 .../libsh/keycloak-direct-access.sh           | 75 ++++++++++++++++++
 exgentic_a2a_runner/run-ibac-comparison.sh    | 76 +++++++++++++++++--
 6 files changed, 216 insertions(+), 88 deletions(-)
 create mode 100644 exgentic_a2a_runner/libsh/keycloak-direct-access.sh

diff --git a/exgentic_a2a_runner/delete-all-deployments.sh b/exgentic_a2a_runner/delete-all-deployments.sh
index 2d664fc..687192e 100755
--- a/exgentic_a2a_runner/delete-all-deployments.sh
+++ b/exgentic_a2a_runner/delete-all-deployments.sh
@@ -100,6 +100,8 @@ export CLUSTER_MODE INGRESS_DOMAIN
 source "$SCRIPT_DIR/libsh/urls.sh"
 # shellcheck source=libsh/check-kubectl-context.sh
 source "$SCRIPT_DIR/libsh/check-kubectl-context.sh"
+# shellcheck source=libsh/keycloak-direct-access.sh
+source "$SCRIPT_DIR/libsh/keycloak-direct-access.sh"
 check_kubectl_context
 
 ROSSOCTL_API="$(rossoctl_api_url)"
@@ -117,6 +119,10 @@ echo ""
 # Step 1: Get Keycloak authentication token
 echo "Step 1: Getting Keycloak authentication token..."
 
+# The token requests below use grant_type=password against the rossoctl client,
+# which requires Direct Access Grants to be enabled on that client.
+enable_direct_access_grants
+
 if [ "$KEYCLOAK_PASSWORD" = "unknown" ]; then
     echo "Step 1.5: Attempting to fetch Keycloak password from cluster..."
     ROSSOCTL_PASSWORD=$("$KUBECTL_BIN" get secret rossoctl-test-user -n keycloak -o jsonpath='{.data.password}' 2>/dev/null | base64 -d 2>/dev/null || echo "")
diff --git a/exgentic_a2a_runner/deploy-agent.sh b/exgentic_a2a_runner/deploy-agent.sh
index abab018..411929a 100755
--- a/exgentic_a2a_runner/deploy-agent.sh
+++ b/exgentic_a2a_runner/deploy-agent.sh
@@ -240,6 +240,8 @@ source "$SCRIPT_DIR/libsh/urls.sh"
 KUBECTL_BIN="${KUBECTL_BIN:-kubectl}"
 # shellcheck source=libsh/check-kubectl-context.sh
 source "$SCRIPT_DIR/libsh/check-kubectl-context.sh"
+# shellcheck source=libsh/keycloak-direct-access.sh
+source "$SCRIPT_DIR/libsh/keycloak-direct-access.sh"
 check_kubectl_context
 
 ROSSOCTL_API="$(rossoctl_api_url)"
@@ -331,35 +333,7 @@ fi
 
 # Step 2: Enable Direct Access Grants for rossoctl client if needed
 echo "Step 2: Enabling Direct Access Grants for rossoctl client..."
-
-# Get admin token first (use "admin" password for master realm)
-ADMIN_TOKEN_RESPONSE=$(curl -s -X POST "$KEYCLOAK_API/realms/master/protocol/openid-connect/token" \
-    -H "Content-Type: application/x-www-form-urlencoded" \
-    -d "username=admin" \
-    -d "password=admin" \
-    -d "grant_type=password" \
-    -d "client_id=admin-cli" 2>/dev/null || echo "TOKEN_ERROR")
-
-if [ "$ADMIN_TOKEN_RESPONSE" != "TOKEN_ERROR" ]; then
-    ADMIN_TOKEN=$(echo "$ADMIN_TOKEN_RESPONSE" | grep -o '"access_token":"[^"]*"' | sed 's/"access_token":"\([^"]*\)"/\1/')
-    
-    if [ -n "$ADMIN_TOKEN" ]; then
-        # Get rossoctl client configuration
-        CLIENT_CONFIG=$(curl -s "$KEYCLOAK_API/admin/realms/rossoctl/clients?clientId=rossoctl" \
-            -H "Authorization: Bearer $ADMIN_TOKEN" 2>/dev/null)
-        
-        CLIENT_ID=$(echo "$CLIENT_CONFIG" | grep -o '"id":"[^"]*"' | head -1 | sed 's/"id":"\([^"]*\)"/\1/')
-        
-        if [ -n "$CLIENT_ID" ]; then
-            # Enable direct access grants
-            curl -s -X PUT "$KEYCLOAK_API/admin/realms/rossoctl/clients/$CLIENT_ID" \
-                -H "Authorization: Bearer $ADMIN_TOKEN" \
-                -H "Content-Type: application/json" \
-                -d '{"directAccessGrantsEnabled": true}' >/dev/null 2>&1
-            echo "✓ Direct access grants enabled for rossoctl client"
-        fi
-    fi
-fi
+enable_direct_access_grants
 
 echo ""
 
diff --git a/exgentic_a2a_runner/deploy-benchmark.sh b/exgentic_a2a_runner/deploy-benchmark.sh
index 1b50d29..b81e499 100755
--- a/exgentic_a2a_runner/deploy-benchmark.sh
+++ b/exgentic_a2a_runner/deploy-benchmark.sh
@@ -122,6 +122,8 @@ source "$SCRIPT_DIR_BENCH/libsh/urls.sh"
 KUBECTL_BIN="${KUBECTL_BIN:-kubectl}"
 # shellcheck source=libsh/check-kubectl-context.sh
 source "$SCRIPT_DIR_BENCH/libsh/check-kubectl-context.sh"
+# shellcheck source=libsh/keycloak-direct-access.sh
+source "$SCRIPT_DIR_BENCH/libsh/keycloak-direct-access.sh"
 check_kubectl_context
 
 # Default to Exgentic registry, can be overridden with environment variable
@@ -195,55 +197,7 @@ fi
 
 # Step 4: Enable Direct Access Grants for rossoctl client if needed
 echo "Step 4: Enabling Direct Access Grants for rossoctl client..."
-
-# Resolve master-realm admin credentials: prefer env vars, fall back to the
-# keycloak-initial-admin secret (RHBK operator), then defaults.
-KEYCLOAK_ADMIN_USERNAME="${KEYCLOAK_ADMIN_USERNAME:-}"
-KEYCLOAK_ADMIN_PASSWORD="${KEYCLOAK_ADMIN_PASSWORD:-}"
-if [ -z "$KEYCLOAK_ADMIN_USERNAME" ] || [ -z "$KEYCLOAK_ADMIN_PASSWORD" ]; then
-    KC_ADMIN_USERNAME=$(kubectl get secret keycloak-initial-admin -n keycloak \
-        -o jsonpath='{.data.username}' 2>/dev/null | base64 -d 2>/dev/null || true)
-    KC_ADMIN_PASSWORD=$(kubectl get secret keycloak-initial-admin -n keycloak \
-        -o jsonpath='{.data.password}' 2>/dev/null | base64 -d 2>/dev/null || true)
-    KEYCLOAK_ADMIN_USERNAME="${KEYCLOAK_ADMIN_USERNAME:-${KC_ADMIN_USERNAME:-admin}}"
-    KEYCLOAK_ADMIN_PASSWORD="${KEYCLOAK_ADMIN_PASSWORD:-${KC_ADMIN_PASSWORD:-admin}}"
-fi
-
-ADMIN_TOKEN_RESPONSE=$(curl -s -X POST "$KEYCLOAK_API/realms/master/protocol/openid-connect/token" \
-    -H "Content-Type: application/x-www-form-urlencoded" \
-    -d "username=${KEYCLOAK_ADMIN_USERNAME}" \
-    -d "password=${KEYCLOAK_ADMIN_PASSWORD}" \
-    -d "grant_type=password" \
-    -d "client_id=admin-cli" 2>/dev/null) || true
-
-ADMIN_TOKEN=$(echo "$ADMIN_TOKEN_RESPONSE" | grep -o '"access_token":"[^"]*"' | sed 's/"access_token":"\([^"]*\)"/\1/')
-if [ -z "$ADMIN_TOKEN" ]; then
-    echo "Error: Could not obtain master-realm admin token from Keycloak"
-    echo "  Response: $ADMIN_TOKEN_RESPONSE"
-    echo "  Set KEYCLOAK_ADMIN_PASSWORD in your .env if the master realm admin password is not 'admin'."
-    exit 1
-fi
-
-CLIENT_CONFIG=$(curl -s "$KEYCLOAK_API/admin/realms/rossoctl/clients?clientId=rossoctl" \
-    -H "Authorization: Bearer $ADMIN_TOKEN" 2>/dev/null)
-CLIENT_ID=$(echo "$CLIENT_CONFIG" | grep -o '"id":"[^"]*"' | head -1 | sed 's/"id":"\([^"]*\)"/\1/')
-if [ -z "$CLIENT_ID" ]; then
-    echo "Error: Could not find rossoctl client ID in Keycloak"
-    echo "  Response: $CLIENT_CONFIG"
-    exit 1
-fi
-
-PUT_CODE=$(curl -s -o /tmp/kc_put_response.txt -w "%{http_code}" \
-    -X PUT "$KEYCLOAK_API/admin/realms/rossoctl/clients/$CLIENT_ID" \
-    -H "Authorization: Bearer $ADMIN_TOKEN" \
-    -H "Content-Type: application/json" \
-    -d '{"directAccessGrantsEnabled": true}' 2>/dev/null) || PUT_CODE="000"
-if [ "$PUT_CODE" != "204" ] && [ "$PUT_CODE" != "200" ]; then
-    echo "Error: Failed to enable direct access grants for rossoctl client (HTTP $PUT_CODE)"
-    echo "  Response: $(cat /tmp/kc_put_response.txt 2>/dev/null)"
-    exit 1
-fi
-echo "✓ Direct access grants enabled for rossoctl client"
+enable_direct_access_grants
 
 echo ""
 
diff --git a/exgentic_a2a_runner/evaluate-benchmark.sh b/exgentic_a2a_runner/evaluate-benchmark.sh
index 02d550b..e420f41 100755
--- a/exgentic_a2a_runner/evaluate-benchmark.sh
+++ b/exgentic_a2a_runner/evaluate-benchmark.sh
@@ -265,15 +265,68 @@ PROMETHEUS_LOCAL_PORT="${PROMETHEUS_LOCAL_PORT:-9191}"
 PROMETHEUS_NAMESPACE="istio-system"
 PROMETHEUS_SERVICE="prometheus"
 
+PF_OTEL_COLLECTOR_LOG=""
+PF_PROMETHEUS_LOG=""
+
+# Ensure a local TCP port is free before we try to bind a port-forward to it.
+# A common failure mode is a stale (or suspended) kubectl port-forward left over
+# from a previous run still holding the port; detect that case and clear it.
+free_local_port() {
+    local port="$1"
+    local label="$2"
+
+    if ! command -v lsof >/dev/null 2>&1; then
+        return 0
+    fi
+
+    local pids
+    pids="$(lsof -tiTCP:"$port" -sTCP:LISTEN -n -P 2>/dev/null | sort -u)"
+    if [ -z "$pids" ]; then
+        return 0
+    fi
+
+    # Only auto-kill leftover kubectl port-forwards; refuse to touch anything else.
+    local stale_pids=""
+    local pid
+    for pid in $pids; do
+        local cmd
+        cmd="$(ps -o command= -p "$pid" 2>/dev/null)"
+        if echo "$cmd" | grep -q "port-forward"; then
+            stale_pids="$stale_pids $pid"
+        fi
+    done
+
+    if [ -z "$stale_pids" ]; then
+        echo "Error: local port $port (for $label) is in use by a non-kubectl process:"
+        # shellcheck disable=SC2086
+        ps -o pid=,command= -p $pids 2>/dev/null | sed 's/^/    /'
+        echo "  Free the port or set the corresponding *_LOCAL_PORT env var, then retry."
+        exit 1
+    fi
+
+    echo "  Found stale kubectl port-forward on port $port ($label); cleaning up..."
+    # shellcheck disable=SC2086
+    kill $stale_pids 2>/dev/null || true
+    sleep 1
+    # shellcheck disable=SC2086
+    kill -9 $stale_pids 2>/dev/null || true
+}
+
+# Start a kubectl port-forward, capturing output so we can surface the real
+# error if it dies. Sets the named PID/LOG variables via the caller.
 if [ "$CLUSTER_MODE" != "in-cluster" ]; then
     if [ "$MLFLOW_ENABLED" = "true" ]; then
         echo "Starting port-forward for OTEL collector (traces -> MLflow)..."
-        "$KUBECTL_BIN" port-forward -n $OTEL_COLLECTOR_NAMESPACE svc/$OTEL_COLLECTOR_SERVICE ${OTEL_COLLECTOR_LOCAL_PORT}:4317 >/dev/null 2>&1 &
+        free_local_port "$OTEL_COLLECTOR_LOCAL_PORT" "OTEL collector"
+        PF_OTEL_COLLECTOR_LOG="$(mktemp -t otel-pf.XXXXXX)"
+        "$KUBECTL_BIN" port-forward -n $OTEL_COLLECTOR_NAMESPACE svc/$OTEL_COLLECTOR_SERVICE ${OTEL_COLLECTOR_LOCAL_PORT}:4317 >"$PF_OTEL_COLLECTOR_LOG" 2>&1 &
         PF_OTEL_COLLECTOR_PID=$!
     fi
 
     echo "Starting port-forward for Prometheus..."
-    "$KUBECTL_BIN" port-forward -n $PROMETHEUS_NAMESPACE svc/$PROMETHEUS_SERVICE ${PROMETHEUS_LOCAL_PORT}:9090 >/dev/null 2>&1 &
+    free_local_port "$PROMETHEUS_LOCAL_PORT" "Prometheus"
+    PF_PROMETHEUS_LOG="$(mktemp -t prom-pf.XXXXXX)"
+    "$KUBECTL_BIN" port-forward -n $PROMETHEUS_NAMESPACE svc/$PROMETHEUS_SERVICE ${PROMETHEUS_LOCAL_PORT}:9090 >"$PF_PROMETHEUS_LOG" 2>&1 &
     PF_PROMETHEUS_PID=$!
 
     if [ "$MLFLOW_ENABLED" = "true" ]; then
@@ -282,6 +335,10 @@ if [ "$CLUSTER_MODE" != "in-cluster" ]; then
 
         if ! ps -p $PF_OTEL_COLLECTOR_PID > /dev/null; then
             echo "Error: OTEL collector port-forward failed to start"
+            if [ -n "$PF_OTEL_COLLECTOR_LOG" ] && [ -s "$PF_OTEL_COLLECTOR_LOG" ]; then
+                echo "  kubectl output:"
+                sed 's/^/    /' "$PF_OTEL_COLLECTOR_LOG"
+            fi
             exit 1
         fi
 
@@ -306,6 +363,8 @@ cleanup() {
         echo "Stopping Prometheus port-forward..."
         kill $PF_PROMETHEUS_PID 2>/dev/null || true
     fi
+    [ -n "$PF_OTEL_COLLECTOR_LOG" ] && rm -f "$PF_OTEL_COLLECTOR_LOG"
+    [ -n "$PF_PROMETHEUS_LOG" ] && rm -f "$PF_PROMETHEUS_LOG"
     echo "Done."
 }
 
diff --git a/exgentic_a2a_runner/libsh/keycloak-direct-access.sh b/exgentic_a2a_runner/libsh/keycloak-direct-access.sh
new file mode 100644
index 0000000..aa71cae
--- /dev/null
+++ b/exgentic_a2a_runner/libsh/keycloak-direct-access.sh
@@ -0,0 +1,75 @@
+#!/bin/bash
+# Enable Keycloak "Direct Access Grants" (the OAuth password grant) for the
+# rossoctl client. deploy-agent.sh, deploy-benchmark.sh and delete-all-deployments.sh
+# all authenticate with grant_type=password against the rossoctl client, which only
+# works when the client has directAccessGrantsEnabled=true. This helper flips that
+# flag via the Keycloak admin API so the password-grant token request in each script
+# succeeds.
+#
+# Must be sourced after KEYCLOAK_API is set. Master-realm admin credentials are
+# resolved in priority order:
+#   1. KEYCLOAK_ADMIN_USERNAME / KEYCLOAK_ADMIN_PASSWORD env vars
+#   2. the keycloak-initial-admin secret (RHBK operator)
+#   3. admin / admin defaults
+#
+# On any failure (unreachable admin API, missing client, rejected PUT) it prints a
+# diagnostic and exits 1 — the password grant is a hard prerequisite for every caller.
+enable_direct_access_grants() {
+    echo "Enabling Direct Access Grants for rossoctl client..."
+
+    # Resolve master-realm admin credentials: prefer env vars, fall back to the
+    # keycloak-initial-admin secret (RHBK operator), then defaults.
+    local admin_user="${KEYCLOAK_ADMIN_USERNAME:-}"
+    local admin_pass="${KEYCLOAK_ADMIN_PASSWORD:-}"
+    if [ -z "$admin_user" ] || [ -z "$admin_pass" ]; then
+        local kc_admin_user kc_admin_pass
+        kc_admin_user=$("$KUBECTL_BIN" get secret keycloak-initial-admin -n keycloak \
+            -o jsonpath='{.data.username}' 2>/dev/null | base64 -d 2>/dev/null || true)
+        kc_admin_pass=$("$KUBECTL_BIN" get secret keycloak-initial-admin -n keycloak \
+            -o jsonpath='{.data.password}' 2>/dev/null | base64 -d 2>/dev/null || true)
+        admin_user="${admin_user:-${kc_admin_user:-admin}}"
+        admin_pass="${admin_pass:-${kc_admin_pass:-admin}}"
+    fi
+
+    # Get master-realm admin token.
+    local admin_token_response admin_token
+    admin_token_response=$(curl -s -X POST "$KEYCLOAK_API/realms/master/protocol/openid-connect/token" \
+        -H "Content-Type: application/x-www-form-urlencoded" \
+        -d "username=${admin_user}" \
+        -d "password=${admin_pass}" \
+        -d "grant_type=password" \
+        -d "client_id=admin-cli" 2>/dev/null) || true
+
+    admin_token=$(echo "$admin_token_response" | grep -o '"access_token":"[^"]*"' | sed 's/"access_token":"\([^"]*\)"/\1/')
+    if [ -z "$admin_token" ]; then
+        echo "Error: Could not obtain master-realm admin token from Keycloak" >&2
+        echo "  Response: $admin_token_response" >&2
+        echo "  Set KEYCLOAK_ADMIN_PASSWORD in your .env if the master realm admin password is not 'admin'." >&2
+        exit 1
+    fi
+
+    # Look up the rossoctl client's internal id.
+    local client_config client_id
+    client_config=$(curl -s "$KEYCLOAK_API/admin/realms/rossoctl/clients?clientId=rossoctl" \
+        -H "Authorization: Bearer $admin_token" 2>/dev/null)
+    client_id=$(echo "$client_config" | grep -o '"id":"[^"]*"' | head -1 | sed 's/"id":"\([^"]*\)"/\1/')
+    if [ -z "$client_id" ]; then
+        echo "Error: Could not find rossoctl client ID in Keycloak" >&2
+        echo "  Response: $client_config" >&2
+        exit 1
+    fi
+
+    # Enable direct access grants (the password grant).
+    local put_code
+    put_code=$(curl -s -o /tmp/kc_put_response.txt -w "%{http_code}" \
+        -X PUT "$KEYCLOAK_API/admin/realms/rossoctl/clients/$client_id" \
+        -H "Authorization: Bearer $admin_token" \
+        -H "Content-Type: application/json" \
+        -d '{"directAccessGrantsEnabled": true}' 2>/dev/null) || put_code="000"
+    if [ "$put_code" != "204" ] && [ "$put_code" != "200" ]; then
+        echo "Error: Failed to enable direct access grants for rossoctl client (HTTP $put_code)" >&2
+        echo "  Response: $(cat /tmp/kc_put_response.txt 2>/dev/null)" >&2
+        exit 1
+    fi
+    echo "✓ Direct access grants enabled for rossoctl client"
+}
diff --git a/exgentic_a2a_runner/run-ibac-comparison.sh b/exgentic_a2a_runner/run-ibac-comparison.sh
index 7ff71f5..8bc1254 100755
--- a/exgentic_a2a_runner/run-ibac-comparison.sh
+++ b/exgentic_a2a_runner/run-ibac-comparison.sh
@@ -15,6 +15,10 @@
 #   --max-parallel-sessions N      Number of concurrent evaluation sessions (default: 1)
 #   --plugin-preset PRESET         Plugin preset for the first run:
 #                                  auth-only | ibac-only | full (default: ibac-only)
+#   --kind                         Target a local Kind cluster (default)
+#   --openshift DOMAIN             Target an OpenShift cluster with the given ingress domain
+#   --in-cluster                   Running as a Kubernetes Job inside the cluster
+#   --dry                          Dry run mode - print commands without executing them
 #   -h, --help                     Show this help and exit
 #
 # The judge is configured from the OPENAI_API_BASE / OPENAI_API_KEY environment
@@ -24,6 +28,7 @@
 #   ./run-ibac-comparison.sh
 #   ./run-ibac-comparison.sh --model gcp/gemini-3-flash-preview --benchmark gsm8k --max-tasks 10 --max-parallel-sessions 1
 #   ./run-ibac-comparison.sh --plugin-preset full
+#   ./run-ibac-comparison.sh --plugin-preset auth-only --benchmark gsm8k --openshift apps.mycluster.example.com
 
 set -euo pipefail
 
@@ -36,6 +41,11 @@ AGENT="tool_calling"
 MAX_TASKS=10
 MAX_PARALLEL_SESSIONS=1
 PLUGIN_PRESET="ibac-only"
+DRY_RUN="false"
+
+# Cluster-mode flag forwarded verbatim to deploy-and-evaluate.sh (which validates
+# it). Empty means "unset" — the sub-scripts then apply their own default.
+CLUSTER_FLAG=()
 
 while [[ $# -gt 0 ]]; do
     case "$1" in
@@ -63,8 +73,24 @@ while [[ $# -gt 0 ]]; do
             PLUGIN_PRESET="$2"
             shift 2
             ;;
+        --kind)
+            CLUSTER_FLAG=(--kind)
+            shift
+            ;;
+        --openshift)
+            CLUSTER_FLAG=(--openshift "$2")
+            shift 2
+            ;;
+        --in-cluster)
+            CLUSTER_FLAG=(--in-cluster)
+            shift
+            ;;
+        --dry)
+            DRY_RUN="true"
+            shift
+            ;;
         -h|--help)
-            sed -n '2,26p' "${BASH_SOURCE[0]}" | sed 's/^# \{0,1\}//'
+            sed -n '2,31p' "${BASH_SOURCE[0]}" | sed 's/^# \{0,1\}//'
             exit 0
             ;;
         *)
@@ -119,11 +145,35 @@ echo "Plugin preset:          $PLUGIN_PRESET"
 echo "Experiment id:          $EXPERIMENT_ID"
 echo "Plugin experiment:      $EXPERIMENT_PLUGIN"
 echo "Baseline experiment:    $EXPERIMENT_BASE"
+if [ ${#CLUSTER_FLAG[@]} -gt 0 ]; then
+    echo "Cluster mode:           ${CLUSTER_FLAG[*]}"
+else
+    echo "Cluster mode:           <default>"
+fi
+echo "Dry run:                $DRY_RUN"
 echo "=========================================="
 
+# In dry-run mode, print each command instead of executing it, and forward
+# --dry to deploy-and-evaluate.sh so its own steps are printed rather than run.
+DRY_FLAG=()
+if [ "$DRY_RUN" = "true" ]; then
+    DRY_FLAG=(--dry)
+fi
+
+run_step() {
+    if [ "$DRY_RUN" = "true" ]; then
+        echo "[DRY RUN] Would execute:"
+        printf '%q ' "$@"
+        echo ""
+        echo ""
+    else
+        "$@"
+    fi
+}
+
 # Run 1: with the selected plugin preset.
-"$SCRIPT_DIR/delete-all-deployments.sh"
-env IBAC_JUDGE_ENDPOINT="$OPENAI_API_BASE" \
+run_step "$SCRIPT_DIR/delete-all-deployments.sh" ${CLUSTER_FLAG[@]+"${CLUSTER_FLAG[@]}"}
+run_step env IBAC_JUDGE_ENDPOINT="$OPENAI_API_BASE" \
     IBAC_JUDGE_MODEL="$MODEL" \
     JUDGE_BEARER="$OPENAI_API_KEY" \
     "$SCRIPT_DIR/deploy-and-evaluate.sh" \
@@ -133,11 +183,13 @@ env IBAC_JUDGE_ENDPOINT="$OPENAI_API_BASE" \
         --max-tasks "$MAX_TASKS" \
         --max-parallel-sessions "$MAX_PARALLEL_SESSIONS" \
         --plugin-preset "$PLUGIN_PRESET" \
-        --experiment "$EXPERIMENT_PLUGIN"
+        --experiment "$EXPERIMENT_PLUGIN" \
+        ${CLUSTER_FLAG[@]+"${CLUSTER_FLAG[@]}"} \
+        ${DRY_FLAG[@]+"${DRY_FLAG[@]}"}
 
 # Run 2: baseline (no plugin preset).
-"$SCRIPT_DIR/delete-all-deployments.sh"
-env IBAC_JUDGE_ENDPOINT="$OPENAI_API_BASE" \
+run_step "$SCRIPT_DIR/delete-all-deployments.sh" ${CLUSTER_FLAG[@]+"${CLUSTER_FLAG[@]}"}
+run_step env IBAC_JUDGE_ENDPOINT="$OPENAI_API_BASE" \
     IBAC_JUDGE_MODEL="$MODEL" \
     JUDGE_BEARER="$OPENAI_API_KEY" \
     "$SCRIPT_DIR/deploy-and-evaluate.sh" \
@@ -146,7 +198,15 @@ env IBAC_JUDGE_ENDPOINT="$OPENAI_API_BASE" \
         --model "openai/$MODEL" \
         --max-tasks "$MAX_TASKS" \
         --max-parallel-sessions "$MAX_PARALLEL_SESSIONS" \
-        --experiment "$EXPERIMENT_BASE"
+        --experiment "$EXPERIMENT_BASE" \
+        ${CLUSTER_FLAG[@]+"${CLUSTER_FLAG[@]}"} \
+        ${DRY_FLAG[@]+"${DRY_FLAG[@]}"}
 
 # Compare the two runs.
-"$SCRIPT_DIR/analyze-run.sh" -c "${EXPERIMENT_PLUGIN},${EXPERIMENT_BASE}"
+run_step "$SCRIPT_DIR/analyze-run.sh" -c "${EXPERIMENT_PLUGIN},${EXPERIMENT_BASE}"
+
+if [ "$DRY_RUN" = "true" ]; then
+    echo "=========================================="
+    echo "✓ Dry run completed - no commands executed"
+    echo "=========================================="
+fi

From 2a434a94eea36941510207d63644c1dfb22c1144 Mon Sep 17 00:00:00 2001
From: Yoav Katz <katz@il.ibm.com>
Date: Tue, 28 Jul 2026 14:09:18 +0300
Subject: [PATCH 2/3] Fix analyze-run secret-mode auth: password grant against
 mlflow client

secret mode previously ran a client_credentials grant for the mlflow
service account, but mlflow-oidc-auth authorizes reads from its own user
DB (not the token's group claim), so the traces API returned 403. Switch
to a password (direct-access) grant for a real MLflow user (default:
admin), matching how the other deploy scripts authenticate.

- analyze-run.sh: obtain the token via a password grant built from
  keycloak_api_url() and the mlflow-oauth-secret client id/secret; resolve
  the user password from KEYCLOAK_PASSWORD or the rossoctl-test-user
  secret; enable Direct Access Grants on the mlflow client first.
- Default MLFLOW_LOCAL_PORT to 8085: 8080 is served by the kind ingress
  for keycloak.localtest.me (127.0.0.1), so binding there would shadow
  Keycloak and send the token request to MLflow.
- Detect the downloader's exit 75 (token rejected) via PIPESTATUS and
  print an actionable hint (log into the MLflow UI once), re-echoing the
  original invocation.
- download_mlflow_traces.py: raise MLflowAuthError on 401/403 and exit 75
  so the wrapper can distinguish "no MLflow user record yet" from a
  generic failure.
- keycloak-direct-access.sh: parameterize enable_direct_access_grants with
  an optional CLIENT_ID (defaults to rossoctl) so analyze-run can target
  the mlflow client.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Signed-off-by: Yoav Katz <katz@il.ibm.com>
---
 exgentic_a2a_runner/analyze-run.sh            | 114 ++++++++++++++----
 exgentic_a2a_runner/download_mlflow_traces.py |  39 +++++-
 .../libsh/keycloak-direct-access.sh           |  27 +++--
 3 files changed, 140 insertions(+), 40 deletions(-)

diff --git a/exgentic_a2a_runner/analyze-run.sh b/exgentic_a2a_runner/analyze-run.sh
index 7a53397..a85b7ed 100755
--- a/exgentic_a2a_runner/analyze-run.sh
+++ b/exgentic_a2a_runner/analyze-run.sh
@@ -27,10 +27,23 @@ WINDOW="${WINDOW:-3h}"
 MLFLOW_NAMESPACE="${MLFLOW_NAMESPACE:-}"
 MLFLOW_SERVICE="${MLFLOW_SERVICE:-}"
 MLFLOW_REMOTE_PORT="${MLFLOW_REMOTE_PORT:-}"
-MLFLOW_LOCAL_PORT="${MLFLOW_LOCAL_PORT:-8080}"
+# Local port for the MLflow port-forward. Must NOT be 8080: the kind ingress
+# serves keycloak.localtest.me (and other *.localtest.me hosts) on 8080, and
+# keycloak.localtest.me resolves to 127.0.0.1 — so binding the port-forward to
+# localhost:8080 would shadow Keycloak, and the secret-mode token request (a
+# password grant against Keycloak) would hit MLflow instead.
+MLFLOW_LOCAL_PORT="${MLFLOW_LOCAL_PORT:-8085}"
 MLFLOW_TLS="${MLFLOW_TLS:-}"
 MLFLOW_WORKSPACE="${MLFLOW_WORKSPACE:-}"
 AUTH_MODE="${AUTH_MODE:-}"
+# secret-mode (kind) auth: a password grant against the mlflow Keycloak client as
+# an MLflow user. MLflow's mlflow-oidc-auth authorizes reads from its own user DB
+# (where "admin" is seeded as a global admin), NOT from the token's Keycloak group
+# claim — so the user must be one MLflow knows. Defaults to admin; override with
+# MLFLOW_USER. The password defaults to the rossoctl-test-user secret in the
+# keycloak namespace (which holds admin's password); override with KEYCLOAK_PASSWORD.
+MLFLOW_USER="${MLFLOW_USER:-admin}"
+KEYCLOAK_PASSWORD="${KEYCLOAK_PASSWORD:-}"
 KUBECTL_BIN="${KUBECTL_BIN:-kubectl}"
 # `whoami -t` is an OpenShift (oc) extension, not a kubectl subcommand, so the
 # token command is separate from KUBECTL_BIN. Override with OC_BIN if needed.
@@ -93,6 +106,10 @@ EOF
     exit 1
 }
 
+# Capture the original invocation so the auth-failure hint can print the exact
+# command to re-run (the arg loop below consumes "$@" via shift).
+ORIGINAL_INVOCATION=("$0" "$@")
+
 while [[ $# -gt 0 ]]; do
     case $1 in
         -u|--url)           MLFLOW_URL="$2"; shift 2 ;;
@@ -238,6 +255,14 @@ source "$SCRIPT_DIR/libsh/check-kubectl-context.sh"
 check_kubectl_context
 echo ""
 
+# urls.sh provides keycloak_api_url (CLUSTER_MODE must be exported, done above);
+# keycloak-direct-access.sh provides enable_direct_access_grants. Both are used
+# by the secret-mode token flow (password grant against the mlflow client).
+# shellcheck source=libsh/urls.sh
+source "$SCRIPT_DIR/libsh/urls.sh"
+# shellcheck source=libsh/keycloak-direct-access.sh
+source "$SCRIPT_DIR/libsh/keycloak-direct-access.sh"
+
 # --- Helper functions ---
 
 OAUTH_TOKEN=""
@@ -277,10 +302,21 @@ cleanup_port_forward() {
     fi
 }
 
-# secret mode: rossoctl's client-credentials flow. Reads mlflow-oauth-secret and
-# execs into the MLflow pod to exchange it for an access token.
+# secret mode: password (direct-access) grant against the mlflow Keycloak client.
+#
+# Why not client_credentials (the previous approach): that mints a token for the
+# mlflow *service account*, which mlflow-oidc-auth does not grant experiment reads
+# to, so the traces API returns 403. mlflow-oidc-auth authorizes from its own user
+# DB, where "admin" is seeded as a global admin — so we obtain a token for a real
+# MLflow user (default: admin) instead. The mlflow client already carries a groups
+# protocol mapper and is the confidential client MLflow trusts.
+#
+# The client id/secret come from mlflow-oauth-secret; the user password defaults to
+# the rossoctl-test-user secret (admin's password). The token endpoint is built from
+# keycloak_api_url (the OIDC_TOKEN_URL in the secret points at the in-cluster
+# Keycloak service, which is not reachable from the laptop).
 get_token_from_secret() {
-    echo "Obtaining OAuth token via mlflow-oauth-secret..."
+    echo "Obtaining OAuth token via password grant against the mlflow client..."
 
     # Note: under `set -e`, a failing command substitution aborts the script
     # before the following `if` can run. Capture status explicitly so the
@@ -293,36 +329,43 @@ get_token_from_secret() {
         return 1
     fi
 
-    local client_id client_secret token_url
+    local client_id client_secret
     client_id=$(echo "$secret_json" | jq -r '.data["OIDC_CLIENT_ID"]' | base64 -d) || true
     client_secret=$(echo "$secret_json" | jq -r '.data["OIDC_CLIENT_SECRET"]' | base64 -d) || true
-    token_url=$(echo "$secret_json" | jq -r '.data["OIDC_TOKEN_URL"]' | base64 -d) || true
-
-    if [ -z "$client_id" ] || [ -z "$client_secret" ] || [ -z "$token_url" ]; then
-        echo "Error: Could not extract OAuth credentials from secret"
+    if [ -z "$client_id" ] || [ -z "$client_secret" ]; then
+        echo "Error: Could not extract OIDC client id/secret from mlflow-oauth-secret"
         return 1
     fi
 
-    local mlflow_pod
-    mlflow_pod=$("$KUBECTL_BIN" get pod -n "$MLFLOW_NAMESPACE" -l app=mlflow -o jsonpath='{.items[0].metadata.name}' 2>/dev/null) || true
-    if [ -z "$mlflow_pod" ]; then
-        echo "Error: Could not find MLflow pod"
+    # Resolve the MLflow user's password: explicit KEYCLOAK_PASSWORD wins, else the
+    # rossoctl-test-user secret (holds admin's password) in the keycloak namespace.
+    local user_password="$KEYCLOAK_PASSWORD"
+    if [ -z "$user_password" ]; then
+        user_password=$("$KUBECTL_BIN" get secret rossoctl-test-user -n keycloak \
+            -o jsonpath='{.data.password}' 2>/dev/null | base64 -d 2>/dev/null || true)
+    fi
+    if [ -z "$user_password" ]; then
+        echo "Error: Could not resolve a password for MLflow user '$MLFLOW_USER'"
+        echo "Hint: set KEYCLOAK_PASSWORD, or confirm the rossoctl-test-user secret exists in the keycloak namespace"
         return 1
     fi
 
+    # The mlflow client needs Direct Access Grants enabled for the password grant.
+    local KEYCLOAK_API
+    KEYCLOAK_API="$(keycloak_api_url)"
+    export KEYCLOAK_API
+    enable_direct_access_grants "$client_id"
+
+    local token_url="$KEYCLOAK_API/realms/rossoctl/protocol/openid-connect/token"
+    echo "Requesting token for user '$MLFLOW_USER' (client '$client_id')..."
     local token_response
-    token_response=$("$KUBECTL_BIN" exec -n "$MLFLOW_NAMESPACE" "$mlflow_pod" -- \
-        python3 -c "
-import urllib.request, urllib.parse, json
-data = urllib.parse.urlencode({
-    'grant_type': 'client_credentials',
-    'client_id': '${client_id}',
-    'client_secret': '${client_secret}'
-}).encode()
-req = urllib.request.Request('${token_url}', data=data, headers={'Content-Type': 'application/x-www-form-urlencoded'})
-resp = urllib.request.urlopen(req)
-print(resp.read().decode())
-" 2>/dev/null) || true
+    token_response=$(curl -s -X POST "$token_url" \
+        -H "Content-Type: application/x-www-form-urlencoded" \
+        -d "grant_type=password" \
+        -d "client_id=${client_id}" \
+        -d "client_secret=${client_secret}" \
+        -d "username=${MLFLOW_USER}" \
+        -d "password=${user_password}" 2>/dev/null) || true
 
     OAUTH_TOKEN=$(echo "$token_response" | jq -r '.access_token' 2>/dev/null) || true
     if [ -z "$OAUTH_TOKEN" ] || [ "$OAUTH_TOKEN" = "null" ]; then
@@ -425,4 +468,25 @@ if [ -n "$COMPARE_EXPERIMENTS" ]; then
     PYTHON_ARGS="--compare"
 fi
 
+# download_mlflow_traces.py exits 75 when MLflow rejects the token (a valid
+# token still gets 403 until the user has logged into the MLflow UI once, which
+# is what populates mlflow-oidc-auth's permission DB). Capture the downloader's
+# status via PIPESTATUS so we can print an actionable hint instead of a raw
+# HTTP 403 traceback.
+set +e
 python3 "$SCRIPT_DIR/download_mlflow_traces.py" | python3 "$SCRIPT_DIR/analyze_traces.py" $PYTHON_ARGS
+DOWNLOAD_STATUS=${PIPESTATUS[0]}
+set -e
+
+if [ "$DOWNLOAD_STATUS" -eq 75 ]; then
+    echo ""
+    echo "MLflow authentication succeeded but access was denied."
+    echo "Log into the MLflow UI once (this registers your user with MLflow's"
+    echo "permission system), then re-run the analysis:"
+    echo ""
+    printf '    '; printf '%q ' "${ORIGINAL_INVOCATION[@]}"; echo
+    echo ""
+    exit 75
+fi
+
+exit "$DOWNLOAD_STATUS"
diff --git a/exgentic_a2a_runner/download_mlflow_traces.py b/exgentic_a2a_runner/download_mlflow_traces.py
index cf3fd42..6654081 100644
--- a/exgentic_a2a_runner/download_mlflow_traces.py
+++ b/exgentic_a2a_runner/download_mlflow_traces.py
@@ -46,6 +46,17 @@
 
 CUTOFF_MS = int(time.time() * 1000) - WINDOW_MS
 
+
+class MLflowAuthError(Exception):
+    """Raised when MLflow rejects the token (401/403).
+
+    A valid token can still be refused: mlflow-oidc-auth only populates its
+    user/permission DB during the interactive browser login (the OIDC
+    authorization-code callback). A raw bearer token from a fresh cluster maps
+    to no user record, so the traces API returns 403 until someone logs into
+    the MLflow UI once. The shell wrapper turns this into an actionable hint.
+    """
+
 # When talking to a port-forwarded reencrypt HTTPS endpoint the cert won't
 # validate against localhost, so allow opting out of verification.
 _SSL_CONTEXT = ssl._create_unverified_context() if MLFLOW_INSECURE_TLS else None
@@ -61,8 +72,16 @@ def mlflow_get(path: str) -> dict:
     if MLFLOW_WORKSPACE:
         headers["x-mlflow-workspace"] = MLFLOW_WORKSPACE
     req = urllib.request.Request(url, headers=headers)
-    with urllib.request.urlopen(req, timeout=30, context=_SSL_CONTEXT) as resp:
-        return json.loads(resp.read())
+    try:
+        with urllib.request.urlopen(req, timeout=30, context=_SSL_CONTEXT) as resp:
+            return json.loads(resp.read())
+    except urllib.error.HTTPError as e:
+        # A valid token can still be refused (403) when MLflow has no user
+        # record for it yet — see MLflowAuthError. Surface that distinctly so
+        # the wrapper can point the user at the UI-login fix.
+        if e.code in (401, 403):
+            raise MLflowAuthError(f"HTTP {e.code} {e.reason}") from e
+        raise
 
 
 def _is_long_enough(trace: dict) -> bool:
@@ -279,5 +298,19 @@ def main() -> int:
     return 0
 
 
+# Distinct exit code for "token rejected by MLflow" so the shell wrapper can
+# print the UI-login fix instead of a generic failure.
+AUTH_EXIT_CODE = 75
+
 if __name__ == "__main__":
-    sys.exit(main())
+    try:
+        sys.exit(main())
+    except MLflowAuthError as e:
+        print(
+            f"\nMLflow rejected the token ({e}).\n"
+            "Even a valid token has no permissions until MLflow records the "
+            "user, which mlflow-oidc-auth only does after an interactive UI "
+            "login. Log into the MLflow web UI once, then re-run this analysis.",
+            file=sys.stderr,
+        )
+        sys.exit(AUTH_EXIT_CODE)
diff --git a/exgentic_a2a_runner/libsh/keycloak-direct-access.sh b/exgentic_a2a_runner/libsh/keycloak-direct-access.sh
index aa71cae..ed588a5 100644
--- a/exgentic_a2a_runner/libsh/keycloak-direct-access.sh
+++ b/exgentic_a2a_runner/libsh/keycloak-direct-access.sh
@@ -1,10 +1,12 @@
 #!/bin/bash
-# Enable Keycloak "Direct Access Grants" (the OAuth password grant) for the
-# rossoctl client. deploy-agent.sh, deploy-benchmark.sh and delete-all-deployments.sh
-# all authenticate with grant_type=password against the rossoctl client, which only
-# works when the client has directAccessGrantsEnabled=true. This helper flips that
-# flag via the Keycloak admin API so the password-grant token request in each script
-# succeeds.
+# Enable Keycloak "Direct Access Grants" (the OAuth password grant) for a client.
+# deploy-agent.sh, deploy-benchmark.sh and delete-all-deployments.sh authenticate
+# with grant_type=password against the rossoctl client (the default); analyze-run.sh
+# does the same against the mlflow client. The password grant only works when the
+# client has directAccessGrantsEnabled=true, so this helper flips that flag via the
+# Keycloak admin API before the token request runs.
+#
+# Usage: enable_direct_access_grants [CLIENT_ID]   (CLIENT_ID defaults to "rossoctl")
 #
 # Must be sourced after KEYCLOAK_API is set. Master-realm admin credentials are
 # resolved in priority order:
@@ -15,7 +17,8 @@
 # On any failure (unreachable admin API, missing client, rejected PUT) it prints a
 # diagnostic and exits 1 — the password grant is a hard prerequisite for every caller.
 enable_direct_access_grants() {
-    echo "Enabling Direct Access Grants for rossoctl client..."
+    local target_client="${1:-rossoctl}"
+    echo "Enabling Direct Access Grants for ${target_client} client..."
 
     # Resolve master-realm admin credentials: prefer env vars, fall back to the
     # keycloak-initial-admin secret (RHBK operator), then defaults.
@@ -48,13 +51,13 @@ enable_direct_access_grants() {
         exit 1
     fi
 
-    # Look up the rossoctl client's internal id.
+    # Look up the target client's internal id.
     local client_config client_id
-    client_config=$(curl -s "$KEYCLOAK_API/admin/realms/rossoctl/clients?clientId=rossoctl" \
+    client_config=$(curl -s "$KEYCLOAK_API/admin/realms/rossoctl/clients?clientId=${target_client}" \
         -H "Authorization: Bearer $admin_token" 2>/dev/null)
     client_id=$(echo "$client_config" | grep -o '"id":"[^"]*"' | head -1 | sed 's/"id":"\([^"]*\)"/\1/')
     if [ -z "$client_id" ]; then
-        echo "Error: Could not find rossoctl client ID in Keycloak" >&2
+        echo "Error: Could not find ${target_client} client ID in Keycloak" >&2
         echo "  Response: $client_config" >&2
         exit 1
     fi
@@ -67,9 +70,9 @@ enable_direct_access_grants() {
         -H "Content-Type: application/json" \
         -d '{"directAccessGrantsEnabled": true}' 2>/dev/null) || put_code="000"
     if [ "$put_code" != "204" ] && [ "$put_code" != "200" ]; then
-        echo "Error: Failed to enable direct access grants for rossoctl client (HTTP $put_code)" >&2
+        echo "Error: Failed to enable direct access grants for ${target_client} client (HTTP $put_code)" >&2
         echo "  Response: $(cat /tmp/kc_put_response.txt 2>/dev/null)" >&2
         exit 1
     fi
-    echo "✓ Direct access grants enabled for rossoctl client"
+    echo "✓ Direct access grants enabled for ${target_client} client"
 }

From cea46ff294b45177e50d667d7405e7b3f214c471 Mon Sep 17 00:00:00 2001
From: Yoav Katz <katz@il.ibm.com>
Date: Wed, 29 Jul 2026 12:54:12 +0300
Subject: [PATCH 3/3] Add --save-traces flag to analyze-run.sh to persist
 downloaded traces

Optionally tee the raw traces JSON from download_mlflow_traces.py into a
timestamped file under the given directory before piping to the analyzer,
so the saved copy and the analysis come from the same download. Off by
default; the directory is created if needed.

Also drop a stray leftover `env` debug line that dumped the full
environment before the download pipeline.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Signed-off-by: Yoav Katz <katz@il.ibm.com>
---
 exgentic_a2a_runner/analyze-run.sh | 37 ++++++++++++++++++++++++++++--
 1 file changed, 35 insertions(+), 2 deletions(-)

diff --git a/exgentic_a2a_runner/analyze-run.sh b/exgentic_a2a_runner/analyze-run.sh
index a85b7ed..ef72d2b 100755
--- a/exgentic_a2a_runner/analyze-run.sh
+++ b/exgentic_a2a_runner/analyze-run.sh
@@ -56,6 +56,8 @@ EXPERIMENT_FILTER=""
 COMPARE_EXPERIMENTS=""
 CLUSTER_MODE=""
 INGRESS_DOMAIN=""
+# Directory to save the raw downloaded traces JSON into. Empty = don't save.
+SAVE_TRACES_DIR="${SAVE_TRACES_DIR:-}"
 
 usage() {
     cat << EOF
@@ -75,6 +77,7 @@ Options:
     --mlflow-tls               MLflow serves HTTPS on the forwarded port
     --mlflow-workspace NAME    Send x-mlflow-workspace header
     --auth-mode MODE           Token source: secret (rossoctl oauth secret) or oc-token (oc whoami -t)
+    --save-traces DIR          Save the raw downloaded traces JSON into DIR (created if needed)
     -h, --help                 Show this help message
 
 The MLflow location, TLS, workspace, auth mode, and experiment id all DEFAULT
@@ -102,6 +105,7 @@ Examples:
     $0 --openshift apps.mycluster.example.com
     $0 --openshift apps.mycluster.example.com --experiment-id 3 --compare baseline,test1
     $0 -u http://mlflow.localtest.me:8080 --window 2d
+    $0 --window 6h --save-traces ./traces
 EOF
     exit 1
 }
@@ -123,6 +127,7 @@ while [[ $# -gt 0 ]]; do
         --mlflow-tls)       MLFLOW_TLS="true"; shift ;;
         --mlflow-workspace) MLFLOW_WORKSPACE="$2"; shift 2 ;;
         --auth-mode)        AUTH_MODE="$2"; shift 2 ;;
+        --save-traces|-save-traces) SAVE_TRACES_DIR="$2"; shift 2 ;;
         --kind)             CLUSTER_MODE="kind"; shift ;;
         --openshift)
             CLUSTER_MODE="openshift"
@@ -224,6 +229,15 @@ if ! WINDOW_MS=$(parse_window_ms "$WINDOW"); then
     exit 1
 fi
 
+# If --save-traces was given, make sure the target directory exists (create it
+# if needed) so the downloader's output can be written there.
+if [ -n "$SAVE_TRACES_DIR" ]; then
+    if ! mkdir -p "$SAVE_TRACES_DIR" 2>/dev/null; then
+        echo "Error: could not create traces directory '$SAVE_TRACES_DIR'"
+        exit 1
+    fi
+fi
+
 echo "=== MLflow Trace Analysis ==="
 echo "Cluster mode: $CLUSTER_MODE"
 if [ "$USE_PORT_FORWARD" = "true" ]; then
@@ -243,6 +257,9 @@ fi
 if [ -n "$COMPARE_EXPERIMENTS" ]; then
     echo "Comparing Experiments: $COMPARE_EXPERIMENTS"
 fi
+if [ -n "$SAVE_TRACES_DIR" ]; then
+    echo "Saving traces to: $SAVE_TRACES_DIR"
+fi
 echo ""
 
 # --- Verify kubectl points at the cluster matching CLUSTER_MODE ---
@@ -474,10 +491,26 @@ fi
 # status via PIPESTATUS so we can print an actionable hint instead of a raw
 # HTTP 403 traceback.
 set +e
-python3 "$SCRIPT_DIR/download_mlflow_traces.py" | python3 "$SCRIPT_DIR/analyze_traces.py" $PYTHON_ARGS
-DOWNLOAD_STATUS=${PIPESTATUS[0]}
+if [ -n "$SAVE_TRACES_DIR" ]; then
+    # tee the downloader's stdout (the raw traces JSON) into a timestamped file
+    # in SAVE_TRACES_DIR before it is piped to the analyzer, so both the saved
+    # copy and the analysis come from the same download.
+    SAVE_TRACES_FILE="$SAVE_TRACES_DIR/traces-$(date +%Y%m%d-%H%M%S).json"
+    python3 "$SCRIPT_DIR/download_mlflow_traces.py" \
+        | tee "$SAVE_TRACES_FILE" \
+        | python3 "$SCRIPT_DIR/analyze_traces.py" $PYTHON_ARGS
+    DOWNLOAD_STATUS=${PIPESTATUS[0]}
+else
+    python3 "$SCRIPT_DIR/download_mlflow_traces.py" | python3 "$SCRIPT_DIR/analyze_traces.py" $PYTHON_ARGS
+    DOWNLOAD_STATUS=${PIPESTATUS[0]}
+fi
 set -e
 
+if [ -n "$SAVE_TRACES_DIR" ] && [ "$DOWNLOAD_STATUS" -eq 0 ]; then
+    echo ""
+    echo "✓ Saved traces to $SAVE_TRACES_FILE"
+fi
+
 if [ "$DOWNLOAD_STATUS" -eq 75 ]; then
     echo ""
     echo "MLflow authentication succeeded but access was denied."