diff --git a/packages/amplify-graphql-conversation-transformer/src/resolvers/assistant-mutation-resolver-fn.js b/packages/amplify-graphql-conversation-transformer/src/resolvers/assistant-mutation-resolver-fn.js index d50e1c141b..288b483083 100644 --- a/packages/amplify-graphql-conversation-transformer/src/resolvers/assistant-mutation-resolver-fn.js +++ b/packages/amplify-graphql-conversation-transformer/src/resolvers/assistant-mutation-resolver-fn.js @@ -10,6 +10,7 @@ export function request(ctx) { const { conversationId, content, associatedUserMessageId } = ctx.args.input; const { owner } = ctx.args; const defaultValues = ctx.stash.defaultValues ?? {}; + const id = defaultValues.id; const message = { __typename: '[[CONVERSATION_MESSAGE_TYPE_NAME]]', @@ -23,26 +24,6 @@ export function request(ctx) { }; return ddb.put({ key: { id }, item: message }); - - // const expression = 'SET #assistantContent = :assistantContent, #updatedAt = :updatedAt'; - // const expressionNames = { '#assistantContent': 'assistantContent', '#updatedAt': 'updatedAt' }; - // const expressionValues = { ':assistantContent': content, ':updatedAt': updatedAt }; - // const condition = JSON.parse( - // util.transform.toDynamoDBConditionExpression({ - // owner: { eq: owner }, - // conversationId: { eq: conversationId }, - // }), - // ); - // return { - // operation: 'UpdateItem', - // key: util.dynamodb.toMapValues({ id: associatedUserMessageId }), - // condition, - // update: { - // expression, - // expressionNames, - // expressionValues: util.dynamodb.toMapValues(expressionValues), - // }, - // }; } /** @@ -56,16 +37,5 @@ export function response(ctx) { util.error(ctx.error.message, ctx.error.type); } - const { conversationId, content, associatedUserMessageId } = ctx.args.input; - const { createdAt, updatedAt } = ctx.result; - - return { - id: associatedUserMessageId, - content, - conversationId, - role: 'assistant', - owner: ctx.stash.owner, - createdAt, - updatedAt, - }; + return ctx.result; } diff --git a/packages/amplify-graphql-conversation-transformer/src/resolvers/verify-session-owner-resolver-fn.js b/packages/amplify-graphql-conversation-transformer/src/resolvers/verify-session-owner-resolver-fn.js index 7a05ff10b7..25ce987425 100644 --- a/packages/amplify-graphql-conversation-transformer/src/resolvers/verify-session-owner-resolver-fn.js +++ b/packages/amplify-graphql-conversation-transformer/src/resolvers/verify-session-owner-resolver-fn.js @@ -1,10 +1,11 @@ export function request(ctx) { const { authFilter } = ctx.stash; + const { conversationId } = [[CONVERSATION_ID_PARENT]]; const query = { expression: 'id = :id', expressionValues: util.dynamodb.toMapValues({ - ':id': ctx.args.conversationId, + ':id': conversationId, }), }; diff --git a/packages/amplify-graphql-conversation-transformer/src/resolvers/verify-session-owner-resolver.ts b/packages/amplify-graphql-conversation-transformer/src/resolvers/verify-session-owner-resolver.ts index 8b8eb5b652..873bd270fb 100644 --- a/packages/amplify-graphql-conversation-transformer/src/resolvers/verify-session-owner-resolver.ts +++ b/packages/amplify-graphql-conversation-transformer/src/resolvers/verify-session-owner-resolver.ts @@ -9,8 +9,33 @@ import { ConversationDirectiveConfiguration } from '../grapqhl-conversation-tran * * @returns {MappingTemplateProvider} An object containing request and response MappingTemplateProviders. */ -export const verifySessionOwnerMappingTemplate = (config: ConversationDirectiveConfiguration): MappingTemplateProvider => { - const resolver = fs.readFileSync(path.join(__dirname, 'verify-session-owner-resolver-fn.js'), 'utf8'); +export const verifySessionOwnerSendMessageMappingTemplate = (config: ConversationDirectiveConfiguration): MappingTemplateProvider => { + const substitutions = { + CONVERSATION_ID_PARENT: 'ctx.args', + }; const templateName = `Mutation.${config.field.name.value}.verify-session-owner.js`; + return verifySessionOwnerMappingTemplate(templateName, substitutions); +}; + +/** + * Creates a mapping template for verifying the session owner in a conversation. + * + * @returns {MappingTemplateProvider} An object containing request and response MappingTemplateProviders. + */ +export const verifySessionOwnerAssistantResponseMappingTemplate = (config: ConversationDirectiveConfiguration): MappingTemplateProvider => { + const substitutions = { + CONVERSATION_ID_PARENT: 'ctx.args.input', + }; + const templateName = `Mutation.${config.field.name.value}AssistantResponse.verify-session-owner.js`; + return verifySessionOwnerMappingTemplate(templateName, substitutions); +}; + +const verifySessionOwnerMappingTemplate = (name: string, substitute: Record) => { + let resolver = fs.readFileSync(path.join(__dirname, 'verify-session-owner-resolver-fn.js'), 'utf8'); + Object.entries(substitute).forEach(([key, value]) => { + const replaced = resolver.replace(new RegExp(`\\[\\[${key}\\]\\]`, 'g'), value); + resolver = replaced; + }); + const templateName = `Mutation.${name}.verify-session-owner.js`; return MappingTemplate.s3MappingFunctionCodeFromString(resolver, templateName); }; diff --git a/packages/amplify-graphql-conversation-transformer/src/transformer-steps/conversation-resolver-generator.ts b/packages/amplify-graphql-conversation-transformer/src/transformer-steps/conversation-resolver-generator.ts index 7ab6f95e1a..1d24da48dd 100644 --- a/packages/amplify-graphql-conversation-transformer/src/transformer-steps/conversation-resolver-generator.ts +++ b/packages/amplify-graphql-conversation-transformer/src/transformer-steps/conversation-resolver-generator.ts @@ -9,7 +9,10 @@ import { IFunction, Function } from 'aws-cdk-lib/aws-lambda'; import { getModelDataSourceNameForTypeName, getTable } from '@aws-amplify/graphql-transformer-core'; import { initMappingTemplate } from '../resolvers/init-resolver'; import { authMappingTemplate } from '../resolvers/auth-resolver'; -import { verifySessionOwnerMappingTemplate } from '../resolvers/verify-session-owner-resolver'; +import { + verifySessionOwnerSendMessageMappingTemplate, + verifySessionOwnerAssistantResponseMappingTemplate, +} from '../resolvers/verify-session-owner-resolver'; import { writeMessageToTableMappingTemplate } from '../resolvers/write-message-to-table-resolver'; import { invokeLambdaMappingTemplate } from '../resolvers/invoke-lambda-resolver'; import { assistantMutationResolver } from '../resolvers/assistant-mutation-resolver'; @@ -54,8 +57,8 @@ export class ConversationResolverGenerator { this.setupMessageTableIndex(ctx, directive); const initResolverFunction = initMappingTemplate(directive); const authResolverFunction = authMappingTemplate(directive); - const verifySessionOwnerResolverFunction = verifySessionOwnerMappingTemplate(directive); - + const verifySessionOwnerSendMessageResolverFunction = verifySessionOwnerSendMessageMappingTemplate(directive); + const verifySessionOwnerAssistantResponseResolverFunction = verifySessionOwnerAssistantResponseMappingTemplate(directive); this.createConversationPipelineResolver( ctx, @@ -66,10 +69,17 @@ export class ConversationResolverGenerator { invokeLambdaFunction, initResolverFunction, authResolverFunction, - verifySessionOwnerResolverFunction, + verifySessionOwnerSendMessageResolverFunction, ); - this.createAssistantResponseResolver(ctx, directive, capitalizedFieldName, initResolverFunction, authResolverFunction, verifySessionOwnerResolverFunction); + this.createAssistantResponseResolver( + ctx, + directive, + capitalizedFieldName, + initResolverFunction, + authResolverFunction, + verifySessionOwnerAssistantResponseResolverFunction, + ); this.createAssistantResponseSubscriptionResolver(ctx, directive, capitalizedFieldName); } @@ -193,7 +203,14 @@ export class ConversationResolverGenerator { runtime, ); - this.addPipelineResolverFunctions(ctx, conversationPipelineResolver, capitalizedFieldName, initResolverFunction, authResolverFunction, verifySessionOwnerResolverFunction); + this.addPipelineResolverFunctions( + ctx, + conversationPipelineResolver, + capitalizedFieldName, + initResolverFunction, + authResolverFunction, + verifySessionOwnerResolverFunction, + ); ctx.resolvers.addResolver(parentName, fieldName, conversationPipelineResolver); } @@ -256,13 +273,12 @@ export class ConversationResolverGenerator { directive.responseMutationName, assistantResponseResolverResourceId, { codeMappingTemplate: assistantResponseResolverFunction }, - ['init', 'auth', 'verifySessionOwner',], + ['init', 'auth', 'verifySessionOwner'], [], conversationMessageDataSource as any, APPSYNC_JS_RUNTIME, ); - // Add init function resolver.addJsFunctionToSlot('init', initResolverFunction);