azure-ai-agents LRO poller util fix for custom end states (#48157)

jpalvarezl · web-flow · commit be92b1279674 · 2026-03-04T06:26:18.000-08:00
* Poller fix

* Polled operation fix changelog entries

* PR copilot feedback

* New recordings

* Using TSP defined custom status for PollUtils

* Removed comments

* Using right value for header from enum

* Added comment with warning

* Applying poller customization through AgentsCustomization.java
diff --git a/sdk/ai/azure-ai-agents/CHANGELOG.md b/sdk/ai/azure-ai-agents/CHANGELOG.md
@@ -32,10 +32,13 @@
 
 ### Bugs Fixed
 
+- Fixed Memory Stores long-running operations (e.g. `beginUpdateMemories`) failing because the required `Foundry-Features` header was not included in poll requests, and custom LRO terminal states (`"completed"`, `"superseded"`) were not mapped to standard `LongRunningOperationStatus` values, causing pollers to hang indefinitely.
 - Fixed request parameter name from `"agent"` to `"agent_reference"` in `ResponsesClient` and `ResponsesAsyncClient` methods `createWithAgent` and `createWithAgentConversation`
 
 ### Other Changes
 
+- Enabled and stabilised `MemoryStoresTests` and `MemoryStoresAsyncTests` (previously `@Disabled`), with timeout guards to prevent hanging.
+
 ## 2.0.0-beta.1 (2026-02-25)
 
 ### Features Added
diff --git a/sdk/ai/azure-ai-agents/assets.json b/sdk/ai/azure-ai-agents/assets.json
@@ -2,5 +2,5 @@
   "AssetsRepo": "Azure/azure-sdk-assets",
   "AssetsRepoPrefixPath": "java",
   "TagPrefix": "java/ai/azure-ai-agents",
-  "Tag": "java/ai/azure-ai-agents_e4777fbd74"
+  "Tag": "java/ai/azure-ai-agents_34d0d1c5d4"
 }
diff --git a/sdk/ai/azure-ai-agents/customizations/src/main/java/AgentsCustomizations.java b/sdk/ai/azure-ai-agents/customizations/src/main/java/AgentsCustomizations.java
@@ -1,5 +1,7 @@
 import com.azure.autorest.customization.Customization;
 import com.azure.autorest.customization.LibraryCustomization;
+import com.github.javaparser.StaticJavaParser;
+import com.github.javaparser.ast.body.MethodDeclaration;
 import org.slf4j.Logger;
 
 
@@ -12,6 +14,7 @@ public class AgentsCustomizations extends Customization {
     @Override
     public void customize(LibraryCustomization libraryCustomization, Logger logger) {
         renameImageGenToolSize(libraryCustomization, logger);
+        modifyPollingStrategies(libraryCustomization, logger);
     }
 
     private void renameImageGenToolSize(LibraryCustomization customization, Logger logger) {
@@ -30,4 +33,24 @@ private void renameImageGenToolSize(LibraryCustomization customization, Logger l
                 .filter(entry -> "ONE_FIVE_THREE_SIXX_ONE_ZERO_TWO_FOUR".equals(entry.getName().getIdentifier()))
                 .forEach(entry -> entry.setName("RESOLUTION_1536_X_1024"))));
     }
+
+    private void modifyPollingStrategies(LibraryCustomization customization, Logger logger) {
+        customization.getClass("com.azure.ai.agents.implementation", "OperationLocationPollingStrategy")
+            .customizeAst(ast -> ast.getClassByName("OperationLocationPollingStrategy")
+                .ifPresent(clazz -> {
+                    clazz.getConstructors().get(1).getBody().getStatements()
+                        .set(0, StaticJavaParser.parseStatement("super(PollingUtils.OPERATION_LOCATION_HEADER, AgentsServicePollUtils.withFoundryFeatures(pollingStrategyOptions));"));
+
+                    clazz.addMember(StaticJavaParser.parseMethodDeclaration("@Override public Mono<PollResponse<T>> poll(PollingContext<T> pollingContext, TypeReference<T> pollResponseType) { return super.poll(pollingContext, pollResponseType).map(AgentsServicePollUtils::remapStatus); }"));
+                }));
+
+        customization.getClass("com.azure.ai.agents.implementation", "SyncOperationLocationPollingStrategy")
+            .customizeAst(ast -> ast.getClassByName("SyncOperationLocationPollingStrategy")
+                .ifPresent(clazz -> {
+                    clazz.getConstructors().get(1).getBody().getStatements()
+                        .set(0, StaticJavaParser.parseStatement("super(PollingUtils.OPERATION_LOCATION_HEADER, AgentsServicePollUtils.withFoundryFeatures(pollingStrategyOptions));"));
+
+                    clazz.addMember(StaticJavaParser.parseMethodDeclaration("@Override public PollResponse<T> poll(PollingContext<T> pollingContext, TypeReference<T> pollResponseType) { return AgentsServicePollUtils.remapStatus(super.poll(pollingContext, pollResponseType)); }"));
+                }));
+    }
 }
