diff --git a/ee/app/assets/javascripts/ai/tanuki_bot/components/app.vue b/ee/app/assets/javascripts/ai/tanuki_bot/components/app.vue
index e98c62eadc45c0654086d0bf1a3734ce7a6848a7..0f5bf4c984096f542b50281d890be29a2c19f7a9 100644
--- a/ee/app/assets/javascripts/ai/tanuki_bot/components/app.vue
+++ b/ee/app/assets/javascripts/ai/tanuki_bot/components/app.vue
@@ -8,8 +8,6 @@ import { renderGFM } from '~/behaviors/markdown/render_gfm';
import { helpPagePath } from '~/helpers/help_page_helper';
import { duoChatGlobalState } from '~/super_sidebar/constants';
import { clearDuoChatCommands } from 'ee/ai/utils';
-import aiResponseSubscription from 'ee/graphql_shared/subscriptions/ai_completion_response.subscription.graphql';
-import aiResponseStreamSubscription from 'ee/graphql_shared/subscriptions/ai_completion_response_stream.subscription.graphql';
import DuoChatCallout from 'ee/ai/components/global_callout/duo_chat_callout.vue';
import getAiMessages from 'ee/ai/graphql/get_ai_messages.query.graphql';
import chatMutation from 'ee/ai/graphql/chat.mutation.graphql';
@@ -21,7 +19,8 @@ import {
GENIE_CHAT_CLEAN_MESSAGE,
GENIE_CHAT_CLEAR_MESSAGE,
} from 'ee/ai/constants';
-import { TANUKI_BOT_TRACKING_EVENT_NAME, MESSAGE_TYPES, SLASH_COMMANDS } from '../constants';
+import { TANUKI_BOT_TRACKING_EVENT_NAME, SLASH_COMMANDS, MESSAGE_TYPES } from '../constants';
+import TanukiBotSubscriptions from './tanuki_bot_subscriptions.vue';
export default {
name: 'TanukiBotChatApp',
@@ -46,6 +45,7 @@ export default {
components: {
GlDuoChat,
DuoChatCallout,
+ TanukiBotSubscriptions,
},
mixins: [Tracking.mixin()],
provide() {
@@ -65,72 +65,6 @@ export default {
},
},
apollo: {
- $subscribe: {
- // eslint-disable-next-line @gitlab/vue-no-undef-apollo-properties
- aiCompletionResponse: {
- query: aiResponseSubscription,
- variables() {
- return {
- userId: this.userId,
- aiAction: 'CHAT',
- };
- },
- result({ data }) {
- const requestId = data?.aiCompletionResponse?.requestId;
- if (requestId && !this.cancelledRequestIds.includes(requestId)) {
- this.addDuoChatMessage(data.aiCompletionResponse);
- if (data.aiCompletionResponse.role.toLowerCase() === MESSAGE_TYPES.TANUKI) {
- this.responseCompleted = requestId;
- clearDuoChatCommands();
- }
- }
- },
- error(err) {
- this.addDuoChatMessage({ errors: [err.toString()] });
- },
- skip() {
- return !this.duoChatGlobalState.isShown;
- },
- },
- // eslint-disable-next-line @gitlab/vue-no-undef-apollo-properties
- aiCompletionResponseStream: {
- query: aiResponseStreamSubscription,
- variables() {
- return {
- userId: this.userId,
- clientSubscriptionId: this.clientSubscriptionId,
- };
- },
- result({ data }) {
- const requestId = data?.aiCompletionResponse?.requestId;
- if (
- requestId &&
- requestId !== this.responseCompleted &&
- !this.cancelledRequestIds.includes(requestId)
- ) {
- this.addDuoChatMessage(data.aiCompletionResponse);
- }
- if (data?.aiCompletionResponse?.chunkId && !this.isResponseTracked) {
- performance.mark('response-received');
- performance.measure('prompt-to-response', 'prompt-sent', 'response-received');
- const [{ duration }] = performance.getEntriesByName('prompt-to-response');
- this.track('ai_response_time', {
- property: requestId,
- value: duration,
- });
- performance.clearMarks();
- performance.clearMeasures();
- this.isResponseTracked = true;
- }
- },
- error(err) {
- this.addDuoChatMessage({ errors: [err.toString()] });
- },
- skip() {
- return !this.duoChatGlobalState.isShown;
- },
- },
- },
// eslint-disable-next-line @gitlab/vue-no-undef-apollo-properties
aiMessages: {
query: getAiMessages,
@@ -140,7 +74,7 @@ export default {
}
},
error(err) {
- this.addDuoChatMessage({ errors: [err.toString()] });
+ this.onError(err);
},
},
},
@@ -150,9 +84,9 @@ export default {
clientSubscriptionId: uuidv4(),
toolName: i18n.GITLAB_DUO,
error: '',
- responseCompleted: undefined,
isResponseTracked: false,
cancelledRequestIds: [],
+ completedRequestId: null,
};
},
computed: {
@@ -192,9 +126,39 @@ export default {
this.cancelledRequestIds.push(this.messages[this.messages.length - 1].requestId);
this.setLoading(false);
},
+ onMessageReceived(aiCompletionResponse) {
+ this.addDuoChatMessage(aiCompletionResponse);
+ if (aiCompletionResponse.role.toLowerCase() === MESSAGE_TYPES.TANUKI) {
+ this.completedRequestId = aiCompletionResponse.requestId;
+ clearDuoChatCommands();
+ }
+ },
+ onMessageStreamReceived(aiCompletionResponse) {
+ if (aiCompletionResponse.requestId !== this.completedRequestId) {
+ this.addDuoChatMessage(aiCompletionResponse);
+ }
+ },
+ onResponseReceived(requestId) {
+ if (this.isResponseTracked) {
+ return;
+ }
+
+ performance.mark('response-received');
+ performance.measure('prompt-to-response', 'prompt-sent', 'response-received');
+ const [{ duration }] = performance.getEntriesByName('prompt-to-response');
+
+ this.track('ai_response_time', {
+ property: requestId,
+ value: duration,
+ });
+
+ performance.clearMarks();
+ performance.clearMeasures();
+ this.isResponseTracked = true;
+ },
onSendChatPrompt(question, variables = {}) {
- this.responseCompleted = undefined;
performance.mark('prompt-sent');
+ this.completedRequestId = null;
this.isResponseTracked = false;
if (!this.isClearOrResetMessage(question)) {
@@ -229,7 +193,7 @@ export default {
this.addDuoChatMessage({
content: question,
});
- this.addDuoChatMessage({ errors: [err.toString()] });
+ this.onError(err);
this.setLoading(false);
});
},
@@ -274,30 +238,45 @@ export default {
});
}
},
+ onError(err) {
+ this.addDuoChatMessage({ errors: [err.toString()] });
+ },
},
};
</script>
<template>
<div>
- <gl-duo-chat
- v-if="duoChatGlobalState.isShown"
- id="duo-chat"
- :slash-commands="$options.SLASH_COMMANDS"
- :title="$options.i18n.gitlabChat"
- :messages="messages"
- :error="error"
- :is-loading="loading"
- :predefined-prompts="$options.i18n.predefinedPrompts"
- :badge-type="null"
- :tool-name="toolName"
- :canceled-request-ids="cancelledRequestIds"
- class="duo-chat-container"
- @chat-cancel="onChatCancel"
- @send-chat-prompt="onSendChatPrompt"
- @chat-hidden="onChatClose"
- @track-feedback="onTrackFeedback"
- />
+ <div v-if="duoChatGlobalState.isShown">
+ <!-- Renderless component for subscriptions -->
+ <tanuki-bot-subscriptions
+ :user-id="userId"
+ :client-subscription-id="clientSubscriptionId"
+ :cancelled-request-ids="cancelledRequestIds"
+ @message="onMessageReceived"
+ @message-stream="onMessageStreamReceived"
+ @response-received="onResponseReceived"
+ @error="onError"
+ />
+
+ <gl-duo-chat
+ id="duo-chat"
+ :slash-commands="$options.SLASH_COMMANDS"
+ :title="$options.i18n.gitlabChat"
+ :messages="messages"
+ :error="error"
+ :is-loading="loading"
+ :predefined-prompts="$options.i18n.predefinedPrompts"
+ :badge-type="null"
+ :tool-name="toolName"
+ :canceled-request-ids="cancelledRequestIds"
+ class="duo-chat-container"
+ @chat-cancel="onChatCancel"
+ @send-chat-prompt="onSendChatPrompt"
+ @chat-hidden="onChatClose"
+ @track-feedback="onTrackFeedback"
+ />
+ </div>
<duo-chat-callout @callout-dismissed="onCalloutDismissed" />
</div>
</template>
diff --git a/ee/app/assets/javascripts/ai/tanuki_bot/components/tanuki_bot_subscriptions.vue b/ee/app/assets/javascripts/ai/tanuki_bot/components/tanuki_bot_subscriptions.vue
new file mode 100644
index 0000000000000000000000000000000000000000..22bed07eb8a114aaad361013aa726c69985c1165
--- /dev/null
+++ b/ee/app/assets/javascripts/ai/tanuki_bot/components/tanuki_bot_subscriptions.vue
@@ -0,0 +1,73 @@
+<script>
+import aiResponseSubscription from 'ee/graphql_shared/subscriptions/ai_completion_response.subscription.graphql';
+import aiResponseStreamSubscription from 'ee/graphql_shared/subscriptions/ai_completion_response_stream.subscription.graphql';
+
+export default {
+ props: {
+ userId: {
+ type: String,
+ required: true,
+ },
+ clientSubscriptionId: {
+ type: String,
+ required: true,
+ },
+ cancelledRequestIds: {
+ type: Array,
+ default: () => [],
+ required: false,
+ },
+ },
+ render() {
+ return null;
+ },
+ apollo: {
+ $subscribe: {
+ // eslint-disable-next-line @gitlab/vue-no-undef-apollo-properties
+ aiCompletionResponse: {
+ query: aiResponseSubscription,
+ variables() {
+ return {
+ userId: this.userId,
+ aiAction: 'CHAT',
+ };
+ },
+ result({ data }) {
+ const requestId = data?.aiCompletionResponse?.requestId;
+
+ if (requestId && !this.cancelledRequestIds.includes(requestId)) {
+ this.$emit('message', data.aiCompletionResponse);
+ }
+ },
+ error(err) {
+ this.$emit('error', err);
+ },
+ },
+ // eslint-disable-next-line @gitlab/vue-no-undef-apollo-properties
+ aiCompletionResponseStream: {
+ query: aiResponseStreamSubscription,
+ variables() {
+ return {
+ userId: this.userId,
+ clientSubscriptionId: this.clientSubscriptionId,
+ };
+ },
+ result({ data }) {
+ const requestId = data?.aiCompletionResponse?.requestId;
+
+ if (requestId && !this.cancelledRequestIds.includes(requestId)) {
+ this.$emit('message-stream', data.aiCompletionResponse);
+ }
+
+ if (data?.aiCompletionResponse?.chunkId) {
+ this.$emit('response-received', requestId);
+ }
+ },
+ error(err) {
+ this.$emit('error', err);
+ },
+ },
+ },
+ },
+};
+</script>
diff --git a/ee/spec/frontend/ai/tanuki_bot/components/app_spec.js b/ee/spec/frontend/ai/tanuki_bot/components/app_spec.js
index 134d74bf9b829dc5820e94b51814aa8d0157f00a..e364290a7d294ded122ecc9c0b9f92186a74463d 100644
--- a/ee/spec/frontend/ai/tanuki_bot/components/app_spec.js
+++ b/ee/spec/frontend/ai/tanuki_bot/components/app_spec.js
@@ -4,18 +4,16 @@ import { v4 as uuidv4 } from 'uuid';
// eslint-disable-next-line no-restricted-imports
import Vuex from 'vuex';
import VueApollo from 'vue-apollo';
-import { createMockSubscription } from 'mock-apollo-client';
import { sendDuoChatCommand } from 'ee/ai/utils';
import TanukiBotChatApp from 'ee/ai/tanuki_bot/components/app.vue';
import DuoChatCallout from 'ee/ai/components/global_callout/duo_chat_callout.vue';
+import TanukiBotSubscriptions from 'ee/ai/tanuki_bot/components/tanuki_bot_subscriptions.vue';
import {
GENIE_CHAT_RESET_MESSAGE,
GENIE_CHAT_CLEAN_MESSAGE,
GENIE_CHAT_CLEAR_MESSAGE,
} from 'ee/ai/constants';
import { TANUKI_BOT_TRACKING_EVENT_NAME } from 'ee/ai/tanuki_bot/constants';
-import aiResponseSubscription from 'ee/graphql_shared/subscriptions/ai_completion_response.subscription.graphql';
-import aiResponseStreamSubscription from 'ee/graphql_shared/subscriptions/ai_completion_response_stream.subscription.graphql';
import chatMutation from 'ee/ai/graphql/chat.mutation.graphql';
import duoUserFeedbackMutation from 'ee/ai/graphql/duo_user_feedback.mutation.graphql';
import getAiMessages from 'ee/ai/graphql/get_ai_messages.query.graphql';
@@ -31,11 +29,10 @@ import {
MOCK_USER_MESSAGE,
MOCK_USER_ID,
MOCK_RESOURCE_ID,
- MOCK_TANUKI_SUCCESS_RES,
+ MOCK_CHUNK_MESSAGE,
MOCK_TANUKI_BOT_MUTATATION_RES,
- MOCK_CHAT_CACHED_MESSAGES_RES,
GENERATE_MOCK_TANUKI_RES,
- MOCK_CHUNK_MESSAGE,
+ MOCK_CHAT_CACHED_MESSAGES_RES,
MOCK_SLASH_COMMANDS,
} from '../mock_data';
@@ -54,14 +51,14 @@ const skipReason = new SkipReason({
describeSkipVue3(skipReason, () => {
let wrapper;
+ const UUIDMOCK = '123';
+
const actionSpies = {
addDuoChatMessage: jest.fn(),
setMessages: jest.fn(),
setLoading: jest.fn(),
};
- let aiResponseSubscriptionHandler = jest.fn();
- let aiResponseStreamSubscriptionHandler = jest.fn();
const chatMutationHandlerMock = jest.fn().mockResolvedValue(MOCK_TANUKI_BOT_MUTATATION_RES);
const duoUserFeedbackMutationHandlerMock = jest.fn().mockResolvedValue({});
const queryHandlerMock = jest.fn().mockResolvedValue(MOCK_CHAT_CACHED_MESSAGES_RES);
@@ -82,6 +79,7 @@ describeSkipVue3(skipReason, () => {
};
const findCallout = () => wrapper.findComponent(DuoChatCallout);
+ const findSubscriptions = () => wrapper.findComponent(TanukiBotSubscriptions);
const createComponent = ({
initialState = {},
@@ -100,16 +98,6 @@ describeSkipVue3(skipReason, () => {
[getAiMessages, queryHandlerMock],
]);
- apolloProvider.defaultClient.setRequestHandler(
- aiResponseSubscription,
- aiResponseSubscriptionHandler,
- );
-
- apolloProvider.defaultClient.setRequestHandler(
- aiResponseStreamSubscription,
- aiResponseStreamSubscriptionHandler,
- );
-
wrapper = shallowMountExtended(TanukiBotChatApp, {
store,
apolloProvider,
@@ -118,16 +106,10 @@ describeSkipVue3(skipReason, () => {
};
const findGlDuoChat = () => wrapper.findComponent(GlDuoChat);
- let perfTrackingSpy;
+
beforeEach(() => {
- uuidv4.mockImplementation(() => '123');
+ uuidv4.mockImplementation(() => UUIDMOCK);
getMarkdown.mockImplementation(({ text }) => Promise.resolve({ data: { html: text } }));
-
- performance.mark = jest.fn();
- performance.measure = jest.fn();
- performance.getEntriesByName = jest.fn(() => [{ duration: 123 }]);
- performance.clearMarks = jest.fn();
- performance.clearMeasures = jest.fn();
});
afterEach(() => {
@@ -196,7 +178,7 @@ describeSkipVue3(skipReason, () => {
describe('when new commands are added to the global state', () => {
beforeEach(() => {
createComponent();
- perfTrackingSpy = mockTracking(undefined, wrapper.element, jest.spyOn);
+ mockTracking(undefined, wrapper.element, jest.spyOn);
performance.mark = jest.fn();
});
@@ -241,7 +223,7 @@ describeSkipVue3(skipReason, () => {
describe('@send-chat-prompt', () => {
beforeEach(() => {
- perfTrackingSpy = mockTracking(undefined, wrapper.element, jest.spyOn);
+ mockTracking(undefined, wrapper.element, jest.spyOn);
performance.mark = jest.fn();
});
@@ -338,6 +320,41 @@ describeSkipVue3(skipReason, () => {
});
});
+ describe('@response-received', () => {
+ let trackingSpy;
+ beforeEach(() => {
+ trackingSpy = mockTracking(undefined, wrapper.element, jest.spyOn);
+ performance.mark = jest.fn();
+ performance.measure = jest.fn();
+ performance.getEntriesByName = jest.fn(() => [{ duration: 123 }]);
+ performance.clearMarks = jest.fn();
+ performance.clearMeasures = jest.fn();
+ });
+
+ afterEach(() => {
+ unmockTracking();
+ });
+
+ it('tracks time to response on first response-received', () => {
+ findSubscriptions().vm.$emit('response-received', 'request-id-123');
+
+ expect(performance.mark).toHaveBeenCalledWith('response-received');
+
+ expect(trackingSpy).toHaveBeenCalledWith(undefined, 'ai_response_time', {
+ property: 'request-id-123',
+ value: 123,
+ });
+ });
+
+ it('does not track time to response after first chunk was tracked', () => {
+ findSubscriptions().vm.$emit('response-received', 'request-id-123');
+ findSubscriptions().vm.$emit('response-received', 'request-id-123');
+
+ expect(performance.mark).toHaveBeenCalledTimes(1);
+ expect(trackingSpy).toHaveBeenCalledTimes(1);
+ });
+ });
+
describe('@track-feedback', () => {
it('calls the feedback GraphQL mutation when message is passed', async () => {
createComponent();
@@ -421,17 +438,7 @@ describeSkipVue3(skipReason, () => {
});
});
- describe('Subscriptions', () => {
- let mockSubscriptionComplete;
- let mockSubscriptionStream;
-
- beforeEach(() => {
- mockSubscriptionComplete = createMockSubscription();
- mockSubscriptionStream = createMockSubscription();
- aiResponseSubscriptionHandler = () => mockSubscriptionComplete;
- aiResponseStreamSubscriptionHandler = () => mockSubscriptionStream;
- });
-
+ describe('Subscription Component', () => {
afterEach(() => {
duoChatGlobalState.isShown = false;
if (wrapper) {
@@ -440,203 +447,112 @@ describeSkipVue3(skipReason, () => {
jest.clearAllMocks();
});
- it('activates subscriptions when isShown is true', async () => {
+ it('renders AiResponseSubscription component with correct props when isShown is true', async () => {
duoChatGlobalState.isShown = true;
createComponent();
await waitForPromises();
- expect(mockSubscriptionComplete.closed).toBe(false);
- expect(mockSubscriptionStream.closed).toBe(false);
+ expect(findSubscriptions().exists()).toBe(true);
+ expect(findSubscriptions().props('userId')).toBe(MOCK_USER_ID);
+ expect(findSubscriptions().props('clientSubscriptionId')).toBe(UUIDMOCK);
+ expect(findSubscriptions().props('cancelledRequestIds')).toHaveLength(0);
});
- it('does not activate subscriptions when isShown is false', async () => {
+ it('does not render AiResponseSubscription component when isShown is false', async () => {
duoChatGlobalState.isShown = false;
createComponent();
await waitForPromises();
- expect(mockSubscriptionComplete.closed).toBe(true);
- expect(mockSubscriptionStream.closed).toBe(true);
- });
-
- it('stops adding new messages when more chunks with the same request ID come in after the full message has already been received', async () => {
- const requestId = '123';
- const firstChunk = MOCK_CHUNK_MESSAGE('first chunk', 1, requestId);
- const secondChunk = MOCK_CHUNK_MESSAGE('second chunk', 2, requestId);
- const successResponse = GENERATE_MOCK_TANUKI_RES('', requestId);
-
- duoChatGlobalState.isShown = true;
-
- createComponent();
- await waitForPromises();
-
- // message chunk streaming in
- mockSubscriptionStream.next(firstChunk);
- await waitForPromises();
- expect(actionSpies.addDuoChatMessage).toHaveBeenCalledTimes(1);
- expect(actionSpies.addDuoChatMessage).toHaveBeenCalledWith(
- expect.any(Object),
- firstChunk.data.aiCompletionResponse,
- );
-
- // full message being sent
- mockSubscriptionComplete.next(successResponse);
- await waitForPromises();
- expect(actionSpies.addDuoChatMessage).toHaveBeenCalledTimes(2);
- expect(actionSpies.addDuoChatMessage).toHaveBeenCalledWith(
- expect.any(Object),
- successResponse.data.aiCompletionResponse,
- );
-
- // another chunk with the same request ID
- mockSubscriptionStream.next(secondChunk);
- await waitForPromises();
- // checking that addDuoChatMessage was not called again since full message was already being sent
- expect(actionSpies.addDuoChatMessage).toHaveBeenCalledTimes(2);
- });
-
- it('clears the commands when streaming is done', async () => {
- const requestId = '123';
- const firstChunk = MOCK_CHUNK_MESSAGE('first chunk', 1, requestId);
- const successResponse = GENERATE_MOCK_TANUKI_RES('', requestId);
-
- duoChatGlobalState.isShown = true;
-
- expect(duoChatGlobalState.commands).toHaveLength(0);
- sendDuoChatCommand({ question: '/troubleshoot', resourceId: '1' });
- expect(duoChatGlobalState.commands).toHaveLength(1);
-
- createComponent();
- await waitForPromises();
-
- // message chunk streaming in
- mockSubscriptionStream.next(firstChunk);
- await waitForPromises();
- // No changes to commands
- expect(duoChatGlobalState.commands).toHaveLength(1);
-
- // full message being sent
- mockSubscriptionComplete.next(successResponse);
- await waitForPromises();
- await waitForPromises();
-
- // commands have been cleared out
- expect(duoChatGlobalState.commands).toHaveLength(0);
+ expect(findSubscriptions().exists()).toBe(false);
});
- it('continues to invoke addDuoChatMessage when a new message chunk arrives with a distinct request ID, even after a complete message has been received', async () => {
- const firstChunk = MOCK_CHUNK_MESSAGE('first chunk', 1);
- const firstChunkNewRequest = MOCK_CHUNK_MESSAGE('first chunk', 2, 2);
-
+ it('calls addDuoChatMessage when @message is fired', () => {
duoChatGlobalState.isShown = true;
createComponent();
- await waitForPromises();
-
- // message chunk streaming in
- mockSubscriptionStream.next(firstChunk);
- expect(actionSpies.addDuoChatMessage).toHaveBeenCalledWith(
- expect.any(Object),
- firstChunk.data.aiCompletionResponse,
- );
-
- // full message being sent
- mockSubscriptionComplete.next(MOCK_TANUKI_SUCCESS_RES);
- await waitForPromises();
- expect(actionSpies.addDuoChatMessage).toHaveBeenCalledTimes(2);
- expect(actionSpies.addDuoChatMessage).toHaveBeenCalledWith(
- expect.any(Object),
- MOCK_TANUKI_SUCCESS_RES.data.aiCompletionResponse,
- );
+ const mockMessage = {
+ content: 'test message content',
+ role: 'user',
+ };
- // another chunk with a new request ID
- mockSubscriptionStream.next(firstChunkNewRequest);
- await waitForPromises();
- expect(actionSpies.addDuoChatMessage).toHaveBeenCalledTimes(3);
+ findSubscriptions().vm.$emit('message', mockMessage);
+ expect(actionSpies.addDuoChatMessage).toHaveBeenCalledWith(expect.anything(), mockMessage);
});
- it('stops streaming in new chunks when requestId was canceled', async () => {
- const requestId = '123';
- const firstChunk = MOCK_CHUNK_MESSAGE('first chunk', 1, requestId);
- const secondChunk = MOCK_CHUNK_MESSAGE('second chunk', 2, requestId);
-
- duoChatGlobalState.isShown = true;
-
- createComponent({
- initialState: {
- messages: [
- {
- requestId,
- },
- ],
- },
+ describe('Subscription Component', () => {
+ beforeEach(() => {
+ duoChatGlobalState.isShown = true;
+ createComponent();
+ mockTracking(undefined, wrapper.element, jest.spyOn);
+ performance.mark = jest.fn();
});
- await waitForPromises();
- perfTrackingSpy = mockTracking(undefined, wrapper.element, jest.spyOn);
- // message chunk streaming in
- mockSubscriptionStream.next(firstChunk);
- await waitForPromises();
- expect(actionSpies.addDuoChatMessage).toHaveBeenCalledTimes(1);
- expect(actionSpies.addDuoChatMessage).toHaveBeenCalledWith(
- expect.any(Object),
- firstChunk.data.aiCompletionResponse,
- );
- findGlDuoChat().vm.$emit('chat-cancel');
+ it('stops adding new messages when more chunks with the same request ID come in after the full message has already been received', () => {
+ const requestId = '123';
+ const firstChunk = MOCK_CHUNK_MESSAGE('first chunk', 1, requestId);
+ const secondChunk = MOCK_CHUNK_MESSAGE('second chunk', 2, requestId);
+ const successResponse = GENERATE_MOCK_TANUKI_RES('', requestId);
- // another chunk with the same request ID
- mockSubscriptionStream.next(secondChunk);
- await waitForPromises();
- // checking that addDuoChatMessage was not called again since request id was canceled
- expect(actionSpies.addDuoChatMessage).toHaveBeenCalledTimes(1);
- });
+ // message chunk streaming in
+ findSubscriptions().vm.$emit('message-stream', firstChunk);
+ expect(actionSpies.addDuoChatMessage).toHaveBeenCalledWith(expect.anything(), firstChunk);
- it('stops adding new message when requestId was canceled', async () => {
- const requestId = '123';
-
- duoChatGlobalState.isShown = true;
- createComponent({
- initialState: {
- messages: [
- {
- requestId,
- },
- ],
- },
+ // full message being sent
+ findSubscriptions().vm.$emit('message', successResponse);
+ expect(actionSpies.addDuoChatMessage).toHaveBeenCalledWith(
+ expect.anything(),
+ successResponse,
+ );
+ // another chunk with the same request ID
+ findSubscriptions().vm.$emit('message-stream', secondChunk);
+ // addDuoChatMessage should not be called since the full message was already sent
+ expect(actionSpies.addDuoChatMessage).toHaveBeenCalledTimes(2);
});
- await waitForPromises();
- findGlDuoChat().vm.$emit('chat-cancel');
+ it('continues to invoke addDuoChatMessage when a new message chunk arrives with a distinct request ID, even after a complete message has been received', () => {
+ const firstRequestId = '123';
+ const secondRequestId = '124';
+ const firstChunk = MOCK_CHUNK_MESSAGE('first chunk', 1, firstRequestId);
+ const secondChunk = MOCK_CHUNK_MESSAGE('second chunk', 2, firstRequestId);
+ const successResponse = GENERATE_MOCK_TANUKI_RES('', secondRequestId);
- // full message being sent
- mockSubscriptionComplete.next(GENERATE_MOCK_TANUKI_RES('', requestId));
- await waitForPromises();
- // checking that addDuoChatMessage was not called since request id was canceled
- expect(actionSpies.addDuoChatMessage).toHaveBeenCalledTimes(0);
- });
+ // message chunk streaming in
+ findSubscriptions().vm.$emit('message-stream', firstChunk);
+ expect(actionSpies.addDuoChatMessage).toHaveBeenCalledWith(expect.anything(), firstChunk);
- it('tracks performance metrics correctly when a chunk is received', async () => {
- const chunkMessage = MOCK_CHUNK_MESSAGE('chunk content', 1, 'requestId-123');
-
- duoChatGlobalState.isShown = true;
- createComponent();
- perfTrackingSpy = mockTracking(undefined, wrapper.element, jest.spyOn);
+ // full message being sent
+ findSubscriptions().vm.$emit('message', successResponse);
+ expect(actionSpies.addDuoChatMessage).toHaveBeenCalledWith(
+ expect.anything(),
+ successResponse,
+ );
+ // another chunk with a new request ID
+ findSubscriptions().vm.$emit('message-stream', secondChunk);
+ // addDuoChatMessage should be called since the second chunk has a new requestId
+ expect(actionSpies.addDuoChatMessage).toHaveBeenCalledWith(
+ expect.anything(),
+ successResponse,
+ );
+ });
- await waitForPromises();
+ it('clears the commands when streaming is done', () => {
+ const requestId = '123';
+ const firstChunk = MOCK_CHUNK_MESSAGE('first chunk', 1, requestId);
+ const successResponse = GENERATE_MOCK_TANUKI_RES('', requestId);
- mockSubscriptionStream.next(chunkMessage);
+ expect(duoChatGlobalState.commands).toHaveLength(0);
+ sendDuoChatCommand({ question: '/troubleshoot', resourceId: '1' });
+ expect(duoChatGlobalState.commands).toHaveLength(1);
- expect(performance.mark).toHaveBeenCalledWith('response-received');
- expect(performance.measure).toHaveBeenCalledWith(
- 'prompt-to-response',
- 'prompt-sent',
- 'response-received',
- );
- expect(performance.getEntriesByName).toHaveBeenCalledWith('prompt-to-response');
- expect(performance.clearMarks).toHaveBeenCalled();
- expect(performance.clearMeasures).toHaveBeenCalled();
+ createComponent();
- expect(perfTrackingSpy).toHaveBeenCalledWith(undefined, 'ai_response_time', {
- property: chunkMessage.data.aiCompletionResponse.requestId,
- value: 123,
+ // message chunk streaming in
+ findSubscriptions().vm.$emit('message-stream', firstChunk);
+ // No changes to commands
+ expect(duoChatGlobalState.commands).toHaveLength(1);
+ // full message being sent
+ findSubscriptions().vm.$emit('message', successResponse);
+ // commands have been cleared out
+ expect(duoChatGlobalState.commands).toHaveLength(0);
});
});
});
diff --git a/ee/spec/frontend/ai/tanuki_bot/components/tanuki_bot_subscriptions_spec.js b/ee/spec/frontend/ai/tanuki_bot/components/tanuki_bot_subscriptions_spec.js
new file mode 100644
index 0000000000000000000000000000000000000000..ef85e79c767e4f713fed7e5f9335909c1fd66164
--- /dev/null
+++ b/ee/spec/frontend/ai/tanuki_bot/components/tanuki_bot_subscriptions_spec.js
@@ -0,0 +1,145 @@
+import Vue from 'vue';
+import VueApollo from 'vue-apollo';
+import { createMockSubscription } from 'mock-apollo-client';
+import AiResponseSubscription from 'ee/ai/tanuki_bot/components/tanuki_bot_subscriptions.vue';
+import aiResponseSubscription from 'ee/graphql_shared/subscriptions/ai_completion_response.subscription.graphql';
+import aiResponseStreamSubscription from 'ee/graphql_shared/subscriptions/ai_completion_response_stream.subscription.graphql';
+import { shallowMountExtended } from 'helpers/vue_test_utils_helper';
+import createMockApollo from 'helpers/mock_apollo_helper';
+import waitForPromises from 'helpers/wait_for_promises';
+import { duoChatGlobalState } from '~/super_sidebar/constants';
+
+import {
+ MOCK_USER_ID,
+ GENERATE_MOCK_TANUKI_RES,
+ MOCK_CHUNK_MESSAGE,
+ MOCK_CLIENT_SUBSCRIPTION_ID,
+} from '../mock_data';
+
+Vue.use(VueApollo);
+
+describe('Ai Response Subscriptions', () => {
+ let wrapper;
+
+ let mockSubscriptionComplete;
+ let mockSubscriptionStream;
+ let aiResponseSubscriptionHandler;
+ let aiResponseStreamSubscriptionHandler;
+
+ const createComponent = ({ propsData = {} } = {}) => {
+ const apolloProvider = createMockApollo();
+
+ apolloProvider.defaultClient.setRequestHandler(
+ aiResponseSubscription,
+ aiResponseSubscriptionHandler,
+ );
+
+ apolloProvider.defaultClient.setRequestHandler(
+ aiResponseStreamSubscription,
+ aiResponseStreamSubscriptionHandler,
+ );
+
+ wrapper = shallowMountExtended(AiResponseSubscription, {
+ apolloProvider,
+ propsData: {
+ userId: MOCK_USER_ID,
+ clientSubscriptionId: MOCK_CLIENT_SUBSCRIPTION_ID,
+ ...propsData,
+ },
+ });
+ };
+
+ beforeEach(() => {
+ mockSubscriptionComplete = createMockSubscription();
+ mockSubscriptionStream = createMockSubscription();
+ aiResponseSubscriptionHandler = jest.fn(() => mockSubscriptionComplete);
+ aiResponseStreamSubscriptionHandler = jest.fn(() => mockSubscriptionStream);
+ });
+
+ afterEach(() => {
+ jest.clearAllMocks();
+ duoChatGlobalState.commands = [];
+
+ if (wrapper) {
+ wrapper.destroy();
+ }
+ });
+
+ describe('Subscriptions', () => {
+ it('passes the correct variables to the subscription queries', async () => {
+ createComponent();
+ await waitForPromises();
+
+ expect(aiResponseSubscriptionHandler).toHaveBeenCalledWith(
+ expect.objectContaining({
+ userId: MOCK_USER_ID,
+ aiAction: 'CHAT',
+ }),
+ );
+
+ expect(aiResponseStreamSubscriptionHandler).toHaveBeenCalledWith(
+ expect.objectContaining({
+ userId: MOCK_USER_ID,
+ clientSubscriptionId: MOCK_CLIENT_SUBSCRIPTION_ID,
+ }),
+ );
+ });
+
+ describe('aiCompletionResponseStream', () => {
+ it('emits message stream event', async () => {
+ const requestId = '123';
+ const firstChunk = {
+ data: { aiCompletionResponse: MOCK_CHUNK_MESSAGE('first chunk', 1, requestId) },
+ };
+
+ createComponent();
+ await waitForPromises();
+
+ // message chunk streaming in
+ mockSubscriptionStream.next(firstChunk);
+ await waitForPromises();
+
+ const emittedEvents = wrapper.emitted('message-stream');
+ expect(emittedEvents).toHaveLength(1);
+ expect(emittedEvents[0]).toEqual([firstChunk.data.aiCompletionResponse]);
+ });
+
+ it('emits response-received event', async () => {
+ const requestId = '123';
+ const firstChunk = {
+ data: { aiCompletionResponse: MOCK_CHUNK_MESSAGE('first chunk', 1, requestId) },
+ };
+
+ createComponent();
+ await waitForPromises();
+
+ // message chunk streaming in
+ mockSubscriptionStream.next(firstChunk);
+ await waitForPromises();
+
+ const emittedEvents = wrapper.emitted('response-received');
+ expect(emittedEvents).toHaveLength(1);
+ expect(emittedEvents[0]).toEqual([requestId]);
+ });
+ });
+
+ describe('aiCompletionResponse', () => {
+ it('emits message event', async () => {
+ const requestId = '123';
+ const successResponse = {
+ data: { aiCompletionResponse: GENERATE_MOCK_TANUKI_RES('', requestId) },
+ };
+ createComponent();
+ await waitForPromises();
+
+ // message chunk streaming in
+ mockSubscriptionComplete.next(successResponse);
+ await waitForPromises();
+
+ const emittedEvents = wrapper.emitted('message');
+ expect(emittedEvents).toHaveLength(1);
+ expect(emittedEvents[0]).toEqual([successResponse.data.aiCompletionResponse]);
+ });
+ });
+ });
+});
diff --git a/ee/spec/frontend/ai/tanuki_bot/mock_data.js b/ee/spec/frontend/ai/tanuki_bot/mock_data.js
index 82a01b89074a8b0ae05993ce5451927d9f765476..80bdb6d71a5fc06f746d7b644d447d671942ee59 100644
--- a/ee/spec/frontend/ai/tanuki_bot/mock_data.js
+++ b/ee/spec/frontend/ai/tanuki_bot/mock_data.js
@@ -70,20 +70,16 @@ export const MOCK_FAILING_USER_MESSAGE = {
export const MOCK_CHUNK_MESSAGE = (content = '', chunkId = 0, requestId = 1) => {
return {
- data: {
- aiCompletionResponse: {
- id: '611363bc-c75a-44e2-80cd-f22ab5e665be',
- requestId,
- content,
- errors: [],
- role: 'ASSISTANT',
- timestamp: '2024-05-29T17:17:06Z',
- type: null,
- chunkId,
- extras: {
- sources: null,
- },
- },
+ id: '611363bc-c75a-44e2-80cd-f22ab5e665be',
+ requestId,
+ content,
+ errors: [],
+ role: 'ASSISTANT',
+ timestamp: '2024-05-29T17:17:06Z',
+ type: null,
+ chunkId,
+ extras: {
+ sources: null,
},
};
};
@@ -93,20 +89,16 @@ export const GENERATE_MOCK_TANUKI_RES = (
requestId = '987',
) => {
return {
- data: {
- aiCompletionResponse: {
- id: '123',
- content: body,
- contentHtml: `<p>${body}</p>`,
- errors: [],
- requestId,
- role: MOCK_TANUKI_MESSAGE.role,
- timestamp: '2021-04-21T12:00:00.000Z',
- type: null,
- chunkId: null,
- extras: null,
- },
- },
+ id: '123',
+ content: body,
+ contentHtml: `<p>${body}</p>`,
+ errors: [],
+ requestId,
+ role: MOCK_TANUKI_MESSAGE.role,
+ timestamp: '2021-04-21T12:00:00.000Z',
+ type: null,
+ chunkId: null,
+ extras: null,
};
};
@@ -138,4 +130,5 @@ export const MOCK_TANUKI_BOT_MUTATATION_RES = {
};
export const MOCK_USER_ID = 'gid://gitlab/User/1';
+export const MOCK_CLIENT_SUBSCRIPTION_ID = '123';
export const MOCK_RESOURCE_ID = 'gid://gitlab/Issue/1';