fix(components): bound follow-up prompt generation latency

packages/components/src/followUpPrompts.test.ts

@@ -0,0 +1,100 @@
+import { FollowUpPromptConfig, FollowUpPromptProvider, ICommonObject } from './Interface'
+import { generateFollowUpPrompts } from './followUpPrompts'
+import { getCredentialData } from './utils'
+import { ChatOpenAI } from '@langchain/openai'
+
+jest.mock('./utils', () => ({
+    getCredentialData: jest.fn()
+}))
+
+const invokeMock = jest.fn()
+const withStructuredOutputMock = jest.fn(() => ({
+    invoke: invokeMock
+}))
+
+jest.mock('@langchain/openai', () => ({
+    ChatOpenAI: jest.fn().mockImplementation(() => ({
+        withStructuredOutput: withStructuredOutputMock
+    })),
+    AzureChatOpenAI: jest.fn()
+}))
+
+const mockedGetCredentialData = getCredentialData as jest.MockedFunction<typeof getCredentialData>
+const mockedChatOpenAI = ChatOpenAI as jest.MockedClass<typeof ChatOpenAI>
+
+const createFollowUpPromptConfig = (status = true): FollowUpPromptConfig => {
+    const providerConfig = {
+        credentialId: 'credential-id',
+        modelName: 'gpt-4.1-mini',
+        baseUrl: '',
+        prompt: 'Suggest follow-up questions for {history}',
+        temperature: '0.2'
+    }
+
+    return {
+        status,
+        selectedProvider: FollowUpPromptProvider.OPENAI,
+        [FollowUpPromptProvider.ANTHROPIC]: providerConfig,
+        [FollowUpPromptProvider.AZURE_OPENAI]: providerConfig,
+        [FollowUpPromptProvider.GOOGLE_GENAI]: providerConfig,
+        [FollowUpPromptProvider.MISTRALAI]: providerConfig,
+        [FollowUpPromptProvider.OPENAI]: providerConfig,
+        [FollowUpPromptProvider.GROQ]: providerConfig,
+        [FollowUpPromptProvider.OLLAMA]: providerConfig
+    }
+}
+
+describe('generateFollowUpPrompts', () => {
+    beforeEach(() => {
+        jest.useRealTimers()
+        jest.clearAllMocks()
+        mockedGetCredentialData.mockResolvedValue({ openAIApiKey: 'openai-key' })
+    })
+
+    it('returns follow-up questions when the provider succeeds', async () => {
+        const questions = ['What is next?', 'Can you expand?', 'Any constraints?']
+        invokeMock.mockResolvedValue({ questions })
+
+        const result = await generateFollowUpPrompts(createFollowUpPromptConfig(), 'history', {} as ICommonObject)
+
+        expect(result).toEqual({ questions })
+        expect(mockedChatOpenAI).toHaveBeenCalledWith({
+            apiKey: 'openai-key',
+            model: 'gpt-4.1-mini',
+            temperature: 0.2,
+            useResponsesApi: true
+        })
+        expect(invokeMock).toHaveBeenCalledWith('Suggest follow-up questions for history')
+    })
+
+    it('retries once when the first provider call fails', async () => {
+        const questions = ['Question one?', 'Question two?', 'Question three?']
+        invokeMock.mockRejectedValueOnce(new Error('temporary provider failure')).mockResolvedValueOnce({ questions })
+
+        const result = await generateFollowUpPrompts(createFollowUpPromptConfig(), 'history', {} as ICommonObject)
+
+        expect(result).toEqual({ questions })
+        expect(invokeMock).toHaveBeenCalledTimes(2)
+    })
+
+    it('returns undefined when provider calls time out', async () => {
+        jest.useFakeTimers()
+        invokeMock.mockImplementation(() => new Promise(() => undefined))
+
+        const resultPromise = generateFollowUpPrompts(createFollowUpPromptConfig(), 'history', {} as ICommonObject)
+        await Promise.resolve()
+
+        await jest.advanceTimersByTimeAsync(20000)
+
+        await expect(resultPromise).resolves.toBeUndefined()
+        expect(invokeMock).toHaveBeenCalledTimes(2)
+    })
+
+    it('does not read credentials or call a provider when follow-up prompts are disabled', async () => {
+        const result = await generateFollowUpPrompts(createFollowUpPromptConfig(false), 'history', {} as ICommonObject)
+
+        expect(result).toBeUndefined()
+        expect(mockedGetCredentialData).not.toHaveBeenCalled()
+        expect(mockedChatOpenAI).not.toHaveBeenCalled()
+    })
+})

