Merge pull request #42 from aws-samples/fix/iam-session-caching

Fix caching Q client issue.

Author: sreedevs-aws

src/functions/slack-event-handler.ts

@@ -332,7 +332,15 @@ export const handler = async (
   }
 
   const [output, slackMessage] = await Promise.all([
-    chat(prompt, attachments, dependencies, slackEventsEnv, iamSessionCreds, context),
+    chat(
+      body.event.user,
+      prompt,
+      attachments,
+      dependencies,
+      slackEventsEnv,
+      iamSessionCreds,
+      context
+    ),
     dependencies.sendSlackMessage(
       slackEventsEnv,
       body.event.channel,

src/functions/slack-interaction-handler.ts

@@ -1,7 +1,6 @@
 import { APIGatewayProxyResult, Callback, Context } from 'aws-lambda';
 import {
   createModal,
-  getMarkdownBlock,
   getMarkdownBlocks,
   openModal,
   SLACK_ACTION,
@@ -193,6 +192,7 @@ export const handler = async (
       }
 
       await dependencies.submitFeedbackRequest(
+        payload.user.id,
         slackInteractionsEnv,
         iamSessionCreds,
         {

src/helpers/amazon-q/amazon-q-client.ts

@@ -3,35 +3,44 @@ import { SlackInteractionsEnv } from '@functions/slack-interaction-handler';
 import { makeLogger } from '@src/logging';
 import { v4 as uuid } from 'uuid';
 import {
-  QBusinessClient,
+  AttachmentInput,
   ChatSyncCommand,
+  ChatSyncCommandOutput,
+  MessageUsefulness,
+  MessageUsefulnessReason,
   PutFeedbackCommand,
   PutFeedbackCommandInput,
-  MessageUsefulnessReason,
-  MessageUsefulness,
   PutFeedbackCommandOutput,
-  ChatSyncCommandOutput,
-  AttachmentInput
+  QBusinessClient
 } from '@aws-sdk/client-qbusiness';
 import { Credentials } from 'aws-sdk';
 
 const logger = makeLogger('amazon-q-client');
 
-let amazonQClient: QBusinessClient | null = null;
+const amazonQClientBySlackUserId: { [key: string]: QBusinessClient } = {};
 
-export const getClient = (env: SlackEventsEnv, iamSessionCreds: Credentials) => {
-  if (amazonQClient === null) {
-    logger.debug(`Initiating AmazonQ client with region ${env.AMAZON_Q_REGION}`);
-    amazonQClient = new QBusinessClient({
-      credentials: iamSessionCreds,
-      region: env.AMAZON_Q_REGION
-    });
+export const getClient = (
+  env: SlackEventsEnv,
+  slackUserId: string,
+  iamSessionCreds: Credentials
+) => {
+  logger.debug(`Initiating AmazonQ client with region ${env.AMAZON_Q_REGION}`);
+  if (amazonQClientBySlackUserId[slackUserId]) {
+    return amazonQClientBySlackUserId[slackUserId];
   }
 
-  return amazonQClient;
+  const newClient = new QBusinessClient({
+    credentials: iamSessionCreds,
+    region: env.AMAZON_Q_REGION
+  });
+
+  amazonQClientBySlackUserId[slackUserId] = newClient;
+
+  return newClient;
 };
 
 export const callClient = async (
+  slackUserId: string,
   message: string,
   attachments: AttachmentInput[],
   env: SlackEventsEnv,
@@ -50,10 +59,11 @@ export const callClient = async (
   };
 
   logger.debug(`callClient input ${JSON.stringify(input)}`);
-  return await getClient(env, iamSessionCreds).send(new ChatSyncCommand(input));
+  return await getClient(env, slackUserId, iamSessionCreds).send(new ChatSyncCommand(input));
 };
 
 export const submitFeedbackRequest = async (
+  slackUserId: string,
   env: SlackInteractionsEnv,
   iamSessionCreds: Credentials,
   context: {
@@ -76,7 +86,9 @@ export const submitFeedbackRequest = async (
   };
 
   logger.debug(`putFeedbackRequest input ${JSON.stringify(input)}`);
-  const response = await getClient(env, iamSessionCreds).send(new PutFeedbackCommand(input));
+  const response = await getClient(env, slackUserId, iamSessionCreds).send(
+    new PutFeedbackCommand(input)
+  );
   logger.debug(`putFeedbackRequest output ${JSON.stringify(response)}`);
 
   return response;

src/helpers/amazon-q/amazon-q-helpers.ts

@@ -15,6 +15,7 @@ const AMAZON_Q_MSG_LIMIT = 7000;
 const WARN_TRUNCATED = `| Please note that you do not have all the conversation history due to limitation`;
 
 export const chat = async (
+  slackUserId: string,
   incomingMessage: string,
   attachments: AttachmentInput[],
   dependencies: ChatDependencies,
@@ -37,6 +38,7 @@ export const chat = async (
         : incomingMessage;
 
     const response = await dependencies.callClient(
+      slackUserId,
       inputMessage,
       attachments,
       env,

src/helpers/idc/session-helpers.ts

@@ -277,8 +277,6 @@ const exchangeIdPTokenForIAMSessionCreds = async (
 
   const idcIdToken = idcResponse.idToken!;
 
-  console.log(`IdC response ${JSON.stringify(idcResponse)}`);
-
   // decode the jwt token
   const decodedToken = jwt.decode(idcIdToken, { complete: true });
 
@@ -294,10 +292,6 @@ const exchangeIdPTokenForIAMSessionCreds = async (
   // Extract 'sts:identity-context' claim using type assertion
   const identityContext = (decodedToken.payload as Payload)['sts:identity_context'];
 
-  console.log(identityContext);
-
-  logger.debug(`IdC response ${JSON.stringify(idcResponse)}`);
-
   // call sts assume role
   const stsClient = new STS({ region: env.region });
   const assumeRoleRequest: AssumeRoleRequest = {

tst/helpers/amazon-q/amazon-q-helpers.test.ts

@@ -14,7 +14,7 @@ import { ChatSyncCommandOutput } from '@aws-sdk/client-qbusiness';
 
 describe('AmazonQ helpers test', () => {
   test('Should get a response as block with context', async () => {
-    const response = await chat('message', [], MOCK_DEPENDENCIES, MOCK_ENV, MOCK_IAM_SESSION_CREDS);
+    const response = await chat('slackUserId', 'message', [], MOCK_DEPENDENCIES, MOCK_ENV, MOCK_IAM_SESSION_CREDS);
     expect(response).toEqual(amazonQValidResponse1);
   });