diff --git a/sdk/ai/ai-projects/src/api/memoryStores/memoryStoreUpdatePoller.ts b/sdk/ai/ai-projects/src/api/memoryStores/memoryStoreUpdatePoller.ts index d46f230ca6c8..b989c44e63f5 100644 --- a/sdk/ai/ai-projects/src/api/memoryStores/memoryStoreUpdatePoller.ts +++ b/sdk/ai/ai-projects/src/api/memoryStores/memoryStoreUpdatePoller.ts @@ -135,6 +135,7 @@ function buildRunningOperation( client: Client, expectedStatuses: string[], getInitialResponse?: () => PromiseLike, + options?: CreateMemoryStoreUpdatePollerOptions, ): RunningOperation { const pollAbortController = new AbortController(); return { @@ -146,13 +147,19 @@ function buildRunningOperation( return toOperationResponse(initialResponse, expectedStatuses); }, sendPollRequest: async (path: string, pollOptions?: { abortSignal?: AbortSignalLike }) => { + // The poll request will both listen to the user provided abort signal and the poller's own abort signal function abortListener(): void { pollAbortController.abort(); } const abortSignal = pollAbortController.signal; - if (pollOptions?.abortSignal?.aborted) { + if (options?.abortSignal?.aborted) { + pollAbortController.abort(); + } else if (pollOptions?.abortSignal?.aborted) { pollAbortController.abort(); } else if (!abortSignal.aborted) { + options?.abortSignal?.addEventListener("abort", abortListener, { + once: true, + }); pollOptions?.abortSignal?.addEventListener("abort", abortListener, { once: true, }); @@ -162,6 +169,7 @@ function buildRunningOperation( try { response = (await client.pathUnchecked(path).get({ abortSignal })) as PathUncheckedResponse; } finally { + options?.abortSignal?.removeEventListener("abort", abortListener); pollOptions?.abortSignal?.removeEventListener("abort", abortListener); } @@ -191,6 +199,7 @@ export function createMemoryStoreUpdatePoller( return initialResponse; } : undefined, + options, ), { intervalInMs: options?.updateIntervalInMs,