Skip to content

Add clientSubscriptionId to Ai completion subscription

Nicolas Dular requested to merge nd/ai-subscription into master

What does this MR do and why?

Related #419531 (closed)

This adds an optional clientSubscriptionId to the ai_completion_response subscription and the AiActionInput. In addition to that it fixes the GraphqlTriggers to be able to deal with optional subscription arguments.

This prepares us to allow listening only to a specific clientSubscriptionId on the websocket, and to only broadcast messages based on a user_id in the future

This has no breaking changes, nor changes how the subscription gets used.

Changelog: changed EE: true

Screenshots or screen recordings

Screenshots are required for UI changes, and strongly recommended for all other merge requests.

Before After

How to set up and validate locally

diff --git a/ee/app/assets/javascripts/ai/graphql/chat.mutation.graphql b/ee/app/assets/javascripts/ai/graphql/chat.mutation.graphql
index c26d8270de37..6ae83926f55f 100644
--- a/ee/app/assets/javascripts/ai/graphql/chat.mutation.graphql
+++ b/ee/app/assets/javascripts/ai/graphql/chat.mutation.graphql
@@ -1,5 +1,10 @@
-mutation chat($question: String!, $resourceId: AiModelID!) {
-  aiAction(input: { chat: { resourceId: $resourceId, content: $question } }) {
+mutation chat($question: String!, $resourceId: AiModelID!, $clientSubscriptionId: String) {
+  aiAction(
+    input: {
+      chat: { resourceId: $resourceId, content: $question }
+      clientSubscriptionId: $clientSubscriptionId
+    }
+  ) {
     requestId
     errors
   }
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 abbcba8fa280..44b9aca8e862 100644
--- a/ee/app/assets/javascripts/ai/tanuki_bot/components/app.vue
+++ b/ee/app/assets/javascripts/ai/tanuki_bot/components/app.vue
@@ -60,6 +60,7 @@ export default {
           return {
             resourceId: this.resourceId || this.userId,
             userId: this.userId,
+            clientSubscriptionId: 'someGeneratedId',
           };
         },
         result({ data }) {
@@ -107,6 +108,7 @@ export default {
           variables: {
             question,
             resourceId: this.resourceId || this.userId,
+            clientSubscriptionId: 'someGeneratedId',
           },
         })
         .then(({ data: { aiAction = {} } = {} }) => {
diff --git a/ee/app/assets/javascripts/graphql_shared/subscriptions/ai_completion_response.subscription.graphql b/ee/app/assets/javascripts/graphql_shared/subscriptions/ai_completion_response.subscription.graphql
index c7878c1e57a1..fe1a8f155382 100644
--- a/ee/app/assets/javascripts/graphql_shared/subscriptions/ai_completion_response.subscription.graphql
+++ b/ee/app/assets/javascripts/graphql_shared/subscriptions/ai_completion_response.subscription.graphql
@@ -1,5 +1,13 @@
-subscription aiCompletionResponse($userId: UserID, $resourceId: AiModelID!) {
-  aiCompletionResponse(userId: $userId, resourceId: $resourceId) {
+subscription aiCompletionResponse(
+  $userId: UserID
+  $resourceId: AiModelID
+  $clientSubscriptionId: String
+) {
+  aiCompletionResponse(
+    userId: $userId
+    resourceId: $resourceId
+    clientSubscriptionId: $clientSubscriptionId
+  ) {
     requestId
     responseBody
     errors

MR acceptance checklist

This checklist encourages us to confirm any changes have been analyzed to reduce risks in quality, performance, reliability, security, and maintainability.

Edited by Nicolas Dular

Merge request reports