packages/components/src/followUpPrompts.ts

@@ -20,6 +20,60 @@ export interface FollowUpPromptResult {
     questions: string[]
 }
 
+interface FollowUpPromptGenerationOptions {
+    timeoutMs?: number
+    maxRetries?: number
+}
+
+type FollowUpPromptTask<T> = () => Promise<T>
+
+const DEFAULT_FOLLOW_UP_PROMPT_TIMEOUT_MS = 10000
+const DEFAULT_FOLLOW_UP_PROMPT_MAX_RETRIES = 1
+
+class FollowUpPromptTimeoutError extends Error {
+    constructor(timeoutMs: number) {
+        super(`Follow-up prompt generation timed out after ${timeoutMs}ms`)
+        this.name = 'FollowUpPromptTimeoutError'
+    }
+}
+
+const createTimeoutPromise = <T>(timeoutMs: number): { promise: Promise<T>; clear: () => void } => {
+    let timeout: NodeJS.Timeout
+    const promise = new Promise<T>((_, reject) => {
+        timeout = setTimeout(() => reject(new FollowUpPromptTimeoutError(timeoutMs)), timeoutMs)
+    })
+
+    return {
+        promise,
+        clear: () => clearTimeout(timeout)
+    }
+}
+
+const invokeWithTimeoutAndRetry = async <T>(
+    task: FollowUpPromptTask<T>,
+    options: FollowUpPromptGenerationOptions = {}
+): Promise<T | undefined> => {
+    const timeoutMs = options.timeoutMs ?? DEFAULT_FOLLOW_UP_PROMPT_TIMEOUT_MS
+    const maxRetries = options.maxRetries ?? DEFAULT_FOLLOW_UP_PROMPT_MAX_RETRIES
+
+    for (let attempt = 0; attempt <= maxRetries; attempt += 1) {
+        const timeout = createTimeoutPromise<T>(timeoutMs)
+
+        try {
+            const taskPromise = task()
+            taskPromise.catch(() => undefined)
+
+            return await Promise.race([taskPromise, timeout.promise])
+        } catch {
+            if (attempt === maxRetries) return undefined
+        } finally {
+            timeout.clear()
+        }
+    }
+
+    return undefined
+}
+
 export const generateFollowUpPrompts = async (
     followUpPromptsConfig: FollowUpPromptConfig,
     apiMessageContent: string,
@@ -44,8 +98,10 @@ export const generateFollowUpPrompts = async (
                 const structuredLLM = llm.withStructuredOutput(FollowUpPromptType, {
                     method: 'functionCalling'
                 })
-                const structuredResponse = await structuredLLM.invoke(followUpPromptsPrompt)
-                return structuredResponse as FollowUpPromptResult
+                return await invokeWithTimeoutAndRetry(async () => {
+                    const structuredResponse = await structuredLLM.invoke(followUpPromptsPrompt)
+                    return structuredResponse as FollowUpPromptResult
+                })
             }
             case FollowUpPromptProvider.AZURE_OPENAI: {
                 const azureOpenAIApiKey = credentialData['azureOpenAIApiKey']
@@ -70,11 +126,13 @@ export const generateFollowUpPrompts = async (
                     {format_instructions}
                 `)
                 const chain = prompt.pipe(llm).pipe(parser)
-                const structuredResponse = await chain.invoke({
-                    history: apiMessageContent,
-                    format_instructions: formatInstructions
+                return await invokeWithTimeoutAndRetry(async () => {
+                    const structuredResponse = await chain.invoke({
+                        history: apiMessageContent,
+                        format_instructions: formatInstructions
+                    })
+                    return structuredResponse as FollowUpPromptResult
                 })
-                return structuredResponse as FollowUpPromptResult
             }
             case FollowUpPromptProvider.GOOGLE_GENAI: {
                 const model = new ChatGoogleGenerativeAI({
@@ -85,8 +143,10 @@ export const generateFollowUpPrompts = async (
                 const structuredLLM = model.withStructuredOutput(FollowUpPromptType, {
                     method: 'functionCalling'
                 })
-                const structuredResponse = await structuredLLM.invoke(followUpPromptsPrompt)
-                return structuredResponse as FollowUpPromptResult
+                return await invokeWithTimeoutAndRetry(async () => {
+                    const structuredResponse = await structuredLLM.invoke(followUpPromptsPrompt)
+                    return structuredResponse as FollowUpPromptResult
+                })
             }
             case FollowUpPromptProvider.MISTRALAI: {
                 const model = new ChatMistralAI({
@@ -98,8 +158,10 @@ export const generateFollowUpPrompts = async (
                 const structuredLLM = model.withStructuredOutput(FollowUpPromptType, {
                     method: 'functionCalling'
                 })
-                const structuredResponse = await structuredLLM.invoke(followUpPromptsPrompt)
-                return structuredResponse as FollowUpPromptResult
+                return await invokeWithTimeoutAndRetry(async () => {
+                    const structuredResponse = await structuredLLM.invoke(followUpPromptsPrompt)
+                    return structuredResponse as FollowUpPromptResult
+                })
             }
             case FollowUpPromptProvider.OPENAI: {
                 const model = new ChatOpenAI({
@@ -112,8 +174,10 @@ export const generateFollowUpPrompts = async (
                 const structuredLLM = model.withStructuredOutput(FollowUpPromptType, {
                     method: 'functionCalling'
                 })
-                const structuredResponse = await structuredLLM.invoke(followUpPromptsPrompt)
-                return structuredResponse as FollowUpPromptResult
+                return await invokeWithTimeoutAndRetry(async () => {
+                    const structuredResponse = await structuredLLM.invoke(followUpPromptsPrompt)
+                    return structuredResponse as FollowUpPromptResult
+                })
             }
             case FollowUpPromptProvider.GROQ: {
                 const llm = new ChatGroq({
@@ -124,44 +188,48 @@ export const generateFollowUpPrompts = async (
                 const structuredLLM = llm.withStructuredOutput(FollowUpPromptType, {
                     method: 'functionCalling'
                 })
-                const structuredResponse = await structuredLLM.invoke(followUpPromptsPrompt)
-                return structuredResponse as FollowUpPromptResult
+                return await invokeWithTimeoutAndRetry(async () => {
+                    const structuredResponse = await structuredLLM.invoke(followUpPromptsPrompt)
+                    return structuredResponse as FollowUpPromptResult
+                })
             }
             case FollowUpPromptProvider.OLLAMA: {
                 const ollamaClient = new Ollama({
                     host: providerConfig.baseUrl || 'http://127.0.0.1:11434'
                 })
 
-                const response = await ollamaClient.chat({
-                    model: providerConfig.modelName,
-                    messages: [
-                        {
-                            role: 'user',
-                            content: followUpPromptsPrompt
-                        }
-                    ],
-                    format: {
-                        type: 'object',
-                        properties: {
-                            questions: {
-                                type: 'array',
-                                items: {
-                                    type: 'string'
-                                },
-                                minItems: 3,
-                                maxItems: 3,
-                                description: 'Three follow-up questions based on the conversation history'
+                return await invokeWithTimeoutAndRetry(async () => {
+                    const response = await ollamaClient.chat({
+                        model: providerConfig.modelName,
+                        messages: [
+                            {
+                                role: 'user',
+                                content: followUpPromptsPrompt
                             }
+                        ],
+                        format: {
+                            type: 'object',
+                            properties: {
+                                questions: {
+                                    type: 'array',
+                                    items: {
+                                        type: 'string'
+                                    },
+                                    minItems: 3,
+                                    maxItems: 3,
+                                    description: 'Three follow-up questions based on the conversation history'
+                                }
+                            },
+                            required: ['questions'],
+                            additionalProperties: false
                         },
-                        required: ['questions'],
-                        additionalProperties: false
-                    },
-                    options: {
-                        temperature: parseFloat(`${providerConfig.temperature}`)
-                    }
-                })
-                const result = FollowUpPromptType.parse(JSON.parse(response.message.content))
-                return result
+                        options: {
+                            temperature: parseFloat(`${providerConfig.temperature}`)
+                        }
+                    })
+                    const result = FollowUpPromptType.parse(JSON.parse(response.message.content))
+                    return result
+                })
             }
         }
     } else {