diff --git a/packages/core/src/core/baseLlmClient.test.ts b/packages/core/src/core/baseLlmClient.test.ts index 82196ffa84a..ec9b0d70d5d 100644 --- a/packages/core/src/core/baseLlmClient.test.ts +++ b/packages/core/src/core/baseLlmClient.test.ts @@ -575,6 +575,28 @@ describe('BaseLlmClient', () => { ); }); + it('should forward the caller-provided abortSignal to retryWithBackoff so the overall retry loop honors it', async () => { + const mockResponse = createMockResponse('This is the content.'); + mockGenerateContent.mockResolvedValue(mockResponse); + + const options = { + modelConfigKey: { model: 'test-model' }, + contents: [{ role: 'user', parts: [{ text: 'Give me content.' }] }], + abortSignal: abortController.signal, + promptId: 'content-prompt-id', + role: LlmRole.UTILITY_TOOL, + }; + + await client.generateContent(options); + + expect(retryWithBackoff).toHaveBeenCalledWith( + expect.any(Function), + expect.objectContaining({ + signal: options.abortSignal, + }), + ); + }); + it('should validate content using shouldRetryOnContent function', async () => { const mockResponse = createMockResponse('Some valid content.'); mockGenerateContent.mockResolvedValue(mockResponse); diff --git a/packages/core/src/core/baseLlmClient.ts b/packages/core/src/core/baseLlmClient.ts index a23b8c353c7..edaa54b31e5 100644 --- a/packages/core/src/core/baseLlmClient.ts +++ b/packages/core/src/core/baseLlmClient.ts @@ -357,6 +357,7 @@ export class BaseLlmClient { shouldRetryOnContent, maxAttempts: availabilityMaxAttempts ?? maxAttempts ?? DEFAULT_MAX_ATTEMPTS, + signal: abortSignal, getAvailabilityContext, onPersistent429: this.config.isInteractive() ? (authType, error) =>