From 6c43ab75d36c7e937b55176f5a3b7a060d72623b Mon Sep 17 00:00:00 2001
From: Ryan Forsyth <forsyth2@llnl.gov>
Date: Wed, 26 Nov 2025 18:47:29 -0600
Subject: [PATCH 1/5] Make token file more robust

---
 zstash/globus_utils.py | 168 ++++++++++++++++++++++++++++++++---------
 1 file changed, 131 insertions(+), 37 deletions(-)

diff --git a/zstash/globus_utils.py b/zstash/globus_utils.py
index e5346f69..102eecad 100644
--- a/zstash/globus_utils.py
+++ b/zstash/globus_utils.py
@@ -43,7 +43,28 @@
 # State files
 GLOBUS_CFG: str = os.path.expanduser("~/.globus-native-apps.cfg")
 INI_PATH: str = os.path.expanduser("~/.zstash.ini")
-TOKEN_FILE = os.path.expanduser("~/.zstash_globus_tokens.json")
+# Default token file - can be overridden via environment variable
+DEFAULT_TOKEN_FILE = os.path.expanduser("~/.zstash_globus_tokens.json")
+
+
+# Helper functions for token file management #################################
+def get_token_file_path() -> str:
+    """
+    Get the token file path, checking environment variable first,
+    then falling back to default.
+    """
+    return os.environ.get("ZSTASH_GLOBUS_TOKEN_FILE", DEFAULT_TOKEN_FILE)
+
+
+def get_endpoint_key(endpoints: List[Optional[str]]) -> str:
+    """
+    Generate a unique key for a pair of endpoints.
+    Sorts endpoints to ensure consistency regardless of order.
+    """
+    # Filter out None values and sort to ensure consistent key
+    sorted_eps = sorted([ep for ep in endpoints if ep is not None])
+    return ":".join(sorted_eps)
+
 
 # Independent functions #######################################################
 # The functions here don't rely on the global variables defined in globus.py.
@@ -51,6 +72,8 @@
 
 # Primarily used by globus_activate ###########################################
 def check_state_files():
+    token_file = get_token_file_path()
+
     if os.path.exists(GLOBUS_CFG):
         logger.warning(
             f"Globus CFG {GLOBUS_CFG} exists. This may be left over from earlier versions of zstash, and may cause issues. Consider deleting."
@@ -65,13 +88,13 @@ def check_state_files():
             f"{INI_PATH} does NOT exist. This means we won't be able to read the local endpoint ID from it."
         )
 
-    if os.path.exists(TOKEN_FILE):
+    if os.path.exists(token_file):
         logger.info(
-            f"Token file {TOKEN_FILE} exists. We can try to load tokens from it."
+            f"Token file {token_file} exists. We can try to load tokens from it."
         )
     else:
         logger.info(
-            f"Token file {TOKEN_FILE} does NOT exist. This means we won't be able to load tokens from it."
+            f"Token file {token_file} does NOT exist. This means we won't be able to load tokens from it."
         )
 
 
@@ -126,31 +149,40 @@ def get_local_endpoint_id(local_endpoint_id: Optional[str]) -> str:
 def get_transfer_client_with_auth(
     both_endpoints: List[Optional[str]],
 ) -> TransferClient:
+    endpoint_key = get_endpoint_key(both_endpoints)
+
     tokens = load_tokens()
 