diff --git a/sdk/ai/azure-ai-agents/src/main/java/com/azure/ai/agents/implementation/AgentsServicePollUtils.java b/sdk/ai/azure-ai-agents/src/main/java/com/azure/ai/agents/implementation/AgentsServicePollUtils.java
@@ -0,0 +1,85 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+package com.azure.ai.agents.implementation;
+
+import com.azure.ai.agents.models.FoundryFeaturesOptInKeys;
+import com.azure.ai.agents.models.MemoryStoreUpdateStatus;
+import com.azure.core.http.HttpHeaderName;
+import com.azure.core.http.HttpHeaders;
+import com.azure.core.http.policy.AddHeadersFromContextPolicy;
+import com.azure.core.util.Context;
+import com.azure.core.util.polling.LongRunningOperationStatus;
+import com.azure.core.util.polling.PollResponse;
+import com.azure.core.util.polling.PollingStrategyOptions;
+
+/**
+ * Shared polling helpers for the Agents SDK.
+ *
+ * <p>The generated {@code OperationLocationPollingStrategy} / {@code SyncOperationLocationPollingStrategy}
+ * delegate here so that the two strategies stay in sync and only minimal edits are needed in the
+ * generated files.</p>
+ *
+ * <p>This class is package-private; it is <b>not</b> part of the public API.</p>
+ */
+final class AgentsServicePollUtils {
+
+    /** Required preview-feature header for Memory Stores operations. */
+    private static final HttpHeaderName FOUNDRY_FEATURES = HttpHeaderName.fromString("Foundry-Features");
+    private static final String FOUNDRY_FEATURES_VALUE = FoundryFeaturesOptInKeys.MEMORY_STORES_V1_PREVIEW.toString();
+
+    private AgentsServicePollUtils() {
+    }
+
+    /**
+     * Adds the {@code Foundry-Features} header to the given {@link PollingStrategyOptions}'s
+     * {@link Context}.  If the context already carries {@link HttpHeaders} under the
+     * {@link AddHeadersFromContextPolicy} key they are preserved; the {@code Foundry-Features}
+     * entry is merged in.  Because the pipeline already contains
+     * {@link AddHeadersFromContextPolicy}, the header is automatically added to every HTTP
+     * request the parent strategy makes (initial, poll, and final-result GETs).
+     *
+     * <p><strong>Note:</strong> this method mutates and returns the same
+     * {@code PollingStrategyOptions} instance.</p>
+     */
+    static PollingStrategyOptions withFoundryFeatures(PollingStrategyOptions options) {
+        Context context = options.getContext() != null ? options.getContext() : Context.NONE;
+        Object existing = context.getData(AddHeadersFromContextPolicy.AZURE_REQUEST_HTTP_HEADERS_KEY).orElse(null);
+        HttpHeaders headers
+            = (existing instanceof HttpHeaders) ? new HttpHeaders((HttpHeaders) existing) : new HttpHeaders();
+        headers.set(FOUNDRY_FEATURES, FOUNDRY_FEATURES_VALUE);
+        return options.setContext(context.addData(AddHeadersFromContextPolicy.AZURE_REQUEST_HTTP_HEADERS_KEY, headers));
+    }
+
+    /**
+     * Remaps a {@link PollResponse} whose status may contain a custom service terminal state
+     * ({@code "completed"}, {@code "superseded"}) that the base {@code OperationResourcePollingStrategy}
+     * cannot recognize.  If no remapping is needed the original response is returned as-is.
+     *
+     * <p>The Memory Stores Azure core defines:</p>
+     * <ul>
+     *   <li>{@code "completed"} {@link LongRunningOperationStatus#SUCCESSFULLY_COMPLETED}</li>
+     *   <li>{@code "superseded"}  {@link LongRunningOperationStatus#USER_CANCELLED}</li>
+     * </ul>
+     */
+    static <T> PollResponse<T> remapStatus(PollResponse<T> response) {
+        LongRunningOperationStatus status = response.getStatus();
+        LongRunningOperationStatus mapped = mapCustomStatus(status);
+        if (mapped == status) {
+            return response;
+        }
+        return new PollResponse<>(mapped, response.getValue(), response.getRetryAfter());
+    }
+
+    private static LongRunningOperationStatus mapCustomStatus(LongRunningOperationStatus status) {
+        // Standard statuses (Succeeded, Failed, Canceled, InProgress, NotStarted) are already
+        // mapped correctly by the parent's PollResult; only remap the custom ones.
+        String name = status.toString();
+        if (MemoryStoreUpdateStatus.COMPLETED.toString().equalsIgnoreCase(name)) {
+            return LongRunningOperationStatus.SUCCESSFULLY_COMPLETED;
+        } else if (MemoryStoreUpdateStatus.SUPERSEDED.toString().equalsIgnoreCase(name)) {
+            return LongRunningOperationStatus.USER_CANCELLED;
+        }
+        return status;
+    }
+}
diff --git a/sdk/ai/azure-ai-agents/src/main/java/com/azure/ai/agents/implementation/OperationLocationPollingStrategy.java b/sdk/ai/azure-ai-agents/src/main/java/com/azure/ai/agents/implementation/OperationLocationPollingStrategy.java
@@ -1,7 +1,6 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT License.
 // Code generated by Microsoft (R) TypeSpec Code Generator.
