Merged in feature/cognito-shared-environments (pull request #211)

cognito and permit shared prod environment support

* code complete

* user creattion tool

test harness

cmd/auth_related/cognito.auth.harness/main.go

@@ -231,7 +231,7 @@ func main() {
 	config.SetAuthRoutePermissions(routePermissions)
 
 	// Register authentication routes
-	cognitoauth.RegisterRoutes(e, config)
+	cognitoauth.RegisterRoutes(e, config, nil)
 
 	// Register your application routes
 	// These will be protected based on the permissions set above

cmd/auth_related/user.creation.tool/cognito-permit-sync.go

@@ -18,6 +18,8 @@ import (
 
 	"github.com/aws/aws-sdk-go-v2/service/cognitoidentityprovider"
 	"github.com/joho/godotenv"
+
+	"queryorchestration/internal/usermanagement"
 )
 
 type ErrorResponse struct {
@@ -72,8 +74,9 @@ type AppConfig struct {
 	PermitBaseURL   string
 
 	// Cognito config
-	CognitoUserPoolID string
-	CognitoRegion     string
+	CognitoUserPoolID    string
+	CognitoRegion        string
+	CognitoEnvironmentID string
 
 	// Other
 	CSVPath      string
@@ -326,6 +329,9 @@ func parseFlags() AppConfig {
 	flag.StringVar(&config.PermitEnvID, "env", "", "Permit.io environment ID or key (required unless using 'use-key' with env-scoped API key)")
 	flag.StringVar(&config.PermitBaseURL, "api-url", "https://api.permit.io", "Permit.io API base URL")
 
+	// Cognito environment isolation flag
+	flag.StringVar(&config.CognitoEnvironmentID, "cognito-env-id", "", "Cognito environment ID to tag users with (falls back to COGNITO_ENVIRONMENT_ID env var)")
+
 	// File flag
 	flag.StringVar(&config.CSVPath, "csv", "", "Path to CSV file containing users (required)")
 
@@ -349,6 +355,7 @@ func parseFlags() AppConfig {
 		fmt.Fprintf(os.Stderr, "  COGNITO_USER_POOL_ID    AWS Cognito User Pool ID (required)\n")
 		fmt.Fprintf(os.Stderr, "  COGNITO_CLIENT_ID       AWS Cognito Client ID (required)\n")
 		fmt.Fprintf(os.Stderr, "  COGNITO_CLIENT_SECRET   AWS Cognito Client Secret (optional)\n")
+		fmt.Fprintf(os.Stderr, "  COGNITO_ENVIRONMENT_ID  Cognito environment ID for user tagging (optional)\n")
 		fmt.Fprintf(os.Stderr, "  AWS_REGION             AWS Region (defaults to us-east-1)\n")
 		fmt.Fprintf(os.Stderr, "\nCSV Format:\n")
 		fmt.Fprintf(os.Stderr, "  email,first_name,last_name,role1,role2,...,roleN\n")
@@ -383,6 +390,11 @@ func parseFlags() AppConfig {
 		config.CognitoRegion = "us-east-1"
 	}
 
+	// Fall back to env var for cognito environment ID if flag not set
+	if config.CognitoEnvironmentID == "" {
+		config.CognitoEnvironmentID = os.Getenv(usermanagement.CognitoEnvIDEnvVar)
+	}
+
 	return config
 }
 
@@ -405,6 +417,11 @@ func validateConfig(config AppConfig) error {
 		return fmt.Errorf("COGNITO_USER_POOL_ID environment variable is not set")
 	}
 
+	// Validate Cognito environment ID if set
+	if err := usermanagement.ValidateCognitoEnvironmentID(config.CognitoEnvironmentID); err != nil {
+		return fmt.Errorf("invalid cognito environment ID: %w", err)
+	}
+
 	// Validate CSV
 	if config.CSVPath == "" {
 		return fmt.Errorf("CSV file path is required")
@@ -725,16 +742,17 @@ func ensureCognitoUser(ctx context.Context, procCtx *UserProcessingContext) (*Co
 		auditLogger.LogAction(ActionCreate, SystemCognito, user["email"], ResultSuccess, "DRY RUN - Would create user")
 		// Create a mock user for dry run
 		return &CognitoUserResponse{
-			SubjectID: "mock-subject-id-" + user["email"],
-			Username:  user["email"],
-			Email:     user["email"],
-			FirstName: user["first_name"],
-			LastName:  user["last_name"],
-			Status:    "CONFIRMED",
+			SubjectID:     "mock-subject-id-" + user["email"],
+			Username:      user["email"],
+			Email:         user["email"],
+			FirstName:     user["first_name"],
+			LastName:      user["last_name"],
+			Status:        "CONFIRMED",
+			EnvironmentID: config.CognitoEnvironmentID,
 		}, nil
 	}
 
-	cognitoUser, err := createCognitoUser(ctx, cognitoClient, cognitoConfig, user)
+	cognitoUser, err := createCognitoUser(ctx, cognitoClient, cognitoConfig, user, config.CognitoEnvironmentID)
 	if err != nil {
 		auditLogger.LogAction(ActionCreate, SystemCognito, user["email"], ResultFailure, fmt.Sprintf("Failed to create: %v", err))
 		fmt.Printf("  ❌ Cognito creation failed: %v\n", err)

cmd/auth_related/user.creation.tool/cognito.go

@@ -16,6 +16,8 @@ import (
 	"github.com/aws/aws-sdk-go-v2/service/cognitoidentityprovider"
 	"github.com/aws/aws-sdk-go-v2/service/cognitoidentityprovider/types"
 	"github.com/aws/aws-sdk-go-v2/service/sts"
+
+	"queryorchestration/internal/usermanagement"
 )
 
 // Cognito types
@@ -31,13 +33,14 @@ type UserInput struct {
 }
 
 type CognitoUserResponse struct {
-	SubjectID string
-	Username  string
-	Email     string
-	FirstName string
-	LastName  string
-	Status    string
-	Enabled   bool
+	SubjectID     string
+	Username      string
+	Email         string
+	FirstName     string
+	LastName      string
+	Status        string
+	Enabled       bool
+	EnvironmentID string
 }
 
 // sanityCheckCredentials validates that both AWS and Permit.io credentials are properly configured.
@@ -107,7 +110,7 @@ func initializeCognitoClient() (*cognitoidentityprovider.Client, error) {
 	return client, nil
 }
 
-func createCognitoUser(ctx context.Context, client *cognitoidentityprovider.Client, config *CognitoConfig, user map[string]string) (*CognitoUserResponse, error) {
+func createCognitoUser(ctx context.Context, client *cognitoidentityprovider.Client, config *CognitoConfig, user map[string]string, environmentID string) (*CognitoUserResponse, error) {
 	// Use email as username
 	username := user["email"]
 
@@ -131,6 +134,14 @@ func createCognitoUser(ctx context.Context, client *cognitoidentityprovider.Clie
 		},
 	}
 
+	// Append environment ID attribute when configured
+	if environmentID != "" {
+		userAttributes = append(userAttributes, types.AttributeType{
+			Name:  aws.String(usermanagement.CognitoEnvIDAttrName),
+			Value: aws.String(environmentID),
+		})
+	}
+
 	// Prepare the AdminCreateUser input
 	createUserInput := &cognitoidentityprovider.AdminCreateUserInput{
 		UserPoolId:             aws.String(config.UserPoolID),
@@ -156,13 +167,14 @@ func createCognitoUser(ctx context.Context, client *cognitoidentityprovider.Clie
 
 	// Build response
 	response := &CognitoUserResponse{
-		SubjectID: subjectID,
-		Username:  aws.ToString(result.User.Username),
-		Email:     user["email"],
-		FirstName: user["first_name"],
-		LastName:  user["last_name"],
-		Status:    string(result.User.UserStatus),
-		Enabled:   result.User.Enabled,
+		SubjectID:     subjectID,
+		Username:      aws.ToString(result.User.Username),
+		Email:         user["email"],
+		FirstName:     user["first_name"],
+		LastName:      user["last_name"],
+		Status:        string(result.User.UserStatus),
+		Enabled:       result.User.Enabled,
+		EnvironmentID: environmentID,
 	}
 
 	return response, nil
@@ -195,7 +207,7 @@ func getCognitoUser(ctx context.Context, client *cognitoidentityprovider.Client,
 	}
 
 	// Extract attributes
-	var subjectID, email, firstName, lastName string
+	var subjectID, email, firstName, lastName, environmentID string
 	for _, attr := range result.UserAttributes {
 		switch aws.ToString(attr.Name) {
 		case "sub":
@@ -206,17 +218,20 @@ func getCognitoUser(ctx context.Context, client *cognitoidentityprovider.Client,
 			firstName = aws.ToString(attr.Value)
 		case "family_name":
 			lastName = aws.ToString(attr.Value)
+		case usermanagement.CognitoEnvIDAttrName:
+			environmentID = aws.ToString(attr.Value)
 		}
 	}
 
 	return &CognitoUserResponse{
-		SubjectID: subjectID,
-		Username:  aws.ToString(result.Username),
-		Email:     email,
-		FirstName: firstName,
-		LastName:  lastName,
-		Status:    string(result.UserStatus),
-		Enabled:   result.Enabled,
+		SubjectID:     subjectID,
+		Username:      aws.ToString(result.Username),
+		Email:         email,
+		FirstName:     firstName,
+		LastName:      lastName,
+		Status:        string(result.UserStatus),
+		Enabled:       result.Enabled,
+		EnvironmentID: environmentID,
 	}, nil
 }
 

cmd/auth_related/user.creation.tool/config_test.go

@@ -0,0 +1,116 @@
+package main
+
+import (
+	"flag"
+	"os"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	"queryorchestration/internal/usermanagement"
+)
+
+// TestParseFlags_CognitoEnvIDFlagOverridesEnvVar verifies that the -cognito-env-id flag
+// takes precedence over the COGNITO_ENVIRONMENT_ID environment variable. This tests the
+// fallback logic at cognito-permit-sync.go:393-396.
+func TestParseFlags_CognitoEnvIDFlagOverridesEnvVar(t *testing.T) {
+	// Set the env var that would normally be the fallback
+	t.Setenv(usermanagement.CognitoEnvIDEnvVar, "fromenv")
+
+	// Reset the global flag.CommandLine to allow re-parsing
+	flag.CommandLine = flag.NewFlagSet(os.Args[0], flag.ContinueOnError)
+
+	// Simulate command-line args with the flag set
+	origArgs := os.Args
+	os.Args = []string{"test", "-cognito-env-id=fromflag", "-csv=dummy.csv", "-project=test", "-env=test"}
+	defer func() { os.Args = origArgs }()
+
+	config := parseFlags()
+
+	assert.Equal(t, "fromflag", config.CognitoEnvironmentID,
+		"Flag -cognito-env-id should override COGNITO_ENVIRONMENT_ID env var")
+}
+
+// TestParseFlags_CognitoEnvIDFallsBackToEnvVar verifies that when the -cognito-env-id
+// flag is not provided, the COGNITO_ENVIRONMENT_ID env var is used as a fallback.
+func TestParseFlags_CognitoEnvIDFallsBackToEnvVar(t *testing.T) {
+	t.Setenv(usermanagement.CognitoEnvIDEnvVar, "fromenv")
+
+	// Reset the global flag.CommandLine to allow re-parsing
+	flag.CommandLine = flag.NewFlagSet(os.Args[0], flag.ContinueOnError)
+
+	origArgs := os.Args
+	os.Args = []string{"test", "-csv=dummy.csv", "-project=test", "-env=test"}
+	defer func() { os.Args = origArgs }()
+
+	config := parseFlags()
+
+	assert.Equal(t, "fromenv", config.CognitoEnvironmentID,
+		"COGNITO_ENVIRONMENT_ID env var should be used when flag is not set")
+}
+
+// TestValidateConfig_InvalidCognitoEnvID verifies that validateConfig returns an error
+// when CognitoEnvironmentID contains invalid characters (e.g., uppercase letters).
+func TestValidateConfig_InvalidCognitoEnvID(t *testing.T) {
+	// Create a temp CSV file so file validation passes
+	tmpFile, err := os.CreateTemp(t.TempDir(), "test-*.csv")
+	require.NoError(t, err)
+	defer tmpFile.Close()
+
+	config := AppConfig{
+		PermitAPIKey:         "permit_key_test",
+		PermitProjectID:      "testproject",
+		PermitEnvID:          "testenv",
+		CognitoUserPoolID:    "us-east-1_TestPool",
+		CognitoRegion:        "us-east-1",
+		CognitoEnvironmentID: "UPPERCASE", // invalid: must be lowercase
+		CSVPath:              tmpFile.Name(),
+	}
+
+	err = validateConfig(config)
+	assert.Error(t, err, "Should reject uppercase environment ID")
+	assert.Contains(t, err.Error(), "invalid cognito environment ID")
+}
+
+// TestValidateConfig_ValidCognitoEnvID verifies that validateConfig accepts a valid
+// lowercase environment ID without errors from the env ID validation.
+func TestValidateConfig_ValidCognitoEnvID(t *testing.T) {
+	tmpFile, err := os.CreateTemp(t.TempDir(), "test-*.csv")
+	require.NoError(t, err)
+	defer tmpFile.Close()
+
+	config := AppConfig{
+		PermitAPIKey:         "permit_key_test",
+		PermitProjectID:      "testproject",
+		PermitEnvID:          "testenv",
+		CognitoUserPoolID:    "us-east-1_TestPool",
+		CognitoRegion:        "us-east-1",
+		CognitoEnvironmentID: "acmehealth", // valid
+		CSVPath:              tmpFile.Name(),
+	}
+
+	err = validateConfig(config)
+	assert.NoError(t, err, "Should accept valid lowercase environment ID")
+}
+
+// TestValidateConfig_EmptyCognitoEnvID verifies that an empty CognitoEnvironmentID
+// passes validation (backward compatibility -- not all deployments use env isolation).
+func TestValidateConfig_EmptyCognitoEnvID(t *testing.T) {
+	tmpFile, err := os.CreateTemp(t.TempDir(), "test-*.csv")
+	require.NoError(t, err)
+	defer tmpFile.Close()
+
+	config := AppConfig{
+		PermitAPIKey:         "permit_key_test",
+		PermitProjectID:      "testproject",
+		PermitEnvID:          "testenv",
+		CognitoUserPoolID:    "us-east-1_TestPool",
+		CognitoRegion:        "us-east-1",
+		CognitoEnvironmentID: "", // empty is valid
+		CSVPath:              tmpFile.Name(),
+	}
+
+	err = validateConfig(config)
+	assert.NoError(t, err, "Should accept empty environment ID (backward compat)")
+}

cmd/auth_related/user.creation.tool/disable_users_test.go

@@ -330,7 +330,7 @@ func TestDisableEnableOperations(t *testing.T) {
 	}
 
 	// Create user
-	_, err = createCognitoUser(ctx, cognitoClient, cognitoConfig, testUser)
+	_, err = createCognitoUser(ctx, cognitoClient, cognitoConfig, testUser, "")
 	if err != nil {
 		// If user already exists, that's fine
 		if !strings.Contains(err.Error(), "already exists") && !strings.Contains(err.Error(), "UsernameExistsException") {
@@ -411,7 +411,7 @@ func TestIdempotentDisableEnable(t *testing.T) {
 		Region:     os.Getenv("AWS_REGION"),
 	}
 
-	_, err = createCognitoUser(ctx, cognitoClient, cognitoConfig, testUser)
+	_, err = createCognitoUser(ctx, cognitoClient, cognitoConfig, testUser, "")
 	if err != nil && !strings.Contains(err.Error(), "already exists") && !strings.Contains(err.Error(), "UsernameExistsException") {
 		t.Fatalf("Failed to create test user: %v", err)
 	}

cmd/auth_related/user.creation.tool/list.users.by.env.sh

@@ -0,0 +1,121 @@
+#!/bin/sh
+# list.users.by.env.sh - Lists Cognito users filtered by custom:environment_id attribute.
+#
+# Usage: ./list.users.by.env.sh <user-pool-id> [environment-id]
+#
+# Arguments:
+#   user-pool-id    - AWS Cognito User Pool ID (required)
+#   environment-id  - Environment ID to filter by (optional, falls back to COGNITO_ENVIRONMENT_ID env var)
+#
+# If no environment-id is provided (via argument or env var), lists all users with their
+# environment_id attribute value.
+#
+# Requires: aws cli, jq
+#
+# Examples:
+#   ./list.users.by.env.sh us-east-2_abc123 acmehealth
+#   COGNITO_ENVIRONMENT_ID=acmehealth ./list.users.by.env.sh us-east-2_abc123
+#   ./list.users.by.env.sh us-east-2_abc123   # lists all users with env info
+
+set -e
+
+USER_POOL_ID="$1"
+ENV_ID="${2:-$COGNITO_ENVIRONMENT_ID}"
+
+if [ -z "$USER_POOL_ID" ]; then
+    echo "Usage: $0 <user-pool-id> [environment-id]" >&2
+    echo "" >&2
+    echo "Arguments:" >&2
+    echo "  user-pool-id    AWS Cognito User Pool ID (required)" >&2
+    echo "  environment-id  Environment ID to filter by (optional)" >&2
+    echo "" >&2
+    echo "Falls back to COGNITO_ENVIRONMENT_ID env var if environment-id not provided." >&2
+    echo "If neither is set, lists all users with their environment_id value." >&2
+    exit 1
+fi
+
+# Check for required tools
+if ! command -v aws >/dev/null 2>&1; then
+    echo "Error: aws cli is required but not installed." >&2
+    exit 1
+fi
+if ! command -v jq >/dev/null 2>&1; then
+    echo "Error: jq is required but not installed." >&2
+    exit 1
+fi
+
+# Fetch all users with pagination
+ALL_USERS="[]"
+PAGINATION_TOKEN=""
+
+while true; do
+    if [ -z "$PAGINATION_TOKEN" ]; then
+        RESPONSE=$(aws cognito-idp list-users \
+            --user-pool-id "$USER_POOL_ID" \
+            --max-results 60 \
+            --output json 2>&1)
+    else
+        RESPONSE=$(aws cognito-idp list-users \
+            --user-pool-id "$USER_POOL_ID" \
+            --max-results 60 \
+            --pagination-token "$PAGINATION_TOKEN" \
+            --output json 2>&1)
+    fi
+
+    if [ $? -ne 0 ]; then
+        echo "Error fetching users: $RESPONSE" >&2
+        exit 1
+    fi
+
+    PAGE_USERS=$(echo "$RESPONSE" | jq '.Users')
+    ALL_USERS=$(echo "$ALL_USERS $PAGE_USERS" | jq -s '.[0] + .[1]')
+
+    PAGINATION_TOKEN=$(echo "$RESPONSE" | jq -r '.PaginationToken // empty')
+    if [ -z "$PAGINATION_TOKEN" ]; then
+        break
+    fi
+done
+
+TOTAL=$(echo "$ALL_USERS" | jq 'length')
+
+if [ -n "$ENV_ID" ]; then
+    # Filter: show users matching this environment_id OR users with no environment_id (untagged)
+    FILTERED=$(echo "$ALL_USERS" | jq --arg env "$ENV_ID" '
+        [.[] | select(
+            (.Attributes // [] | map(select(.Name == "custom:environment_id")) | length == 0) or
+            (.Attributes // [] | map(select(.Name == "custom:environment_id" and .Value == $env)) | length > 0)
+        )]
+    ')
+    FILTERED_COUNT=$(echo "$FILTERED" | jq 'length')
+
+    echo "Users for environment: $ENV_ID"
+    echo "Showing $FILTERED_COUNT of $TOTAL total users (matching + untagged)"
+    echo "============================================="
+    echo ""
+
+    echo "$FILTERED" | jq -r '
+        .[] |
+        "Email:          " + (.Attributes // [] | map(select(.Name == "email")) | .[0].Value // "N/A") +
+        "\nUsername:       " + .Username +
+        "\nStatus:        " + .UserStatus +
+        "\nEnabled:       " + (.Enabled | tostring) +
+        "\nEnvironment:   " + (.Attributes // [] | map(select(.Name == "custom:environment_id")) | .[0].Value // "(untagged)") +
+        "\n---"
+    '
+else
+    # No filter: list all users with their environment_id
+    echo "All users (no environment filter)"
+    echo "Total: $TOTAL users"
+    echo "============================================="
+    echo ""
+
+    echo "$ALL_USERS" | jq -r '
+        .[] |
+        "Email:          " + (.Attributes // [] | map(select(.Name == "email")) | .[0].Value // "N/A") +
+        "\nUsername:       " + .Username +
+        "\nStatus:        " + .UserStatus +
+        "\nEnabled:       " + (.Enabled | tostring) +
+        "\nEnvironment:   " + (.Attributes // [] | map(select(.Name == "custom:environment_id")) | .[0].Value // "(untagged)") +
+        "\n---"
+    '
+fi

cmd/auth_related/user.creation.tool/permit.go

@@ -85,15 +85,22 @@ func getAPIKeyScope(client *http.Client, config AppConfig) (*APIKeyScope, error)
 func createPermitUser(client *http.Client, config AppConfig, cognitoUser *CognitoUserResponse) error {
 	url := fmt.Sprintf("%s/v2/facts/%s/%s/users", config.PermitBaseURL, config.PermitProjectID, config.PermitEnvID)
 
+	attributes := map[string]interface{}{
+		"cognito_username": cognitoUser.Username,
+		"cognito_status":   cognitoUser.Status,
+	}
+
+	// Add environment_id to Permit.io attributes when configured
+	if config.CognitoEnvironmentID != "" {
+		attributes["environment_id"] = config.CognitoEnvironmentID
+	}
+
 	userObj := PermitUser{
-		Key:       cognitoUser.SubjectID, // Use Cognito Subject ID as the key
-		Email:     cognitoUser.Email,
-		FirstName: cognitoUser.FirstName,
-		LastName:  cognitoUser.LastName,
-		Attributes: map[string]interface{}{
-			"cognito_username": cognitoUser.Username,
-			"cognito_status":   cognitoUser.Status,
-		},
+		Key:        cognitoUser.SubjectID, // Use Cognito Subject ID as the key
+		Email:      cognitoUser.Email,
+		FirstName:  cognitoUser.FirstName,
+		LastName:   cognitoUser.LastName,
+		Attributes: attributes,
 	}
 
 	jsonData, err := json.Marshal(userObj)

cmd/auth_related/user.creation.tool/readme.md

@@ -120,6 +120,11 @@ PERMIT_KEY=your-permit-api-key-here
 # AWS Cognito Configuration
 COGNITO_USER_POOL_ID=us-east-1_xxxxxxxxx
 
+# Environment Isolation (optional)
+# Tags created users with this environment ID in Cognito and Permit.io.
+# Can also be set via the -cognito-env-id flag (flag takes precedence).
+COGNITO_ENVIRONMENT_ID=acmehealth
+
 # AWS Configuration
 AWS_REGION=us-east-1
 AWS_ACCESS_KEY_ID=your-access-key        # Or use IAM role
@@ -216,6 +221,7 @@ To check what project/environment your API key is scoped to:
 - `-env`: Permit.io environment ID or key (required unless using `use-key` with an environment-scoped API key)
 - `-csv`: Path to CSV file containing users (required)
 - `-api-url`: Permit.io API base URL (default: https://api.permit.io)
+- `-cognito-env-id`: Environment ID to tag users with in Cognito (`custom:environment_id` attribute) and Permit.io. Falls back to `COGNITO_ENVIRONMENT_ID` environment variable if not provided. See [Environment Isolation](#environment-isolation) below.
 - `--delete-users`: Delete users instead of creating/updating them
 - `--disable-users`: Disable users in Cognito instead of creating/updating them
 - `--enable-users`: Enable users in Cognito instead of creating/updating them
@@ -223,6 +229,59 @@ To check what project/environment your API key is scoped to:
 
 **Important**: Only one operation mode can be specified at a time (`--delete-users`, `--disable-users`, or `--enable-users`). If none are specified, the tool defaults to create/update mode.
 
+## Environment Isolation
+
+When operating in a multi-environment deployment with a shared Cognito user pool, the tool supports tagging users with an environment identifier. This scopes users to a specific server environment so that each deployment only sees its own users.
+
+For a full description of the multi-environment user isolation feature, see the [cognitoauth readme](../../../internal/cognitoauth/readme.md#multi-environment-user-isolation).
+
+### Setting the Environment ID
+
+The environment ID can be provided in two ways (flag takes precedence):
+
+1. **Command line flag**: `-cognito-env-id acmehealth`
+2. **Environment variable**: `COGNITO_ENVIRONMENT_ID=acmehealth`
+
+The value must follow the naming rules: lowercase alphanumeric and underscores, must start with a letter, must end with a letter or digit, no consecutive underscores, max 50 characters. The tool validates the value at startup and exits with an error if invalid.
+
+### Usage Examples
+
+```bash
+# Create users tagged with environment ID via flag
+./cognito-permit-sync -project myproject -env production -csv users.csv -cognito-env-id acmehealth
+
+# Create users tagged with environment ID via env var
+export COGNITO_ENVIRONMENT_ID=acmehealth
+./cognito-permit-sync -project myproject -env production -csv users.csv
+
+# Dry run to preview environment tagging
+./cognito-permit-sync -project myproject -env production -csv users.csv -cognito-env-id acmehealth --dry-run
+```
+
+### What Gets Tagged
+
+When an environment ID is configured:
+
+- **Cognito**: Each created user gets a `custom:environment_id` attribute set to the configured value
+- **Permit.io**: Each created user gets an `environment_id` custom attribute set to the configured value
+
+When no environment ID is configured, users are created without any environment tag (backward compatible).
+
+### Listing Users by Environment
+
+The `list.users.by.env.sh` script lists all Cognito users that have a specific `custom:environment_id` attribute value:
+
+```bash
+# List users by environment ID passed as argument
+./list.users.by.env.sh <user-pool-id> [environment-id]
+
+# Falls back to COGNITO_ENVIRONMENT_ID env var if environment-id argument is omitted
+export COGNITO_ENVIRONMENT_ID=acmehealth
+./list.users.by.env.sh us-east-2_21upuTkkT
+```
+
+This script requires the AWS CLI and `jq` to be installed. It uses `aws cognito-idp list-users` and filters the results by the `custom:environment_id` attribute.
+
 ## How It Works
 
 ### Create/Update Mode (Default)
@@ -354,11 +413,14 @@ Before running this tool, the following must already exist in your Permit.io env
   "last_name": "Doe",
   "attributes": {
     "cognito_username": "user@example.com",
-    "cognito_status": "CONFIRMED"
+    "cognito_status": "CONFIRMED",
+    "environment_id": "acmehealth"
   }
 }
 ```
 
+The `environment_id` attribute is included only when `-cognito-env-id` or `COGNITO_ENVIRONMENT_ID` is configured. When not set, this attribute is omitted.
+
 ### What the Tool Deletes from Permit.io
 
 | Resource | API Endpoint | Details |
@@ -383,7 +445,7 @@ Before running this tool, the following must already exist in your Permit.io env
 | **User Key** | Uses Cognito `Subject ID` (the `sub` claim from JWT) as the Permit.io user key - **not** the email address |
 | **Tenant** | All role assignments use a hardcoded `"default"` tenant. Multi-tenant support is not currently implemented. |
 | **Role Sync** | Idempotent - removes roles no longer in CSV, adds missing ones, skips unchanged roles |
-| **User Attributes** | Stores `cognito_username` and `cognito_status` as custom attributes on the Permit.io user |
+| **User Attributes** | Stores `cognito_username`, `cognito_status`, and optionally `environment_id` as custom attributes on the Permit.io user |
 | **Role Validation** | Validates all roles exist **before** processing any users. If any role is missing, the tool exits with an error listing the missing roles. |
 
 ### Creating Required Roles in Permit.io

cmd/auth_related/user.creation.tool/run.tool.sh

@@ -1,5 +1,6 @@
 # permit stuff
 export PERMIT_KEY="permit_key_redacted"
+# development environment only
 
 # Cognito stuff
 export COGNITO_USER_POOL_ID="us-east-2_21upuTkkT"
@@ -12,5 +13,8 @@ export AWS_REGION=us-east-2
 export AWS_PROFILE=aarete
 # must do a project list to get the guid values for the project and environment before making the call to import.
 #go run ./... -project list
-go run ./... -project ebde1e1e9623491cab6f8112e67bd61c -env 9d6801123cfd4a0ea2ef1df2f430e9d3 -csv ./users.to.import.csv --dry-run
+#go run ./... -project ebde1e1e9623491cab6f8112e67bd61c -env 9d6801123cfd4a0ea2ef1df2f430e9d3 -csv ./users.to.import.csv --dry-run
 #go run ./... -project ebde1e1e9623491cab6f8112e67bd61c -env 9d6801123cfd4a0ea2ef1df2f430e9d3 -csv ./users.to.import.csv
+
+AWS_ENDPOINT_URL="" go run ./... -project use-key -csv users.to.import.csv -cognito-env-id dev 
+# --dry-run

cmd/auth_related/user.creation.tool/userlifecycle_test.go

@@ -349,6 +349,83 @@ func getCognitoUserSubject(ctx context.Context, client *cognitoidentityprovider.
 	return "", fmt.Errorf("subject ID not found for user %s", username)
 }
 
+// TestUserLifecycleWithEnvironmentID verifies that users created with COGNITO_ENVIRONMENT_ID
+// have the custom:environment_id attribute set in Cognito. Creates users with the env ID,
+// verifies the attribute, then cleans up.
+func TestUserLifecycleWithEnvironmentID(t *testing.T) {
+	if !checkRequiredEnvVars(t) {
+		return
+	}
+
+	const testCognitoEnvID = "integration-test-env"
+	envCSVPath := "./temp-envid-users.csv"
+
+	// Create a minimal CSV with one test user
+	envCSVContent := "email,first_name,last_name,role1\n"
+	envCSVContent += "envid-test-lifecycle@example.com,EnvTest,Lifecycle,viewer\n"
+	if err := os.WriteFile(envCSVPath, []byte(envCSVContent), 0600); err != nil {
+		t.Fatalf("Failed to create env test CSV: %v", err)
+	}
+	defer func() { _ = os.Remove(envCSVPath) }()
+
+	// Set COGNITO_ENVIRONMENT_ID for the tool to pick up
+	t.Setenv("COGNITO_ENVIRONMENT_ID", testCognitoEnvID)
+
+	t.Run("CreateUsersWithEnvID", func(t *testing.T) {
+		originalArgs := os.Args
+		defer func() { os.Args = originalArgs }()
+
+		os.Args = []string{
+			"cognito-permit-sync",
+			"-project", testProjectID,
+			"-env", testEnvID,
+			"-csv", envCSVPath,
+		}
+		resetFlags()
+
+		if err := run(); err != nil {
+			t.Fatalf("Failed to create users with environment ID: %v", err)
+		}
+	})
+
+	t.Run("VerifyEnvironmentIDAttribute", func(t *testing.T) {
+		ctx := t.Context()
+		cognitoClient, err := initializeCognitoClient()
+		if err != nil {
+			t.Fatalf("Failed to initialize Cognito client: %v", err)
+		}
+
+		userPoolID := os.Getenv("COGNITO_USER_POOL_ID")
+		user, err := getCognitoUser(ctx, cognitoClient, userPoolID, "envid-test-lifecycle@example.com")
+		if err != nil {
+			t.Fatalf("Failed to get user from Cognito: %v", err)
+		}
+
+		if user.EnvironmentID != testCognitoEnvID {
+			t.Fatalf("Expected environment_id=%q, got %q", testCognitoEnvID, user.EnvironmentID)
+		}
+		t.Logf("Verified user has environment_id=%q", user.EnvironmentID)
+	})
+
+	t.Run("CleanupEnvIDTestUsers", func(t *testing.T) {
+		originalArgs := os.Args
+		defer func() { os.Args = originalArgs }()
+
+		os.Args = []string{
+			"cognito-permit-sync",
+			"-project", testProjectID,
+			"-env", testEnvID,
+			"-csv", envCSVPath,
+			"--delete-users",
+		}
+		resetFlags()
+
+		if err := run(); err != nil {
+			t.Fatalf("Failed to delete env test users: %v", err)
+		}
+	})
+}
+
 // Helper function to reset flags for testing
 func resetFlags() {
 	// Reset the default command line flag set

cmd/queryAPI/main.go

@@ -31,9 +31,13 @@ import (
 	"queryorchestration/internal/label"
 	"queryorchestration/internal/server/api"
 	"queryorchestration/internal/serviceconfig"
+	awsc "queryorchestration/internal/serviceconfig/aws"
 	"queryorchestration/internal/serviceconfig/build"
 	"queryorchestration/internal/serviceconfig/objectstore"
 	"queryorchestration/internal/serviceconfig/queue/clientsync"
+	"queryorchestration/internal/usermanagement"
+
+	"github.com/aws/aws-sdk-go-v2/service/cognitoidentityprovider"
 
 	queryapi "queryorchestration/api/queryAPI"
 	documentbatch "queryorchestration/internal/document/batch"
@@ -134,8 +138,26 @@ func main() {
 		os.Exit(1)
 	}
 
+	// Create Cognito SDK client for environment attribute registration and middleware use.
+	// This client is created once and reused for the lifetime of the process.
+	awsCfg, err := awsc.GetAWSConfig(ctx)
+	if err != nil {
+		slog.Error("Failed to load AWS config for Cognito client", "error", err)
+		os.Exit(1)
+	}
+	cognitoClient := cognitoidentityprovider.NewFromConfig(awsCfg)
+
+	// Ensure the custom:environment_id attribute is registered on the Cognito user pool.
+	// This is safe to call on every startup (idempotent). Failure is non-fatal because
+	// the attribute may already exist or the environment may not support this API (e.g. LocalStack).
+	if err := usermanagement.EnsureEnvironmentIDAttribute(ctx, cognitoClient, cfg.GetAuthUserPoolID(), cfg.GetLogger()); err != nil {
+		cfg.GetLogger().Warn("Failed to ensure Cognito environment_id attribute", "error", err)
+	} else {
+		cfg.GetLogger().Info("Cognito custom attribute custom:environment_id is registered")
+	}
+
 	// Initialize S3 client before creating the server (needed for background worker)
-	err := cfg.SetStoreClient(ctx)
+	err = cfg.SetStoreClient(ctx)
 	if err != nil {
 		slog.Error(err.Error())
 		os.Exit(1)