azure-ai-agents LRO poller util fix for custom end states

sdk/ai/azure-ai-agents/CHANGELOG.md

@@ -28,8 +28,12 @@
 
 ### 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.
+
 ### 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

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_8815e6a59a"
 }

sdk/ai/azure-ai-agents/src/main/java/com/azure/ai/agents/implementation/AgentsServicePollUtils.java

@@ -0,0 +1,81 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+package com.azure.ai.agents.implementation;
+
+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 = "MemoryStores=V1Preview";
+
+    private AgentsServicePollUtils() {
+    }
+
+    // ---- header injection -------------------------------------------------
+
+    /**
+     * Returns a copy of the given {@link PollingStrategyOptions} whose {@link Context} includes
+     * the {@code Foundry-Features} header.  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).
+     */
+    static PollingStrategyOptions withFoundryFeatures(PollingStrategyOptions options) {
+        HttpHeaders headers = new HttpHeaders();
+        headers.set(FOUNDRY_FEATURES, FOUNDRY_FEATURES_VALUE);
+        Context context = options.getContext() != null ? options.getContext() : Context.NONE;
+        return options.setContext(context.addData(AddHeadersFromContextPolicy.AZURE_REQUEST_HTTP_HEADERS_KEY, headers));
+    }
+
+    // ---- status remapping -------------------------------------------------
+
+    /**
+     * Remaps a {@link PollResponse} whose status may contain a custom service terminal state
+     * ({@code "completed"}, {@code "superseded"}) that the base {@code OperationResourcePollingStrategy}
+     * cannot recognise.  If no remapping is needed the original response is returned as-is.
+     *
+     * <p>The Memory Stores TypeSpec defines:</p>
+     * <ul>
+     *   <li>{@code "completed"} &rarr; {@link LongRunningOperationStatus#SUCCESSFULLY_COMPLETED}</li>
+     *   <li>{@code "superseded"} &rarr; {@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 ("completed".equalsIgnoreCase(name)) {
+            return LongRunningOperationStatus.SUCCESSFULLY_COMPLETED;
+        } else if ("superseded".equalsIgnoreCase(name)) {
+            return LongRunningOperationStatus.USER_CANCELLED;
+        }
+        return status;
+    }
+}

sdk/ai/azure-ai-agents/src/main/java/com/azure/ai/agents/implementation/OperationLocationPollingStrategy.java

@@ -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
@@ -107,6 +108,19 @@ public Mono<PollResponse<T>> onInitialResponse(Response<?> response, PollingCont
         }
     }
 
+    /**
+     * {@inheritDoc}
+     *
+     * <p>Delegates to the parent (which handles URL construction with api-version and sends the
+     * {@code Foundry-Features} header injected via the context) and then remaps custom LRO
+     * terminal states (e.g. "completed", "superseded") to standard
+     * {@link LongRunningOperationStatus} values.</p>
+     */
+    @Override
+    public Mono<PollResponse<T>> poll(PollingContext<T> pollingContext, TypeReference<T> pollResponseType) {
+        return super.poll(pollingContext, pollResponseType).map(AgentsServicePollUtils::remapStatus);
+    }
+
     /**
      * {@inheritDoc}
      */

sdk/ai/azure-ai-agents/src/main/java/com/azure/ai/agents/implementation/SyncOperationLocationPollingStrategy.java

@@ -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
@@ -104,6 +105,19 @@ public PollResponse<T> onInitialResponse(Response<?> response, PollingContext<T>
                 response.getValue())));
     }
 
+    /**
+     * {@inheritDoc}
+     *
+     * <p>Delegates to the parent (which handles URL construction with api-version and sends the
+     * {@code Foundry-Features} header injected via the context) and then remaps custom LRO
+     * terminal states (e.g. "completed", "superseded") to standard
+     * {@link LongRunningOperationStatus} values.</p>
+     */
+    @Override
+    public PollResponse<T> poll(PollingContext<T> pollingContext, TypeReference<T> pollResponseType) {
+        return AgentsServicePollUtils.remapStatus(super.poll(pollingContext, pollResponseType));
+    }
+
     /**
      * {@inheritDoc}
      */

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);
     }
 }

sdk/ai/azure-ai-agents/src/test/java/com/azure/ai/agents/MemoryStoresTests.java

