Use integration tasks for processing Ask AI queries in Slack (#1011)

* Use integration tasks for processing Ask AI queries

* use itty router v4

* Fix isAllowedToResponse

* comment

* use Promise.all

* review

* update integration tsconfig

* changeset

* config

* undo integration.json

* export verifyIntegrationSignature from runtime

* changeset

* comments

Author: taranvohra

.changeset/hip-keys-double.md

@@ -0,0 +1,5 @@
+---
+'@gitbook/integration-slack': patch
+---
+
+Fix AI queries timing out after 30s

.changeset/huge-bars-behave.md

@@ -0,0 +1,5 @@
+---
+'@gitbook/runtime': minor
+---
+
+Export `verifyIntegrationRequestSignature` utility method

bun.lock

@@ -124,7 +124,7 @@
     },
     "integrations/front": {
       "name": "@gitbook/integration-front",
-      "version": "0.0.1",
+      "version": "1.0.0",
       "dependencies": {
         "@gitbook/runtime": "*",
       },
@@ -519,7 +519,7 @@
     },
     "integrations/runllm-widget": {
       "name": "@gitbook/integration-runllm-widget",
-      "version": "0.3.0",
+      "version": "0.4.0",
       "dependencies": {
         "@gitbook/api": "*",
         "@gitbook/runtime": "*",
@@ -567,11 +567,11 @@
     },
     "integrations/slack": {
       "name": "@gitbook/integration-slack",
-      "version": "2.5.0",
+      "version": "2.5.2",
       "dependencies": {
         "@gitbook/api": "*",
         "@gitbook/runtime": "*",
-        "itty-router": "^2.6.1",
+        "itty-router": "^4.0.26",
         "js-sha256": "^0.9.0",
         "remove-markdown": "^0.5.0",
       },
@@ -701,7 +701,7 @@
     },
     "packages/api": {
       "name": "@gitbook/api",
-      "version": "0.143.0",
+      "version": "0.143.1",
       "dependencies": {
         "event-iterator": "^2.0.0",
         "eventsource-parser": "^3.0.0",
@@ -2640,7 +2640,7 @@
 
     "@gitbook/integration-runkit/itty-router": ["[email protected]", "", {}, "sha512-hIPHtXGymCX7Lzb2I4G6JgZFE4QEEQwst9GORK7sMYUpJvLfy4yZJr95r04e8DzoAnj6HcxM2m4TbK+juu+18g=="],
 
-    "@gitbook/integration-slack/itty-router": ["itty-router@2.6.6", "", {}, "sha512-hIPHtXGymCX7Lzb2I4G6JgZFE4QEEQwst9GORK7sMYUpJvLfy4yZJr95r04e8DzoAnj6HcxM2m4TbK+juu+18g=="],
+    "@gitbook/integration-slack/itty-router": ["itty-router@4.2.2", "", {}, "sha512-KegPW0l9SNPadProoFT07AB84uOqLUwzlXQ7HsqkS31WUrxkjdhcemRpTDUuetbMJ89uBtWeQSVoiEmUAu31uw=="],
 
     "@gitbook/integration-va-auth0/itty-router": ["[email protected]", "", {}, "sha512-KegPW0l9SNPadProoFT07AB84uOqLUwzlXQ7HsqkS31WUrxkjdhcemRpTDUuetbMJ89uBtWeQSVoiEmUAu31uw=="],
 

integrations/slack/package.json

@@ -5,7 +5,7 @@
     "dependencies": {
         "@gitbook/runtime": "*",
         "@gitbook/api": "*",
-        "itty-router": "^2.6.1",
+        "itty-router": "^4.0.26",
         "js-sha256": "^0.9.0",
         "remove-markdown": "^0.5.0"
     },

integrations/slack/src/actions/queryAskAI.ts

@@ -1,10 +1,11 @@
-import type {
+import {
     SearchAIAnswer,
     GitBookAPI,
     Revision,
     RevisionPage,
     RevisionPageGroup,
     SearchAIAnswerSource,
+    IntegrationInstallation,
 } from '@gitbook/api';
 
 import {
@@ -14,8 +15,14 @@ import {
 } from '../configuration';
 import { slackAPI } from '../slack';
 import { QueryDisplayBlock, ShareTools, decodeSlackEscapeChars, Spacer, SourcesBlock } from '../ui';
-import { getInstallationApiClient, stripBotName, stripMarkdown } from '../utils';
+import {
+    getInstallationApiClient,
+    getIntegrationInstallationForTeam,
+    stripBotName,
+    stripMarkdown,
+} from '../utils';
 import { Logger } from '@gitbook/runtime';
+import { IntegrationTaskAskAI } from '../types';
 
 const logger = Logger('slack:queryAskAI');
 
@@ -72,10 +79,8 @@ const capitalizeFirstLetter = (text: string) =>
 async function getRelatedSources(params: {
     sources?: SearchAIAnswer['sources'];
     client: GitBookAPI;
-    environment: SlackRuntimeEnvironment;
-    organization: string;
 }): Promise<RelatedSource[]> {
-    const { sources, client, organization } = params;
+    const { sources, client } = params;
 
     if (!sources || sources.length === 0) {
         return [];
@@ -149,59 +154,120 @@ async function getRelatedSources(params: {
 /*
  * Queries GitBook AskAI via the GitBook API and posts the answer in the form of Slack UI Blocks back to the original channel/conversation/thread.
  */
-export async function queryAskAI({
-    channelId,
-    teamId,
-    threadId,
-    userId,
-    text,
-    messageType,
-    context,
-    authorization,
-
-    responseUrl,
-    channelName,
-}: IQueryAskAI) {
-    const { environment, api } = context;
+export async function queryAskAI(params: IQueryAskAI) {
+    const {
+        channelId,
+        teamId,
+        threadId,
+        userId,
+        text,
+        messageType,
+        context,
+        authorization,
+        responseUrl,
+    } = params;
+    const { api } = context;
 
     const askText = `_Asking: ${stripMarkdown(text)}_`;
     logger.info(`${askText} (channelId: ${channelId}, teamId: ${teamId}, userId: ${userId})`);
 
-    const { client, installation } = await getInstallationApiClient(api, teamId);
+    const installation = await getIntegrationInstallationForTeam(context, teamId);
     if (!installation) {
         throw new Error('Installation not found');
     }
-    // Authenticate as the installation
+    // Use the slack access token
     const accessToken = (installation.configuration as SlackInstallationConfiguration)
         .oauth_credentials?.access_token;
 
     // strip a bot name if the user_id from the request is present in the query itself (specifically for a bot mention)
     // @ts-ignore
     const parsedQuery = stripMarkdown(stripBotName(text, authorization?.user_id));
 
-    // async acknowledge the request to the end user early
-    slackAPI(
-        context,
-        {
-            method: 'POST',
-            path: messageType === 'ephemeral' ? 'chat.postEphemeral' : 'chat.postMessage',
-            responseUrl,
-            payload: {
-                channel: channelId,
-                text: askText,
-                ...(userId ? { user: userId } : {}), // actually shouldn't be optional
-                ...(threadId ? { thread_ts: threadId } : {}),
+    await Promise.all([
+        // acknowledge the ask query back to the user
+        slackAPI(
+            context,
+            {
+                method: 'POST',
+                path: messageType === 'ephemeral' ? 'chat.postEphemeral' : 'chat.postMessage',
+                responseUrl,
+                payload: {
+                    channel: channelId,
+                    text: askText,
+                    ...(userId ? { user: userId } : {}), // actually shouldn't be optional
+                    ...(threadId ? { thread_ts: threadId } : {}),
+                },
+            },
+            {
+                accessToken,
             },
+        ),
+        // Queue a task to process the AskAI query asynchronously. Because workers have a 30s timeout on
+        // waitUntil which is not enough for scenarios where AskAI might take longer to respond.
+        queueQueryAskAI({
+            ...params,
+            accessToken,
+            installation,
+            query: parsedQuery,
+        }),
+    ]);
+}
+
+/**
+ * Queues an integration task to process the AskAI query asynchronously.
+ */
+async function queueQueryAskAI(
+    params: IQueryAskAI & {
+        query: string;
+        installation: IntegrationInstallation;
+        accessToken: string | undefined;
+    },
+) {
+    const { accessToken, installation, query, context, ...rest } = params;
+
+    const task: IntegrationTaskAskAI = {
+        type: 'ask:ai',
+        payload: {
+            query,
+            organizationId: installation.target.organization,
+            installationId: installation.id,
+            accessToken,
+            ...rest,
         },
-        {
+    };
+
+    logger.info(`Queue task ${task.type} for installation: ${task.payload.installationId})`);
+
+    await context.api.integrations.queueIntegrationTask(context.environment.integration.name, {
+        task,
+    });
+}
+
+/**
+ * Handle the integration task to process the AskAI query.
+ */
+export async function handleAskAITask(task: IntegrationTaskAskAI, context: SlackRuntimeContext) {
+    const {
+        payload: {
+            channelName,
+            channelId,
+            userId,
+            text,
+            messageType,
+            responseUrl,
+            threadId,
+            query,
+            organizationId,
+            installationId,
             accessToken,
         },
-    );
+    } = task;
 
+    const client = await getInstallationApiClient(context, installationId);
     const result = await client.orgs.askInOrganization(
-        installation.target.organization,
+        organizationId,
         {
-            query: parsedQuery,
+            query,
         },
         {
             format: 'markdown',
@@ -223,8 +289,6 @@ export async function queryAskAI({
         const relatedSources = await getRelatedSources({
             sources: answer.sources,
             client,
-            environment,
-            organization: installation.target.organization,
         });
 
         const header = text.length > 150 ? `${text.slice(0, 140)}...` : text;

integrations/slack/src/router.ts

@@ -1,6 +1,13 @@
 import { Router } from 'itty-router';
 
-import { createOAuthHandler, FetchEventCallback, OAuthResponse } from '@gitbook/runtime';
+import {
+    createOAuthHandler,
+    ExposableError,
+    FetchEventCallback,
+    Logger,
+    OAuthResponse,
+    verifyIntegrationRequestSignature,
+} from '@gitbook/runtime';
 
 import {
     createSlackEventsHandler,
@@ -13,6 +20,10 @@ import {
 import { unfurlLink } from './links';
 import { verifySlackRequest, acknowledgeSlackRequest } from './middlewares';
 import { getChannelsPaginated } from './slack';
+import { IntegrationTask } from './types';
+import { handleAskAITask } from './actions';
+
+const logger = Logger('slack');
 
 /**
  * Handle incoming HTTP requests:
@@ -45,6 +56,39 @@ export const handleFetchEvent: FetchEventCallback = async (request, context) =>
         ].join(' '),
     );
 
+    /**
+     * Handle integration tasks
+     */
+    router.post('/tasks', async (request) => {
+        const verified = await verifyIntegrationRequestSignature(request, environment);
+
+        if (!verified) {
+            const message = `Invalid signature for integration task`;
+            logger.error(message);
+            throw new ExposableError(message);
+        }
+
+        const { task } = JSON.parse(await request.text()) as { task: IntegrationTask };
+        logger.debug('verified & received integration task', task.type);
+
+        switch (task.type) {
+            case 'ask:ai': {
+                await handleAskAITask(task, context);
+                break;
+            }
+            default: {
+                const error = `Unknown integration task type: ${task.type}`;
+                logger.error(error);
+                throw new Error(error);
+            }
+        }
+
+        return new Response(JSON.stringify({ processed: true }), {
+            status: 200,
+            headers: { 'content-type': 'application/json' },
+        });
+    });
+
     /*
      * Authenticate the user using OAuth.
      */

integrations/slack/src/types.ts

@@ -1,3 +1,5 @@
+import type { IQueryAskAI } from './actions';
+
 // Slack Block Kit Types
 export type SlackTextType = 'plain_text' | 'mrkdwn';
 
@@ -21,3 +23,22 @@ export type SlackBlock =
     | { type: 'context'; elements: SlackTextField[] }
     | { type: 'divider' }
     | { type: 'actions'; elements: SlackButtonElement[] };
+
+export type IntegrationTaskType = 'ask:ai';
+
+export type BaseIntegrationTask<Type extends IntegrationTaskType, Payload extends object> = {
+    type: Type;
+    payload: Payload;
+};
+
+export type IntegrationTaskAskAI = BaseIntegrationTask<
+    'ask:ai',
+    {
+        query: string;
+        organizationId: string;
+        installationId: string;
+        accessToken: string | undefined;
+    } & Omit<IQueryAskAI, 'context'>
+>;
+
+export type IntegrationTask = IntegrationTaskAskAI;