-
 package com.azure.ai.agents.implementation;
 
 import com.azure.core.exception.AzureException;
@@ -22,7 +21,6 @@
 import reactor.core.publisher.Mono;
 
 // DO NOT modify this helper class
-
 /**
  * Implements an operation location polling strategy, from Operation-Location.
  *
@@ -35,7 +33,9 @@ public final class OperationLocationPollingStrategy<T, U> extends OperationResou
     private static final ClientLogger LOGGER = new ClientLogger(OperationLocationPollingStrategy.class);
 
     private final ObjectSerializer serializer;
+
     private final String endpoint;
+
     private final String propertyName;
 
     /**
@@ -56,7 +56,8 @@ public OperationLocationPollingStrategy(PollingStrategyOptions pollingStrategyOp
      * @throws NullPointerException if {@code pollingStrategyOptions} is null.
      */
     public OperationLocationPollingStrategy(PollingStrategyOptions pollingStrategyOptions, String propertyName) {
-        super(PollingUtils.OPERATION_LOCATION_HEADER, pollingStrategyOptions);
+        super(PollingUtils.OPERATION_LOCATION_HEADER,
+            AgentsServicePollUtils.withFoundryFeatures(pollingStrategyOptions));
         this.propertyName = propertyName;
         this.endpoint = pollingStrategyOptions.getEndpoint();
         this.serializer = pollingStrategyOptions.getSerializer() != null
@@ -71,7 +72,6 @@ public OperationLocationPollingStrategy(PollingStrategyOptions pollingStrategyOp
     public Mono<PollResponse<T>> onInitialResponse(Response<?> response, PollingContext<T> pollingContext,
         TypeReference<T> pollResponseType) {
         // Response<?> is Response<BinaryData>
-
         HttpHeader operationLocationHeader = response.getHeaders().get(PollingUtils.OPERATION_LOCATION_HEADER);
         if (operationLocationHeader != null) {
             pollingContext.setData(PollingUtils.OPERATION_LOCATION_HEADER.getCaseSensitiveName(),
@@ -80,7 +80,6 @@ public Mono<PollResponse<T>> onInitialResponse(Response<?> response, PollingCont
         final String httpMethod = response.getRequest().getHttpMethod().name();
         pollingContext.setData(PollingUtils.HTTP_METHOD, httpMethod);
         pollingContext.setData(PollingUtils.REQUEST_URL, response.getRequest().getUrl().toString());
-
         if (response.getStatusCode() == 200
             || response.getStatusCode() == 201
             || response.getStatusCode() == 202
@@ -137,4 +136,9 @@ public Mono<U> getResult(PollingContext<T> pollingContext, TypeReference<U> resu
             return super.getResult(pollingContext, resultType);
         }
     }
+
+    @Override
+    public Mono<PollResponse<T>> poll(PollingContext<T> pollingContext, TypeReference<T> pollResponseType) {
+        return super.poll(pollingContext, pollResponseType).map(AgentsServicePollUtils::remapStatus);
+    }
 }
diff --git a/sdk/ai/azure-ai-agents/src/main/java/com/azure/ai/agents/implementation/SyncOperationLocationPollingStrategy.java b/sdk/ai/azure-ai-agents/src/main/java/com/azure/ai/agents/implementation/SyncOperationLocationPollingStrategy.java
@@ -1,7 +1,6 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT License.
 // Code generated by Microsoft (R) TypeSpec Code Generator.
-
 package com.azure.ai.agents.implementation;
 
 import com.azure.core.exception.AzureException;
@@ -23,7 +22,6 @@
 import java.util.Map;
 
 // DO NOT modify this helper class
-
 /**
  * Implements a synchronous operation location polling strategy, from Operation-Location.
  *
@@ -36,7 +34,9 @@ public final class SyncOperationLocationPollingStrategy<T, U> extends SyncOperat
     private static final ClientLogger LOGGER = new ClientLogger(SyncOperationLocationPollingStrategy.class);
 
     private final ObjectSerializer serializer;
+
     private final String endpoint;
+
     private final String propertyName;
 
     /**
@@ -57,7 +57,8 @@ public SyncOperationLocationPollingStrategy(PollingStrategyOptions pollingStrate
      * @throws NullPointerException if {@code pollingStrategyOptions} is null.
      */
     public SyncOperationLocationPollingStrategy(PollingStrategyOptions pollingStrategyOptions, String propertyName) {
-        super(PollingUtils.OPERATION_LOCATION_HEADER, pollingStrategyOptions);
+        super(PollingUtils.OPERATION_LOCATION_HEADER,
+            AgentsServicePollUtils.withFoundryFeatures(pollingStrategyOptions));
         this.propertyName = propertyName;
         this.endpoint = pollingStrategyOptions.getEndpoint();
         this.serializer = pollingStrategyOptions.getSerializer() != null
@@ -72,7 +73,6 @@ public SyncOperationLocationPollingStrategy(PollingStrategyOptions pollingStrate
     public PollResponse<T> onInitialResponse(Response<?> response, PollingContext<T> pollingContext,
         TypeReference<T> pollResponseType) {
         // Response<?> is Response<BinaryData>
-
         HttpHeader operationLocationHeader = response.getHeaders().get(PollingUtils.OPERATION_LOCATION_HEADER);
         if (operationLocationHeader != null) {
             pollingContext.setData(PollingUtils.OPERATION_LOCATION_HEADER.getCaseSensitiveName(),
@@ -81,7 +81,6 @@ public PollResponse<T> onInitialResponse(Response<?> response, PollingContext<T>
         final String httpMethod = response.getRequest().getHttpMethod().name();
         pollingContext.setData(PollingUtils.HTTP_METHOD, httpMethod);
         pollingContext.setData(PollingUtils.REQUEST_URL, response.getRequest().getUrl().toString());
-
         if (response.getStatusCode() == 200
             || response.getStatusCode() == 201
             || response.getStatusCode() == 202
@@ -97,7 +96,6 @@ public PollResponse<T> onInitialResponse(Response<?> response, PollingContext<T>
             }
             return new PollResponse<>(LongRunningOperationStatus.IN_PROGRESS, initialResponseType, retryAfter);
         }
-
         throw LOGGER.logExceptionAsError(new AzureException(
             String.format("Operation failed or cancelled with status code %d, '%s' header: %s, and response body: %s",
                 response.getStatusCode(), PollingUtils.OPERATION_LOCATION_HEADER, operationLocationHeader,
@@ -130,4 +128,9 @@ public U getResult(PollingContext<T> pollingContext, TypeReference<U> resultType
             return super.getResult(pollingContext, resultType);
         }
     }
+
+    @Override
+    public PollResponse<T> poll(PollingContext<T> pollingContext, TypeReference<T> pollResponseType) {
+        return AgentsServicePollUtils.remapStatus(super.poll(pollingContext, pollResponseType));
+    }
 }
diff --git a/sdk/ai/azure-ai-agents/src/test/java/com/azure/ai/agents/MemoryStoresAsyncTests.java b/sdk/ai/azure-ai-agents/src/test/java/com/azure/ai/agents/MemoryStoresAsyncTests.java
@@ -12,21 +12,20 @@
 import com.azure.ai.agents.models.MemoryStoreDetails;
 import com.azure.ai.agents.models.MemoryStoreUpdateCompletedResult;
 import com.azure.ai.agents.models.MemoryStoreUpdateResponse;
-import com.azure.ai.agents.models.MemoryStoreUpdateStatus;
 import com.azure.ai.agents.models.PageOrder;
 import com.azure.core.exception.ResourceNotFoundException;
 import com.azure.core.http.HttpClient;
 import com.azure.core.util.polling.AsyncPollResponse;
-import com.azure.core.util.polling.LongRunningOperationStatus;
 import com.azure.core.util.polling.PollerFlux;
 import com.openai.models.responses.EasyInputMessage;
 import com.openai.models.responses.ResponseInputItem;
-import org.junit.jupiter.api.Disabled;
+import org.junit.jupiter.api.Timeout;
 import org.junit.jupiter.params.ParameterizedTest;
 import org.junit.jupiter.params.provider.MethodSource;
 import reactor.core.publisher.Mono;
 import reactor.test.StepVerifier;
 
+import java.time.Duration;
 import java.util.Arrays;
 import java.util.Objects;
 
@@ -35,12 +34,9 @@
 import static org.junit.jupiter.api.Assertions.assertNotNull;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 
-@Disabled("Awaiting service versioning consolidation.")
+@Timeout(30)
 public class MemoryStoresAsyncTests extends ClientTestBase {
 
-    private static final LongRunningOperationStatus COMPLETED_OPERATION_STATUS
-        = LongRunningOperationStatus.fromString(MemoryStoreUpdateStatus.COMPLETED.toString(), true);
-
     @ParameterizedTest(name = DISPLAY_NAME_WITH_ARGUMENTS)
     @MethodSource("com.azure.ai.agents.TestUtils#getTestParameters")
     public void basicMemoryStoresCrud(HttpClient httpClient, AgentsServiceVersion serviceVersion) {
@@ -282,15 +278,9 @@ private static Mono<Void> cleanupBeforeTest(MemoryStoresAsyncClient memoryStoreC
     private static Mono<MemoryStoreUpdateCompletedResult>
         waitForUpdateCompletion(PollerFlux<MemoryStoreUpdateResponse, MemoryStoreUpdateCompletedResult> pollerFlux) {
         Objects.requireNonNull(pollerFlux, "pollerFlux cannot be null");
-        return pollerFlux.takeUntil(response -> COMPLETED_OPERATION_STATUS.equals(response.getStatus()))
+        return pollerFlux.takeUntil(response -> response.getStatus().isComplete())
+            .timeout(Duration.ofSeconds(30))
             .last()
-            .map(AsyncPollResponse::getValue)
-            .map(response -> {
-                MemoryStoreUpdateCompletedResult result = response == null ? null : response.getResult();
-                if (result == null) {
-                    throw new IllegalStateException("Memory store update did not complete successfully.");
-                }
-                return result;
-            });
+            .flatMap(AsyncPollResponse::getFinalResult);
     }
 }
diff --git a/sdk/ai/azure-ai-agents/src/test/java/com/azure/ai/agents/MemoryStoresTests.java b/sdk/ai/azure-ai-agents/src/test/java/com/azure/ai/agents/MemoryStoresTests.java
diff --git a/sdk/ai/azure-ai-agents/src/test/java/com/azure/ai/agents/implementation/AgentsServicePollUtilsTest.java b/sdk/ai/azure-ai-agents/src/test/java/com/azure/ai/agents/implementation/AgentsServicePollUtilsTest.java

Original file line number	Diff line number	Diff line change
`@@ -2,5 +2,5 @@`
`2`	`2`	`"AssetsRepo": "Azure/azure-sdk-assets",`
`3`	`3`	`"AssetsRepoPrefixPath": "java",`
`4`	`4`	`"TagPrefix": "java/ai/azure-ai-agents",`
`5`		`- "Tag": "java/ai/azure-ai-agents_e4777fbd74"`
	`5`	`+ "Tag": "java/ai/azure-ai-agents_34d0d1c5d4"`
`6`	`6`	`}`