How to set Cognito Groups in Migration trigger

If someone else runs into this - I solved this using a combination of a User Migration trigger and a Pre Token Generation trigger.

In the User Migration trigger (mostly copied from https://github.com/Collaborne/migrate-cognito-user-pool-lambda) look up and create the user if auth fails/user doesn't exist in the new pool.

In the Pre Token Generation trigger if the user hasn't been added to groups yet look up group membership in the old user pool (adminListGroupsForUser), add them to the new pool (adminAddUserToGroup). The crucial part is to override the group membership claims in the response so that they will be added to the token on the client side (groupsToOverride is just an array of the group names the user is part of):

event.response = {
    "claimsOverrideDetails": {
        "claimsToAddOrOverride": {
            
        },
        "groupOverrideDetails": {
            "groupsToOverride": groupsToOverride,
        }
    }
};

Thank you @BrokenGlass, I used this approach. For anyone else here's an example Typescript preTokenGeneration lambda.

//preTokenGenerations.ts
import { PreTokenGenerationTriggerHandler } from 'aws-lambda';
import { preTokenAuthentication } from '../services/preTokenService';

export const handler: PreTokenGenerationTriggerHandler = async (event, context) => {
  console.log({
    event,
    context,
    request: event.request,
    userAttributes: event.request.userAttributes,
    clientMetadata: event.request.clientMetadata,
    groupConfiguration: event.request.groupConfiguration,
  })

  const OLD_USER_POOL_ID = process.env.OLD_USER_POOL_ID;

  if (!OLD_USER_POOL_ID) {
    throw new Error("OLD_USER_POOL_ID is required for the lambda to work.")
  }

  const {
    userPoolId,
    request: {
      userAttributes: {
        email
      }
    },
    region
  } = event;

  switch (event.triggerSource) {
    case "TokenGeneration_Authentication":
      const groupsToOverride = await preTokenAuthentication({
        userPoolId,
        oldUserPoolId: OLD_USER_POOL_ID,
        username: email,
        region
      })

      event.response = {
        "claimsOverrideDetails": {
          "claimsToAddOrOverride": {

          },
          "groupOverrideDetails": {
            "groupsToOverride": groupsToOverride,
          }
        }
      };

      return event
    default:
      console.log(`Bad triggerSource ${event.triggerSource}`);
      return new Promise((resolve) => {
        resolve(event)
      });
  }
}

// preTokenService.ts
import { getUsersGroups, cognitoIdentityServiceProvider, assignUserToGroup } from "./cognito"

interface IPreTokenAuthentication {
  userPoolId: string;
  oldUserPoolId: string;
  username: string;
  region: string
}

export const preTokenAuthentication = async ({ userPoolId, oldUserPoolId, username, region }: IPreTokenAuthentication): string[] => {
  const cognitoISP = cognitoIdentityServiceProvider({ region });

  const newPoolUsersGroups = await getUsersGroups({
    cognitoISP,
    userPoolId,
    username
  });

  // If the user in the new pool already has groups assigned then exit
  if (newPoolUsersGroups.length !== 0) {
    console.log("No action required user already exists in a group")
    return;
  }

  const oldPoolUsersGroups = await getUsersGroups({
    cognitoISP,
    userPoolId: oldUserPoolId,
    username
  });

  // If the user in the old pool doesn't have any groups then nothing else for this function to do so exit.
  if (oldPoolUsersGroups.length === 0) {
    console.error("No action required user migrated user didn't belong to a group")
    return;
  }

  console.log({ oldPoolUsersGroups, newPoolUsersGroups })

  await assignUserToGroup({
    cognitoISP,
    userPoolId,
    username,
    groups: oldPoolUsersGroups
  })


  return oldPoolUsersGroups;
}

// cognito.ts
import { AdminAddUserToGroupRequest, AdminListGroupsForUserRequest } from "aws-sdk/clients/cognitoidentityserviceprovider";
import { CognitoIdentityServiceProvider } from 'aws-sdk';

interface ICognitoIdentityServiceProvider {
  region: string;
}

export const cognitoIdentityServiceProvider = ({ region }: ICognitoIdentityServiceProvider) => {
  const options: CognitoIdentityServiceProvider.Types.ClientConfiguration = {
    region,
  };
  const cognitoIdentityServiceProvider = new CognitoIdentityServiceProvider(options);
  return cognitoIdentityServiceProvider;
}
interface IGetUsersGroups {
  cognitoISP: CognitoIdentityServiceProvider,
  userPoolId: string,
  username: string
}

export const getUsersGroups = async ({ cognitoISP, userPoolId, username }: IGetUsersGroups): Promise<string[]> => {
  try {
    const params: AdminListGroupsForUserRequest = {
      UserPoolId: userPoolId,
      Username: username,
    }

    const response = await cognitoISP.adminListGroupsForUser(params).promise();
    return response.Groups?.map(group => group.GroupName!) || [];
  } catch (err) {
    console.error(err)
    return [];
  }
}

interface IAssignUserToGroup {
  cognitoISP: CognitoIdentityServiceProvider,
  username: string;
  groups: string[];
  userPoolId: string;
}

/**
 * Use Administration to assign a user to groups
 * @param {
 *  cognitoISP the cognito identity service provider to perform the action on
 *  userPoolId the userPool for which the user is being modified within
 *  username the username or email for which the action is to be performed
 *  groups the groups to assign the user too
 * }
 */
export const assignUserToGroup = async ({ cognitoISP, userPoolId, username, groups }: IAssignUserToGroup) => {
  console.log({ userPoolId, username, groups })
  for (const group of groups) {
    const params: AdminAddUserToGroupRequest = {
      UserPoolId: userPoolId,
      Username: username,
      GroupName: group
    };
    try {
      const response = await cognitoISP.adminAddUserToGroup(params).promise();
      console.log({ response })
    } catch (err) {
      console.error(err)
    }
  }
}

Tips, make sure under the trigger section in Cognito that you have the migration and preToken triggers set. You also need to ensure SRP is not enabled so the lambda can see the password to be able to successfully migrate the user.

Things to test is that when the user is first migrated that they are assigned their groups. And for future logins they are also assigned to their groups.

Let me know if anyone has any feedback or questions, happy to help.