@@ -8,13 +8,15 @@
 import com.azure.core.exception.ResourceNotFoundException;
 import com.azure.core.http.HttpClient;
 import com.azure.core.util.polling.LongRunningOperationStatus;
+import com.azure.core.util.polling.PollResponse;
 import com.azure.core.util.polling.SyncPoller;
 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 java.time.Duration;
 import java.util.Arrays;
 
 import static com.azure.ai.agents.TestUtils.DISPLAY_NAME_WITH_ARGUMENTS;
@@ -23,9 +25,11 @@
 import static org.junit.jupiter.api.Assertions.assertThrows;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 
-@Disabled("Awaiting service versioning consolidation.")
+@Timeout(30)
 public class MemoryStoresTests extends ClientTestBase {
 
+    private static final Duration POLL_TIMEOUT = Duration.ofSeconds(30);
+
     @ParameterizedTest(name = DISPLAY_NAME_WITH_ARGUMENTS)
     @MethodSource("com.azure.ai.agents.TestUtils#getTestParameters")
     public void basicMemoryStoresCrud(HttpClient httpClient, AgentsServiceVersion serviceVersion) {
@@ -126,13 +130,10 @@ public void basicMemoryStores(HttpClient httpClient, AgentsServiceVersion servic
         SyncPoller<MemoryStoreUpdateResponse, MemoryStoreUpdateCompletedResult> updatePoller
             = memoryStoreClient.beginUpdateMemories(memoryStoreName, scope, Arrays.asList(userMessage), null, 0);
 
-        // Wait for the update operation to complete
-        LongRunningOperationStatus status = null;
-        while (status != LongRunningOperationStatus.fromString(MemoryStoreUpdateStatus.COMPLETED.toString(), true)) {
-            sleep(500);
-            System.out.println("Polling status: " + status);
-            status = updatePoller.poll().getStatus();
-        }
+        // Wait for the update operation to complete (with timeout to avoid hanging)
+        PollResponse<MemoryStoreUpdateResponse> pollResponse = updatePoller.waitForCompletion(POLL_TIMEOUT);
+        assertTrue(pollResponse.getStatus().isComplete(),
+            "Polling did not complete within timeout. Last status: " + pollResponse.getStatus());
         MemoryStoreUpdateCompletedResult updateResult = updatePoller.getFinalResult();
         assertNotNull(updateResult);
         assertNotNull(updateResult.getMemoryOperations());
@@ -234,12 +235,9 @@ public void advancedMemoryStores(HttpClient httpClient, AgentsServiceVersion ser
         System.out.println("Superseded first memory update operation (Update ID: " + initialUpdateId + ", Status: "
             + initialPoller.poll().getStatus() + ")");
 
-        LongRunningOperationStatus chainedStatus = null;
-        while (chainedStatus
-            != LongRunningOperationStatus.fromString(MemoryStoreUpdateStatus.COMPLETED.toString(), true)) {
-            sleep(500);
-            chainedStatus = chainedPoller.poll().getStatus();
-        }
+        PollResponse<MemoryStoreUpdateResponse> chainedPollResponse = chainedPoller.waitForCompletion(POLL_TIMEOUT);
+        assertTrue(chainedPollResponse.getStatus().isComplete(),
+            "Chained polling did not complete within timeout. Last status: " + chainedPollResponse.getStatus());
         MemoryStoreUpdateCompletedResult updateResult = chainedPoller.getFinalResult();
         assertNotNull(updateResult);
         assertNotNull(updateResult.getMemoryOperations());
@@ -302,11 +300,9 @@ public void advancedMemoryStores(HttpClient httpClient, AgentsServiceVersion ser
         System.out.println("Deleted memory store `" + memoryStoreName + "`");
     }
 
-    private static void cleanupBeforeTest(MemoryStoresClient memoryStoreClient, String memoryStoreName) {
-        // Ensure clean state: delete if it already exists
+    private void cleanupBeforeTest(MemoryStoresClient memoryStoreClient, String memoryStoreName) {
         try {
-            DeleteMemoryStoreResult deleteExisting = memoryStoreClient.deleteMemoryStore(memoryStoreName);
-            assertNotNull(deleteExisting);
+            memoryStoreClient.deleteMemoryStore(memoryStoreName);
         } catch (ResourceNotFoundException ex) {
             // ok if it does not exist
         }