-    # Check if we have stored refresh tokens
-    if "transfer.api.globus.org" in tokens:
-        token_data = tokens["transfer.api.globus.org"]
-        if "refresh_token" in token_data:
-            logger.info("Found stored refresh token - using it")
-            # Create a simple auth client for the RefreshTokenAuthorizer
-            auth_client = NativeAppAuthClient(ZSTASH_CLIENT_ID)
-            try:
-                transfer_authorizer = RefreshTokenAuthorizer(
-                    refresh_token=token_data["refresh_token"], auth_client=auth_client
-                )
-                transfer_client = TransferClient(authorizer=transfer_authorizer)
-                return transfer_client
-            except AuthAPIError as e:
-                logger.error("Stored refresh token is invalid.")
-                logger.error(
-                    f"One possible cause: {TOKEN_FILE} may be configured for a different Globus endpoint. For example, you may have previously set a different destination endpoint for `--hpss=globus://`."
+    # Check if we have stored refresh tokens for this endpoint pair
+    if endpoint_key in tokens:
+        endpoint_tokens = tokens[endpoint_key]
+        if "transfer.api.globus.org" in endpoint_tokens:
+            token_data = endpoint_tokens["transfer.api.globus.org"]
+            if "refresh_token" in token_data:
+                logger.info(
+                    f"Found stored refresh token for endpoints {endpoint_key} - using it"
                 )
-                logger.error(f"Try deleting {TOKEN_FILE} and re-running.")
-                raise e
-
-    # No stored tokens, need to authenticate
-    logger.info("No stored tokens found - starting authentication")
+                # Create a simple auth client for the RefreshTokenAuthorizer
+                auth_client = NativeAppAuthClient(ZSTASH_CLIENT_ID)
+                try:
+                    transfer_authorizer = RefreshTokenAuthorizer(
+                        refresh_token=token_data["refresh_token"],
+                        auth_client=auth_client,
+                    )
+                    transfer_client = TransferClient(authorizer=transfer_authorizer)
+                    return transfer_client
+                except AuthAPIError:
+                    logger.warning(
+                        f"Stored refresh token for {endpoint_key} is invalid, will re-authenticate."
+                    )
+                    # Remove invalid token entry
+                    del tokens[endpoint_key]
+                    save_tokens_to_file(tokens)
+
+    # No stored tokens for this endpoint pair, need to authenticate
+    logger.info(
+        f"No stored tokens found for endpoints {endpoint_key} - starting authentication"
+    )
 
     # Get the required scopes
     all_scopes = get_all_endpoint_scopes(both_endpoints)
@@ -169,7 +201,7 @@ def get_transfer_client_with_auth(
     token_response = client.oauth2_exchange_code_for_tokens(auth_code)
 
     # Save tokens for next time
-    save_tokens(token_response)
+    save_tokens(token_response, both_endpoints)
 
     # Get the transfer token and create authorizer
     globus_transfer_data = token_response.by_resource_server["transfer.api.globus.org"]
@@ -181,16 +213,65 @@ def get_transfer_client_with_auth(
     return transfer_client
 
 
-def load_tokens():
-    if os.path.exists(TOKEN_FILE):
+def load_tokens() -> Dict:
+    """
+    Load all tokens from the token file.
+    Returns a dict with structure:
+    {
+        "endpoint1:endpoint2": {
+            "transfer.api.globus.org": {
+                "access_token": "...",
+                "refresh_token": "...",
+                "expires_at": ...
+            }
+        },
+        ...
+    }
+
+    Also handles legacy single-token format for backward compatibility.
+    """
+    token_file = get_token_file_path()
+
+    if os.path.exists(token_file):
         try:
-            with open(TOKEN_FILE, "r") as f:
-                return json.load(f)
+            with open(token_file, "r") as f:
+                data = json.load(f)
+
+            # Check if this is the old single-token format
+            if "transfer.api.globus.org" in data:
+                # Legacy format detected - migrate it
+                logger.info("Detected legacy token format, migrating to new format")
+                # We can't determine the original endpoints, so we'll just
+                # return empty dict and let user re-authenticate
+                # Optionally, we could try to keep the old token with a generic key
+                return {}
+
+            return data
         except (json.JSONDecodeError, IOError):
+            logger.warning("Error reading token file")
             return {}
     return {}
 
 
+def save_tokens_to_file(tokens: Dict):
+    """
+    Save the complete token dictionary to file.
+    """
+    token_file = get_token_file_path()
+
+    try:
+        # Create directory if it doesn't exist
+        token_dir = os.path.dirname(token_file)
+        if token_dir and not os.path.exists(token_dir):
+            os.makedirs(token_dir)
+
+        with open(token_file, "w") as f:
+            json.dump(tokens, f, indent=2)
+        logger.info(f"Tokens saved successfully to {token_file}")
+    except IOError as e:
+        logger.error(f"Failed to save tokens: {e}")
+
+
 def get_all_endpoint_scopes(endpoints: List[Optional[str]]) -> str:
     inner = " ".join(
         [
@@ -202,7 +283,16 @@ def get_all_endpoint_scopes(endpoints: List[Optional[str]]) -> str:
     return f"urn:globus:auth:scope:transfer.api.globus.org:all[{inner}]"
 
 
-def save_tokens(token_response):
+def save_tokens(token_response, endpoints: List[Optional[str]]):
+    """
+    Save tokens for a specific endpoint pair.
+    """
+    endpoint_key = get_endpoint_key(endpoints)
+
+    # Load existing tokens
+    all_tokens = load_tokens()
+
+    # Prepare tokens for this endpoint pair
     tokens_to_save = {}
     for resource_server, token_data in token_response.by_resource_server.items():
         tokens_to_save[resource_server] = {
@@ -211,9 +301,11 @@ def save_tokens(token_response):
             "expires_at": token_data.get("expires_at_seconds"),
         }
 
-    with open(TOKEN_FILE, "w") as f:
-        json.dump(tokens_to_save, f, indent=2)
-    logger.info("Tokens saved successfully")
+    # Store under the endpoint key
+    all_tokens[endpoint_key] = tokens_to_save
+
+    # Save everything back to file
+    save_tokens_to_file(all_tokens)
 
 
 # Primarily used by globus_transfer ###########################################
@@ -262,6 +354,8 @@ def set_up_TransferData(
 
 
 def submit_transfer_with_checks(transfer_client, transfer_data) -> GlobusHTTPResponse:
+    token_file = get_token_file_path()
+
     task: GlobusHTTPResponse
     try:
         task = transfer_client.submit_transfer(transfer_data)
@@ -273,9 +367,9 @@ def submit_transfer_with_checks(transfer_client, transfer_data) -> GlobusHTTPRes
             )
 
             logger.error(
-                f"One possible cause: {TOKEN_FILE} may be configured for a different Globus endpoint. For example, you may have previously set a different destination endpoint for `--hpss=globus://`."
+                f"One possible cause: {token_file} may be configured for a different Globus endpoint. For example, you may have previously set a different destination endpoint for `--hpss=globus://`."
             )
-            logger.error(f"Try deleting {TOKEN_FILE} and re-running.")
+            logger.error(f"Try deleting {token_file} and re-running.")
 
             logger.error(
                 "Another possible cause: insufficient Globus consents. It's possible the consent on https://auth.globus.org/v2/web/consents is for a different destination endpoint."

From 30d44bc3b0f30a8fa158efa3747070db4ac3fa05 Mon Sep 17 00:00:00 2001
From: Ryan Forsyth <forsyth2@llnl.gov>
Date: Wed, 26 Nov 2025 18:57:38 -0600
Subject: [PATCH 2/5] Test revision 1

---
 .../bash_tests/run_from_any/globus_auth.bash  | 154 +++++++++++++-----
 1 file changed, 112 insertions(+), 42 deletions(-)

diff --git a/tests/integration/bash_tests/run_from_any/globus_auth.bash b/tests/integration/bash_tests/run_from_any/globus_auth.bash
index 527a0368..12c887cb 100755
--- a/tests/integration/bash_tests/run_from_any/globus_auth.bash
+++ b/tests/integration/bash_tests/run_from_any/globus_auth.bash
@@ -16,7 +16,7 @@ check_log_does_not_have()
     local log_file="${2}"
     grep "${not_expected_grep}" ${log_file}
     if [ $? == 0 ]; then
-        echo "Not-expected grep '${expected_grep}' was found in ${log_file}. Test failed."
+        echo "Not-expected grep '${not_expected_grep}' was found in ${log_file}. Test failed."
         exit 2
     fi
 }
@@ -127,7 +127,8 @@ test_single_auth_code()
     check_log_has "INFO: Writing to empty ${INI_PATH}" ${case_name}.log
     check_log_has "INFO: Setting local_endpoint_id based on" ${case_name}.log
     # From get_transfer_client_with_auth
-    check_log_has "INFO: No stored tokens found - starting authentication" ${case_name}.log
+    check_log_has "INFO: No stored tokens found for endpoints" ${case_name}.log
+    check_log_has "starting authentication" ${case_name}.log
     check_log_has "Please go to this URL and login:" ${case_name}.log # Our one expected authentication prompt
     # From save_tokens
     check_log_has "INFO: Tokens saved successfully" ${case_name}.log
@@ -150,7 +151,8 @@ test_single_auth_code()
     check_log_has "INFO: Setting local_endpoint_id based on ${INI_PATH}" ${case_name}.log # Differs from run1
     check_log_has "INFO: Setting local_endpoint_id based on" ${case_name}.log
     # From get_transfer_client_with_auth
-    check_log_has "INFO: Found stored refresh token - using it" ${case_name}.log # Differs from run1
+    check_log_has "INFO: Found stored refresh token for endpoints" ${case_name}.log # Updated to match new log message
+    check_log_has "using it" ${case_name}.log
     check_log_does_not_have "Please go to this URL and login:" ${case_name}.log # There should be no login prompts for run2!
     # From save_tokens
     check_log_does_not_have "INFO: Tokens saved successfully" ${case_name}.log # Differs from run1
@@ -191,16 +193,25 @@ test_different_endpoint1()
 
     case_name="different_endpoint1"
     setup ${case_name} "${src_dir}"
-    # Expecting to see exactly 1 authentication prompt
+    # With multi-token support, this should now work and prompt for auth
+    # Expecting to see exactly 1 authentication prompt for the new endpoint
     zstash create --hpss=${globus_path}/${case_name} zstash_demo 2>&1 | tee ${case_name}.log
-    check_log_has "INFO: Found stored refresh token - using it" ${case_name}.log
-    check_log_has "ERROR: One possible cause" ${case_name}.log
-    check_log_has "ERROR: Try deleting" ${case_name}.log
-    check_log_has "ERROR: Another possible cause" ${case_name}.log
-    check_log_has "try revoking consents before re-running" ${case_name}.log
-    check_log_has "ERROR: Exception: Insufficient Globus consents" ${case_name}.log
-
-    if ! confirm "Did you avoid having to paste any auth codes on this run?"; then
+    if [ $? != 0 ]; then
+        echo "${case_name} failed. Check ${case_name}.log for details."
+        exit 1
+    fi
+
+    # Should see token not found for this endpoint pair, then authentication
+    check_log_has "INFO: No stored tokens found for endpoints" ${case_name}.log
+    check_log_has "starting authentication" ${case_name}.log
+    check_log_has "Please go to this URL and login:" ${case_name}.log
+    check_log_has "INFO: Tokens saved successfully" ${case_name}.log
+
+    # Should NOT see any errors about insufficient consents
+    check_log_does_not_have "ERROR: One possible cause" ${case_name}.log
+    check_log_does_not_have "ERROR: Insufficient Globus consents" ${case_name}.log
+
+    if ! confirm "Did you have to paste an auth code for this new endpoint?"; then
         echo "test_different_endpoint1 failed"
         exit 1
     fi
@@ -220,34 +231,29 @@ test_different_endpoint2()
     dst_endpoint_uuid=$(get_endpoint ${dst_endpoint})
     globus_path=globus://${dst_endpoint_uuid}/${dst_dir}
 
-    echo "Reset Globus consents:"
-    echo "https://auth.globus.org/v2/web/consents > Globus Endpoint Performance Monitoring > rescind all"
-    if ! confirm "Have you revoked Globus consents?"; then
-        exit 1
-    fi
-
     echo "Running test_different_endpoint2"
+    echo "This test verifies that we can switch back to a previously-used endpoint without re-authenticating"
     echo "Exit codes: 0 -- success, 1 -- failure"
 
-    case_name="different_endpoint2a"
+    case_name="different_endpoint2"
     setup ${case_name} "${src_dir}"
-    zstash create --hpss=${globus_path}/${case_name} zstash_demo 2>&1 | tee ${case_name}.log
-    check_log_has ".zstash_globus_tokens.json exists. We can try to load tokens from it." ${case_name}.log
-    check_log_has ".zstash_globus_tokens.json may be configured for a different Globus endpoint." ${case_name}.log
-    check_log_has "Try deleting" ${case_name}.log
-    check_log_has "globus_sdk.services.auth.errors.AuthAPIError: ('POST', 'https://auth.globus.org/v2/oauth2/token', None, 400, 'Error', 'Bad Request')" ${case_name}.log
-
-    rm -rf ~/.zstash_globus_tokens.json
-    case_name="different_endpoint2b"
-    setup ${case_name} "${src_dir}"
-    # Expecting to see exactly 1 authentication prompt
+    # This endpoint was already authenticated in test_different_endpoint1
+    # We should be able to use the stored token without prompting
     zstash create --hpss=${globus_path}/${case_name} zstash_demo 2>&1 | tee ${case_name}.log
     if [ $? != 0 ]; then
-        echo "${case_name} failed. Check ${case_name}_create.log for details."
+        echo "${case_name} failed. Check ${case_name}.log for details."
         exit 1
     fi
 
-    if ! confirm "Did you only have to paste an auth code once (for 2b, not 2a)?"; then
+    # Should find the stored token for this endpoint pair
+    check_log_has "INFO: Found stored refresh token for endpoints" ${case_name}.log
+    check_log_has "using it" ${case_name}.log
+
+    # Should NOT see authentication prompt
+    check_log_does_not_have "Please go to this URL and login:" ${case_name}.log
+    check_log_does_not_have "starting authentication" ${case_name}.log
+
+    if ! confirm "Did you NOT have to paste any auth codes for this run?"; then
         echo "test_different_endpoint2 failed"
         exit 1
     fi
@@ -268,20 +274,27 @@ test_different_endpoint3()
     globus_path=globus://${dst_endpoint_uuid}/${dst_dir}
 
     echo "Running test_different_endpoint3"
+    echo "This test verifies a third different endpoint (also requires auth)"
     echo "Exit codes: 0 -- success, 1 -- failure"
 
-    rm -rf ~/.zstash_globus_tokens.json
     case_name="different_endpoint3"
     setup ${case_name} "${src_dir}"
+    # This is a third endpoint that hasn't been authenticated yet
     # Expecting to see exactly 1 authentication prompt
     zstash create --hpss=${globus_path}/${case_name} zstash_demo 2>&1 | tee ${case_name}.log
     if [ $? != 0 ]; then
-        echo "${case_name} failed. Check ${case_name}_create.log for details."
+        echo "${case_name} failed. Check ${case_name}.log for details."
         exit 1
     fi
 
-    if ! confirm "Did you only have to paste an auth code once?"; then
-        echo "test_different_endpoint2 failed"
+    # Should see no stored token, then authentication
+    check_log_has "INFO: No stored tokens found for endpoints" ${case_name}.log
+    check_log_has "starting authentication" ${case_name}.log
+    check_log_has "Please go to this URL and login:" ${case_name}.log
+    check_log_has "INFO: Tokens saved successfully" ${case_name}.log
+
+    if ! confirm "Did you have to paste an auth code for this third endpoint?"; then
+        echo "test_different_endpoint3 failed"
         exit 1
     fi
     # Cleanup:
@@ -289,6 +302,60 @@ test_different_endpoint3()
     rm -rf ${path_to_repo}/tests/utils/globus_auth
 }
 
+test_legacy_token_migration()
+{
+    local path_to_repo=$1
+    local dst_endpoint=$2
+    local dst_dir=$3
+
+    src_dir=${path_to_repo}/tests/utils/globus_auth
+    mkdir -p ${src_dir}
+    dst_endpoint_uuid=$(get_endpoint ${dst_endpoint})
+    globus_path=globus://${dst_endpoint_uuid}/${dst_dir}
+
+    TOKEN_FILE=${HOME}/.zstash_globus_tokens.json
+
+    echo "Running test_legacy_token_migration"
+    echo "This test verifies that legacy single-token format is handled gracefully"
+    echo "Exit codes: 0 -- success, 1 -- failure"
+
+    # Create a mock legacy token file (just the structure, not valid tokens)
+    cat > ${TOKEN_FILE} << 'EOF'
+{
+  "transfer.api.globus.org": {
+    "access_token": "fake_legacy_token",
+    "refresh_token": "fake_legacy_refresh",
+    "expires_at": 1234567890
+  }
+}
+EOF
+
+    case_name="legacy_migration"
+    setup ${case_name} "${src_dir}"
+
+    # Should detect legacy format and require re-authentication
+    zstash create --hpss=${globus_path}/${case_name} zstash_demo 2>&1 | tee ${case_name}.log
+    if [ $? != 0 ]; then
+        echo "${case_name} failed. Check ${case_name}.log for details."
+        exit 1
+    fi
+
+    # Should see migration message and re-authentication
+    check_log_has "INFO: Detected legacy token format, migrating to new format" ${case_name}.log
+    check_log_has "INFO: No stored tokens found for endpoints" ${case_name}.log
+    check_log_has "starting authentication" ${case_name}.log
+    check_log_has "Please go to this URL and login:" ${case_name}.log
+
+    if ! confirm "Did you have to paste an auth code to migrate from legacy format?"; then
+        echo "test_legacy_token_migration failed"
+        exit 1
+    fi
+
+    # Cleanup:
+    cd ${path_to_repo}/tests/integration/bash_tests/run_from_any
+    rm -rf ${path_to_repo}/tests/utils/globus_auth
+}
+
 # Follow these directions #####################################################
 
 # Example usage:
@@ -371,16 +438,19 @@ test_single_auth_code ${path_to_repo} NERSC_HPSS_ENDPOINT ${hpss_dst_dir}
 echo "Testing transfer to pic#compy-dtn ######################################"
 test_single_auth_code ${path_to_repo} PIC_COMPY_DTN_ENDPOINT ${compy_dst_dir}
 
-echo "Follow-up tests: behavior when switching to different endpoints"
-echo "NOTE: if you commented out tests above, and your last endpoint used was NERSC_PERLMUTTER_ENDPOINT, the following test will not work properly."
-echo "Test 1: What if we switch to a different endpoint? #####################"
+echo "Follow-up tests: multi-endpoint token storage behavior"
+echo "Test 1: Switch to a different endpoint (should prompt for new auth) ####"
 test_different_endpoint1 ${path_to_repo} ${dst_endpoint_switch1} ${dst_dir_switch1}
-echo "Test 2: What if we try a) revoking consents and then b) removing the token file? ###"
+echo "Test 2: Switch back to a previous endpoint (should use stored token) ###"
 test_different_endpoint2 ${path_to_repo} ${dst_endpoint_switch1} ${dst_dir_switch1}
-echo "Test 3: What if we switch to a different endpoint again, but first remove the token file? ###"
+echo "Test 3: Switch to a third different endpoint (should prompt for auth) ##"
 test_different_endpoint3 ${path_to_repo} ${dst_endpoint_switch2} ${dst_dir_switch2}
-echo "Check https://auth.globus.org/v2/web/consents > Globus Endpoint Performance Monitoring: you should have *two* consents there now."
-if ! confirm "Does https://auth.globus.org/v2/web/consents > Globus Endpoint Performance Monitoring show *two* consents?"; then
+echo "Test 4: Verify legacy token format migration ###########################"
+test_legacy_token_migration ${path_to_repo} LCRC_IMPROV_DTN_ENDPOINT ${chrysalis_dst_dir}
+
+echo "Check https://auth.globus.org/v2/web/consents > Globus Endpoint Performance Monitoring"
+echo "You should now have multiple consents (one for each endpoint pair authenticated)."
+if ! confirm "Does https://auth.globus.org/v2/web/consents show multiple consents as expected?"; then
     exit 1
 fi
 echo "All globus_auth tests completed successfully."

From ecf3c632c33b4c7b85b6f0966f59eb959b3a1996 Mon Sep 17 00:00:00 2001
From: Ryan Forsyth <forsyth2@llnl.gov>
Date: Wed, 26 Nov 2025 19:03:24 -0600
Subject: [PATCH 3/5] Test revision 2

---
 .../bash_tests/run_from_any/globus_auth.bash  | 84 +++++++++++++++++++
 1 file changed, 84 insertions(+)

diff --git a/tests/integration/bash_tests/run_from_any/globus_auth.bash b/tests/integration/bash_tests/run_from_any/globus_auth.bash
index 12c887cb..bc4b6232 100755
--- a/tests/integration/bash_tests/run_from_any/globus_auth.bash
+++ b/tests/integration/bash_tests/run_from_any/globus_auth.bash
@@ -356,6 +356,88 @@ EOF
     rm -rf ${path_to_repo}/tests/utils/globus_auth
 }
 
+test_custom_token_file()
+{
+    local path_to_repo=$1
+    local dst_endpoint=$2
+    local dst_dir=$3
+
+    src_dir=${path_to_repo}/tests/utils/globus_auth
+    mkdir -p ${src_dir}
+    dst_endpoint_uuid=$(get_endpoint ${dst_endpoint})
+    globus_path=globus://${dst_endpoint_uuid}/${dst_dir}
+
+    echo "Running test_custom_token_file"
+    echo "This test verifies that ZSTASH_GLOBUS_TOKEN_FILE environment variable works"
+    echo "Exit codes: 0 -- success, 1 -- failure"
+
+    # Create a temporary directory for test tokens
+    TEMP_TOKEN_DIR=$(mktemp -d)
+    CUSTOM_TOKEN_FILE="${TEMP_TOKEN_DIR}/custom_tokens.json"
+
+    case_name="custom_token_file"
+    setup ${case_name} "${src_dir}"
+
+    # Use custom token file via environment variable
+    export ZSTASH_GLOBUS_TOKEN_FILE="${CUSTOM_TOKEN_FILE}"
+
+    # First run - should authenticate and save to custom location
+    zstash create --hpss=${globus_path}/${case_name}_run1 zstash_demo 2>&1 | tee ${case_name}_run1.log
+    if [ $? != 0 ]; then
+        echo "${case_name} run1 failed. Check ${case_name}_run1.log for details."
+        unset ZSTASH_GLOBUS_TOKEN_FILE
+        rm -rf ${TEMP_TOKEN_DIR}
+        exit 1
+    fi
+
+    # Verify custom token file was created
+    if [ ! -f "${CUSTOM_TOKEN_FILE}" ]; then
+        echo "ERROR: Custom token file ${CUSTOM_TOKEN_FILE} was not created!"
+        unset ZSTASH_GLOBUS_TOKEN_FILE
+        rm -rf ${TEMP_TOKEN_DIR}
+        exit 1
+    fi
+
+    # Verify token was saved to custom location
+    check_log_has "INFO: Tokens saved successfully to ${CUSTOM_TOKEN_FILE}" ${case_name}_run1.log
+    check_log_has "INFO: Token file ${CUSTOM_TOKEN_FILE} exists" ${case_name}_run1.log
+
+    # Second run - should use token from custom location
+    zstash create --hpss=${globus_path}/${case_name}_run2 zstash_demo 2>&1 | tee ${case_name}_run2.log
+    if [ $? != 0 ]; then
+        echo "${case_name} run2 failed. Check ${case_name}_run2.log for details."
+        unset ZSTASH_GLOBUS_TOKEN_FILE
+        rm -rf ${TEMP_TOKEN_DIR}
+        exit 1
+    fi
+
+    # Should have used the stored token
+    check_log_has "INFO: Found stored refresh token for endpoints" ${case_name}_run2.log
+    check_log_does_not_have "Please go to this URL and login:" ${case_name}_run2.log
+
+    # Verify default token file was NOT created
+    DEFAULT_TOKEN_FILE=${HOME}/.zstash_globus_tokens.json
+    if [ -f "${DEFAULT_TOKEN_FILE}" ]; then
+        echo "WARNING: Default token file ${DEFAULT_TOKEN_FILE} exists when it shouldn't"
+        echo "This suggests the custom token file setting may not be working correctly"
+    fi
+
+    if ! confirm "Did you only have to paste an auth code once (for run1, not run2)?"; then
+        echo "test_custom_token_file failed"
+        unset ZSTASH_GLOBUS_TOKEN_FILE
+        rm -rf ${TEMP_TOKEN_DIR}
+        exit 1
+    fi
+
+    # Cleanup
+    unset ZSTASH_GLOBUS_TOKEN_FILE
+    rm -rf ${TEMP_TOKEN_DIR}
+    cd ${path_to_repo}/tests/integration/bash_tests/run_from_any
+    rm -rf ${path_to_repo}/tests/utils/globus_auth
+
+    echo "test_custom_token_file completed successfully"
+}
+
 # Follow these directions #####################################################
 
 # Example usage:
@@ -447,6 +529,8 @@ echo "Test 3: Switch to a third different endpoint (should prompt for auth) ##"
 test_different_endpoint3 ${path_to_repo} ${dst_endpoint_switch2} ${dst_dir_switch2}
 echo "Test 4: Verify legacy token format migration ###########################"
 test_legacy_token_migration ${path_to_repo} LCRC_IMPROV_DTN_ENDPOINT ${chrysalis_dst_dir}
+echo "Test 5: Verify custom token file via environment variable ##############"
+test_custom_token_file ${path_to_repo} LCRC_IMPROV_DTN_ENDPOINT ${chrysalis_dst_dir}
 
 echo "Check https://auth.globus.org/v2/web/consents > Globus Endpoint Performance Monitoring"
 echo "You should now have multiple consents (one for each endpoint pair authenticated)."

From 0ab51bcf0c9060c40e612ffdcdba133fe752833f Mon Sep 17 00:00:00 2001
From: Ryan Forsyth <forsyth2@llnl.gov>
Date: Wed, 26 Nov 2025 19:13:59 -0600
Subject: [PATCH 4/5] Use parameter instead of env variable

---
 .../bash_tests/run_from_any/globus_auth.bash  | 18 ++++++----------
 zstash/create.py                              |  7 ++++++-
 zstash/globus.py                              |  8 ++++++-
 zstash/globus_utils.py                        | 21 ++++++++++++++++---
 4 files changed, 37 insertions(+), 17 deletions(-)

diff --git a/tests/integration/bash_tests/run_from_any/globus_auth.bash b/tests/integration/bash_tests/run_from_any/globus_auth.bash
index bc4b6232..706377a5 100755
--- a/tests/integration/bash_tests/run_from_any/globus_auth.bash
+++ b/tests/integration/bash_tests/run_from_any/globus_auth.bash
@@ -368,7 +368,7 @@ test_custom_token_file()
     globus_path=globus://${dst_endpoint_uuid}/${dst_dir}
 
     echo "Running test_custom_token_file"
-    echo "This test verifies that ZSTASH_GLOBUS_TOKEN_FILE environment variable works"
+    echo "This test verifies that --globus-token-file parameter works"
     echo "Exit codes: 0 -- success, 1 -- failure"
 
     # Create a temporary directory for test tokens
@@ -378,14 +378,10 @@ test_custom_token_file()
     case_name="custom_token_file"
     setup ${case_name} "${src_dir}"
 
-    # Use custom token file via environment variable
-    export ZSTASH_GLOBUS_TOKEN_FILE="${CUSTOM_TOKEN_FILE}"
-
     # First run - should authenticate and save to custom location
-    zstash create --hpss=${globus_path}/${case_name}_run1 zstash_demo 2>&1 | tee ${case_name}_run1.log
+    zstash create --hpss=${globus_path}/${case_name}_run1 --globus-token-file="${CUSTOM_TOKEN_FILE}" zstash_demo 2>&1 | tee ${case_name}_run1.log
     if [ $? != 0 ]; then
         echo "${case_name} run1 failed. Check ${case_name}_run1.log for details."
-        unset ZSTASH_GLOBUS_TOKEN_FILE
         rm -rf ${TEMP_TOKEN_DIR}
         exit 1
     fi
@@ -393,25 +389,25 @@ test_custom_token_file()
     # Verify custom token file was created
     if [ ! -f "${CUSTOM_TOKEN_FILE}" ]; then
         echo "ERROR: Custom token file ${CUSTOM_TOKEN_FILE} was not created!"
-        unset ZSTASH_GLOBUS_TOKEN_FILE
         rm -rf ${TEMP_TOKEN_DIR}
         exit 1
     fi
 
-    # Verify token was saved to custom location
+    # Verify token was saved to custom location and custom path was used
+    check_log_has "INFO: Using custom token file: ${CUSTOM_TOKEN_FILE}" ${case_name}_run1.log
     check_log_has "INFO: Tokens saved successfully to ${CUSTOM_TOKEN_FILE}" ${case_name}_run1.log
     check_log_has "INFO: Token file ${CUSTOM_TOKEN_FILE} exists" ${case_name}_run1.log
 
     # Second run - should use token from custom location
-    zstash create --hpss=${globus_path}/${case_name}_run2 zstash_demo 2>&1 | tee ${case_name}_run2.log
+    zstash create --hpss=${globus_path}/${case_name}_run2 --globus-token-file="${CUSTOM_TOKEN_FILE}" zstash_demo 2>&1 | tee ${case_name}_run2.log
     if [ $? != 0 ]; then
         echo "${case_name} run2 failed. Check ${case_name}_run2.log for details."
-        unset ZSTASH_GLOBUS_TOKEN_FILE
         rm -rf ${TEMP_TOKEN_DIR}
         exit 1
     fi
 
     # Should have used the stored token
+    check_log_has "INFO: Using custom token file: ${CUSTOM_TOKEN_FILE}" ${case_name}_run2.log
     check_log_has "INFO: Found stored refresh token for endpoints" ${case_name}_run2.log
     check_log_does_not_have "Please go to this URL and login:" ${case_name}_run2.log
 
@@ -424,13 +420,11 @@ test_custom_token_file()
 
     if ! confirm "Did you only have to paste an auth code once (for run1, not run2)?"; then
         echo "test_custom_token_file failed"
-        unset ZSTASH_GLOBUS_TOKEN_FILE
         rm -rf ${TEMP_TOKEN_DIR}
         exit 1
     fi
 
     # Cleanup
-    unset ZSTASH_GLOBUS_TOKEN_FILE
     rm -rf ${TEMP_TOKEN_DIR}
     cd ${path_to_repo}/tests/integration/bash_tests/run_from_any
     rm -rf ${path_to_repo}/tests/utils/globus_auth
diff --git a/zstash/create.py b/zstash/create.py
index b502f1e6..1d324f50 100644
--- a/zstash/create.py
+++ b/zstash/create.py
@@ -57,7 +57,7 @@ def create():
         if url.scheme == "globus":
             # identify globus endpoints
             logger.debug(f"{ts_utc()}:Calling globus_activate(hpss)")
-            globus_activate(hpss)
+            globus_activate(hpss, args.globus_token_file)
         else:
             # config.hpss is not "none", so we need to
             # create target HPSS directory
@@ -176,6 +176,11 @@ def setup_create() -> Tuple[str, argparse.Namespace]:
         action="store_true",
         help="FOR ADVANCED USERS ONLY: If a duplicate tar is encountered, overwrite the existing database record with the new one (i.e., it will assume the latest tar is the correct one). If this flag is not set, zstash will permit multiple entries for the same tar in its database.",
     )
+    optional.add_argument(
+        "--globus-token-file",
+        type=str,
+        help="Path to custom Globus token file. If not specified, uses ~/.zstash_globus_tokens.json",
+    )
     optional.add_argument(
         "--for-developers-force-database-corruption",
         type=str,
diff --git a/zstash/globus.py b/zstash/globus.py
index 2cacad5f..f6e39fd6 100644
--- a/zstash/globus.py
+++ b/zstash/globus.py
@@ -13,6 +13,7 @@
     check_state_files,
     get_local_endpoint_id,
     get_transfer_client_with_auth,
+    set_token_file_path,
     set_up_TransferData,
     submit_transfer_with_checks,
 )
@@ -27,7 +28,7 @@
 archive_directory_listing: IterableTransferResponse = None
 
 
-def globus_activate(hpss: str):
+def globus_activate(hpss: str, token_file: Optional[str] = None):
     """
     Read the local globus endpoint UUID from ~/.zstash.ini.
     If the ini file does not exist, create an ini file with empty values,
@@ -40,6 +41,11 @@ def globus_activate(hpss: str):
     url = urlparse(hpss)
     if url.scheme != "globus":
         return
+
+    # Set the token file path if provided
+    if token_file:
+        set_token_file_path(token_file)
+
     check_state_files()
     remote_endpoint = url.netloc
     local_endpoint = get_local_endpoint_id(local_endpoint)
diff --git a/zstash/globus_utils.py b/zstash/globus_utils.py
index 102eecad..bdd9390c 100644
--- a/zstash/globus_utils.py
+++ b/zstash/globus_utils.py
@@ -43,17 +43,32 @@
 # State files
 GLOBUS_CFG: str = os.path.expanduser("~/.globus-native-apps.cfg")
 INI_PATH: str = os.path.expanduser("~/.zstash.ini")
-# Default token file - can be overridden via environment variable
+# Default token file - can be overridden via set_token_file_path()
 DEFAULT_TOKEN_FILE = os.path.expanduser("~/.zstash_globus_tokens.json")
 
+# Module-level variable to store custom token file path
+_custom_token_file: Optional[str] = None
+
 
 # Helper functions for token file management #################################
+def set_token_file_path(token_file: str) -> None:
+    """
+    Set a custom token file path for the current session.
+    This should be called before any Globus operations.
+    """
+    global _custom_token_file
+    _custom_token_file = token_file
+    logger.info(f"Using custom token file: {token_file}")
+
+
 def get_token_file_path() -> str:
     """
-    Get the token file path, checking environment variable first,
+    Get the token file path, checking custom path first,
     then falling back to default.
     """
-    return os.environ.get("ZSTASH_GLOBUS_TOKEN_FILE", DEFAULT_TOKEN_FILE)
+    if _custom_token_file:
+        return _custom_token_file
+    return DEFAULT_TOKEN_FILE
 
 
 def get_endpoint_key(endpoints: List[Optional[str]]) -> str:

From a548292398ccb6887358827f299e7e66728a7f65 Mon Sep 17 00:00:00 2001
From: Ryan Forsyth <forsyth2@llnl.gov>
Date: Wed, 26 Nov 2025 19:25:28 -0600
Subject: [PATCH 5/5] Update other zstash operations

---
 zstash/extract.py | 7 +++++++
 zstash/ls.py      | 7 +++++++
 zstash/update.py  | 7 ++++++-
 3 files changed, 20 insertions(+), 1 deletion(-)

diff --git a/zstash/extract.py b/zstash/extract.py
index 64977aef..bd4f6fd3 100644
--- a/zstash/extract.py
+++ b/zstash/extract.py
@@ -19,6 +19,7 @@
 import _io
 
 from . import parallel
+from .globus import globus_activate
 from .hpss import hpss_get
 from .settings import (
     BLOCK_SIZE,
@@ -98,6 +99,11 @@ def setup_extract() -> Tuple[argparse.Namespace, str]:
         "--retries", type=int, default=1, help="number of times to retry an hsi command"
     )
     optional.add_argument("--tars", type=str, help="specify which tars to process")
+    optional.add_argument(
+        "--globus-token-file",
+        type=str,
+        help="Path to custom Globus token file. If not specified, uses ~/.zstash_globus_tokens.json",
+    )
     optional.add_argument(
         "-v", "--verbose", action="store_true", help="increase output verbosity"
     )
@@ -175,6 +181,7 @@ def extract_database(
                 hpss: str = config.hpss
             else:
                 raise TypeError("Invalid config.hpss={}".format(config.hpss))
+            globus_activate(hpss, args.globus_token_file)
             hpss_get(hpss, get_db_filename(cache), cache)
         else:
             error_str: str = (
diff --git a/zstash/ls.py b/zstash/ls.py
index 8b6ad6e4..ff9f231b 100644
--- a/zstash/ls.py
+++ b/zstash/ls.py
@@ -7,6 +7,7 @@
 import sys
 from typing import List, Tuple, Union
 
+from .globus import globus_activate
 from .hpss import hpss_get
 from .settings import (
     DEFAULT_CACHE,
@@ -70,6 +71,11 @@ def setup_ls() -> Tuple[argparse.Namespace, str]:
         help='the path to the zstash archive on the local file system. The default name is "zstash".',
     )
     optional.add_argument("--tars", action="store_true", help="Display tars")
+    optional.add_argument(
+        "--globus-token-file",
+        type=str,
+        help="Path to custom Globus token file. If not specified, uses ~/.zstash_globus_tokens.json",
+    )
     optional.add_argument(
         "-v", "--verbose", action="store_true", help="increase output verbosity"
     )
@@ -101,6 +107,7 @@ def ls_database(args: argparse.Namespace, cache: str) -> List[FilesRow]:
             else:
                 raise TypeError("Invalid config.hpss={}".format(config.hpss))
             try:
+                globus_activate(hpss, args.globus_token_file)
                 # Retrieve from HPSS
                 hpss_get(hpss, get_db_filename(cache), cache)
             except RuntimeError:
diff --git a/zstash/update.py b/zstash/update.py
index b0f2af40..c0414cd2 100644
--- a/zstash/update.py
+++ b/zstash/update.py
@@ -105,6 +105,11 @@ def setup_update() -> Tuple[argparse.Namespace, str]:
         action="store_true",
         help="do not wait for each Globus transfer until it completes.",
     )
+    optional.add_argument(
+        "--globus-token-file",
+        type=str,
+        help="Path to custom Globus token file. If not specified, uses ~/.zstash_globus_tokens.json",
+    )
     optional.add_argument(
         "--error-on-duplicate-tar",
         action="store_true",
@@ -160,7 +165,7 @@ def update_database(  # noqa: C901
                 hpss: str = config.hpss
             else:
                 raise TypeError("Invalid config.hpss={}".format(config.hpss))
-            globus_activate(hpss)
+            globus_activate(hpss, args.globus_token_file)
             hpss_get(hpss, get_db_filename(cache), cache)
         else:
             error_str: str = (