From 7689800d3632dba4c460436874fee91b61995942 Mon Sep 17 00:00:00 2001
From: Jose Alvarez <josealvar@microsoft.com>
Date: Fri, 27 Feb 2026 14:33:19 +0100
Subject: [PATCH 01/10] Poller fix

---
 sdk/ai/azure-ai-agents/assets.json            |   2 +-
 .../AgentsServicePollUtils.java               |  81 +++++++++++++
 .../OperationLocationPollingStrategy.java     |  16 ++-
 .../SyncOperationLocationPollingStrategy.java |  16 ++-
 .../ai/agents/MemoryStoresAsyncTests.java     |  22 +---
 .../azure/ai/agents/MemoryStoresTests.java    |  34 +++---
 .../AgentsServicePollUtilsTest.java           | 108 ++++++++++++++++++
 7 files changed, 241 insertions(+), 38 deletions(-)
 create mode 100644 sdk/ai/azure-ai-agents/src/main/java/com/azure/ai/agents/implementation/AgentsServicePollUtils.java
 create mode 100644 sdk/ai/azure-ai-agents/src/test/java/com/azure/ai/agents/implementation/AgentsServicePollUtilsTest.java

diff --git a/sdk/ai/azure-ai-agents/assets.json b/sdk/ai/azure-ai-agents/assets.json
index a6061845a097..485dd874b5b0 100644
--- 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_8815e6a59a"
 }
\ No newline at end of file
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
new file mode 100644
index 000000000000..5d95286a0b78
--- /dev/null
+++ b/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;
+    }
+}
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
index df8351347821..9d055893d9f9 100644
--- 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
@@ -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}
      */
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
index 7d249b059e57..c8c2c46834f0 100644
--- 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
@@ -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}
      */
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
index 4b73a6f6eed4..80715460c5ca 100644
--- 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
index 6e067c5bae03..d681b9833913 100644
--- 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
@@ -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
         }
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
new file mode 100644
index 000000000000..3f31ebf97773
--- /dev/null
+++ b/sdk/ai/azure-ai-agents/src/test/java/com/azure/ai/agents/implementation/AgentsServicePollUtilsTest.java
@@ -0,0 +1,108 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+package com.azure.ai.agents.implementation;
+
+import com.azure.core.http.HttpHeaders;
+import com.azure.core.http.HttpPipelineBuilder;
+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;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.Arguments;
+import org.junit.jupiter.params.provider.MethodSource;
+
+import java.util.stream.Stream;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertNotNull;
+import static org.junit.jupiter.api.Assertions.assertSame;
+
+class AgentsServicePollUtilsTest {
+
+    // ---- remapStatus tests ------------------------------------------------
+
+    static Stream<Arguments> remapStatusCases() {
+        return Stream.of(
+            // Custom statuses that need remapping
+            Arguments.of("completed", LongRunningOperationStatus.SUCCESSFULLY_COMPLETED),
+            Arguments.of("Completed", LongRunningOperationStatus.SUCCESSFULLY_COMPLETED),
+            Arguments.of("COMPLETED", LongRunningOperationStatus.SUCCESSFULLY_COMPLETED),
+            Arguments.of("superseded", LongRunningOperationStatus.USER_CANCELLED),
+            Arguments.of("Superseded", LongRunningOperationStatus.USER_CANCELLED));
+    }
+
+    @ParameterizedTest
+    @MethodSource("remapStatusCases")
+    void remapStatusMapsCustomStatuses(String statusName, LongRunningOperationStatus expected) {
+        // The parent's PollResult.setStatus(String) calls fromString(name, false) for unknown statuses
+        LongRunningOperationStatus customStatus = LongRunningOperationStatus.fromString(statusName, false);
+        PollResponse<String> original = new PollResponse<>(customStatus, "value");
+
+        PollResponse<String> remapped = AgentsServicePollUtils.remapStatus(original);
+
+        assertEquals(expected, remapped.getStatus());
+        assertEquals("value", remapped.getValue());
+    }
+
+    static Stream<Arguments> standardStatusCases() {
+        return Stream.of(Arguments.of(LongRunningOperationStatus.SUCCESSFULLY_COMPLETED),
+            Arguments.of(LongRunningOperationStatus.FAILED), Arguments.of(LongRunningOperationStatus.USER_CANCELLED),
+            Arguments.of(LongRunningOperationStatus.IN_PROGRESS), Arguments.of(LongRunningOperationStatus.NOT_STARTED));
+    }
+
+    @ParameterizedTest
+    @MethodSource("standardStatusCases")
+    void remapStatusPassesThroughStandardStatuses(LongRunningOperationStatus status) {
+        PollResponse<String> original = new PollResponse<>(status, "value");
+
+        PollResponse<String> result = AgentsServicePollUtils.remapStatus(original);
+
+        assertSame(original, result, "Standard status should return the same PollResponse instance");
+    }
+
+    @Test
+    void remapStatusPreservesRetryAfter() {
+        LongRunningOperationStatus completed = LongRunningOperationStatus.fromString("completed", false);
+        java.time.Duration retryAfter = java.time.Duration.ofSeconds(5);
+        PollResponse<String> original = new PollResponse<>(completed, "value", retryAfter);
+
+        PollResponse<String> remapped = AgentsServicePollUtils.remapStatus(original);
+
+        assertEquals(LongRunningOperationStatus.SUCCESSFULLY_COMPLETED, remapped.getStatus());
+        assertEquals(retryAfter, remapped.getRetryAfter());
+    }
+
+    // ---- withFoundryFeatures tests ----------------------------------------
+
+    @Test
+    void withFoundryFeaturesAddsHeaderToContext() {
+        PollingStrategyOptions options = new PollingStrategyOptions(new HttpPipelineBuilder().build());
+
+        PollingStrategyOptions result = AgentsServicePollUtils.withFoundryFeatures(options);
+
+        Context context = result.getContext();
+        assertNotNull(context);
+        Object headerObj = context.getData(AddHeadersFromContextPolicy.AZURE_REQUEST_HTTP_HEADERS_KEY).orElse(null);
+        assertNotNull(headerObj, "Context should contain HTTP headers under the AddHeadersFromContextPolicy key");
+        HttpHeaders headers = (HttpHeaders) headerObj;
+        assertEquals("MemoryStores=V1Preview",
+            headers.getValue(com.azure.core.http.HttpHeaderName.fromString("Foundry-Features")));
+    }
+
+    @Test
+    void withFoundryFeaturesPreservesExistingContext() {
+        PollingStrategyOptions options
+            = new PollingStrategyOptions(new HttpPipelineBuilder().build()).setContext(new Context("myKey", "myValue"));
+
+        PollingStrategyOptions result = AgentsServicePollUtils.withFoundryFeatures(options);
+
+        Context context = result.getContext();
+        assertEquals("myValue", context.getData("myKey").orElse(null), "Existing context data should be preserved");
+        assertNotNull(context.getData(AddHeadersFromContextPolicy.AZURE_REQUEST_HTTP_HEADERS_KEY).orElse(null),
+            "Foundry-Features header should also be present");
+    }
+}

From 7a4df25869ae54573baceaf34a295741aacd8aad Mon Sep 17 00:00:00 2001
From: Jose Alvarez <josealvar@microsoft.com>
Date: Fri, 27 Feb 2026 14:36:20 +0100
Subject: [PATCH 02/10] Polled operation fix changelog entries

---
 sdk/ai/azure-ai-agents/CHANGELOG.md | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/sdk/ai/azure-ai-agents/CHANGELOG.md b/sdk/ai/azure-ai-agents/CHANGELOG.md
index 85896d1621b7..759b68fc2895 100644
--- a/sdk/ai/azure-ai-agents/CHANGELOG.md
+++ b/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

From 0216116f6b95d5479db28bbc4ac069def8023157 Mon Sep 17 00:00:00 2001
From: Jose Alvarez <josealvar@microsoft.com>
Date: Fri, 27 Feb 2026 14:57:23 +0100
Subject: [PATCH 03/10] PR copilot feedback

---
 .../AgentsServicePollUtils.java               | 15 ++++--
 .../azure/ai/agents/MemoryStoresTests.java    | 49 +++----------------
 .../AgentsServicePollUtilsTest.java           | 22 +++++++++
 3 files changed, 40 insertions(+), 46 deletions(-)

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
index 5d95286a0b78..1890f53165c5 100644
--- 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
@@ -33,15 +33,22 @@ 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
+     * 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) {
-        HttpHeaders headers = new HttpHeaders();
-        headers.set(FOUNDRY_FEATURES, FOUNDRY_FEATURES_VALUE);
         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));
     }
 
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
index d681b9833913..0258db1cd9e5 100644
--- 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
@@ -7,8 +7,6 @@
 import com.azure.ai.agents.models.DeleteMemoryStoreResult;
 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;
@@ -118,10 +116,6 @@ public void basicMemoryStores(HttpClient httpClient, AgentsServiceVersion servic
         assertNotNull(memoryStore.getId());
         assertEquals(memoryStoreName, memoryStore.getName());
         assertEquals(description, memoryStore.getDescription());
-        System.out.println("Created memory store: " + memoryStore.getName() + " (" + memoryStore.getId() + "): "
-            + memoryStore.getDescription());
-        System.out.println("  - Chat model: " + definition.getChatModel());
-        System.out.println("  - Embedding model: " + definition.getEmbeddingModel());
 
         // Add memories to the memory store
         ResponseInputItem userMessage = ResponseInputItem.ofEasyInputMessage(
@@ -131,19 +125,15 @@ public void basicMemoryStores(HttpClient httpClient, AgentsServiceVersion servic
             = memoryStoreClient.beginUpdateMemories(memoryStoreName, scope, Arrays.asList(userMessage), null, 0);
 
         // 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());
+        updatePoller.waitForCompletion(POLL_TIMEOUT);
         MemoryStoreUpdateCompletedResult updateResult = updatePoller.getFinalResult();
         assertNotNull(updateResult);
         assertNotNull(updateResult.getMemoryOperations());
-        System.out.println("Updated with " + updateResult.getMemoryOperations().size() + " memory operations");
+
         for (MemoryOperation operation : updateResult.getMemoryOperations()) {
             assertNotNull(operation.getKind());
             assertNotNull(operation.getMemoryItem().getMemoryId());
             assertNotNull(operation.getMemoryItem().getContent());
-            System.out.println("  - Operation: " + operation.getKind() + ", Memory ID: "
-                + operation.getMemoryItem().getMemoryId() + ", Content: " + operation.getMemoryItem().getContent());
         }
 
         ResponseInputItem queryMessage = ResponseInputItem.ofEasyInputMessage(
@@ -154,23 +144,19 @@ public void basicMemoryStores(HttpClient httpClient, AgentsServiceVersion servic
             Arrays.asList(queryMessage), null, searchOptions);
         assertNotNull(searchResponse);
         assertNotNull(searchResponse.getMemories());
-        System.out.println("Found " + searchResponse.getMemories().size() + " memories");
+
         for (MemorySearchItem memory : searchResponse.getMemories()) {
             assertNotNull(memory.getMemoryItem().getMemoryId());
             assertNotNull(memory.getMemoryItem().getContent());
-            System.out.println("  - Memory ID: " + memory.getMemoryItem().getMemoryId() + ", Content: "
-                + memory.getMemoryItem().getContent());
         }
 
         // Delete memories for a specific scope
         memoryStoreClient.deleteScope(memoryStoreName, scope);
-        System.out.println("Deleted memories for scope '" + scope + "'");
 
         // Delete memory store
         DeleteMemoryStoreResult deleteResponse = memoryStoreClient.deleteMemoryStore(memoryStoreName);
         assertNotNull(deleteResponse);
         assertTrue(deleteResponse.isDeleted());
-        System.out.println("Deleted memory store `" + memoryStoreName + "`");
     }
 
     @ParameterizedTest(name = DISPLAY_NAME_WITH_ARGUMENTS)
@@ -203,8 +189,6 @@ public void advancedMemoryStores(HttpClient httpClient, AgentsServiceVersion ser
             = memoryStoreClient.createMemoryStore(memoryStoreName, definition, description, null);
         assertNotNull(memoryStore);
         assertEquals(memoryStoreName, memoryStore.getName());
-        System.out.println("Created memory store: " + memoryStore.getName() + " (" + memoryStore.getId() + "): "
-            + memoryStore.getDescription());
 
         ResponseInputItem initialMessage = ResponseInputItem.ofEasyInputMessage(
             EasyInputMessage.builder().role(EasyInputMessage.Role.USER).content(firstMessageContent).build());
@@ -215,8 +199,6 @@ public void advancedMemoryStores(HttpClient httpClient, AgentsServiceVersion ser
         assertNotNull(initialResponse);
         String initialUpdateId = initialResponse.getUpdateId();
         assertNotNull(initialUpdateId);
-        System.out.println("Scheduled memory update operation (Update ID: " + initialUpdateId + ", Status: "
-            + initialPoller.poll().getStatus() + ")");
 
         // Extend the previous update with another update and more messages
         ResponseInputItem chainedMessage = ResponseInputItem.ofEasyInputMessage(
@@ -228,27 +210,16 @@ public void advancedMemoryStores(HttpClient httpClient, AgentsServiceVersion ser
         assertNotNull(chainedResponse);
         String chainedUpdateId = chainedResponse.getUpdateId();
         assertNotNull(chainedUpdateId);
-        System.out.println("Scheduled memory update operation (Update ID: " + chainedUpdateId + ", Status: "
-            + chainedPoller.poll().getStatus() + ")");
-
-        // As first update has not started yet, the new update will cancel the first update and cover both sets of messages
-        System.out.println("Superseded first memory update operation (Update ID: " + initialUpdateId + ", Status: "
-            + initialPoller.poll().getStatus() + ")");
 
-        PollResponse<MemoryStoreUpdateResponse> chainedPollResponse = chainedPoller.waitForCompletion(POLL_TIMEOUT);
-        assertTrue(chainedPollResponse.getStatus().isComplete(),
-            "Chained polling did not complete within timeout. Last status: " + chainedPollResponse.getStatus());
+        chainedPoller.waitForCompletion(POLL_TIMEOUT);
         MemoryStoreUpdateCompletedResult updateResult = chainedPoller.getFinalResult();
         assertNotNull(updateResult);
         assertNotNull(updateResult.getMemoryOperations());
-        System.out.println("Second update " + chainedUpdateId + " completed with "
-            + updateResult.getMemoryOperations().size() + " memory operations");
+
         for (MemoryOperation operation : updateResult.getMemoryOperations()) {
             assertNotNull(operation.getKind());
             assertNotNull(operation.getMemoryItem().getMemoryId());
             assertNotNull(operation.getMemoryItem().getContent());
-            System.out.println("  - Operation: " + operation.getKind() + ", Memory ID: "
-                + operation.getMemoryItem().getMemoryId() + ", Content: " + operation.getMemoryItem().getContent());
         }
 
         // Retrieve memories from the memory store
@@ -261,12 +232,10 @@ public void advancedMemoryStores(HttpClient httpClient, AgentsServiceVersion ser
             = memoryStoreClient.searchMemories(memoryStoreName, scope, Arrays.asList(searchQuery), null, searchOptions);
         assertNotNull(searchResponse);
         assertNotNull(searchResponse.getMemories());
-        System.out.println("Found " + searchResponse.getMemories().size() + " memories");
+
         for (MemorySearchItem memory : searchResponse.getMemories()) {
             assertNotNull(memory.getMemoryItem().getMemoryId());
             assertNotNull(memory.getMemoryItem().getContent());
-            System.out.println("  - Memory ID: " + memory.getMemoryItem().getMemoryId() + ", Content: "
-                + memory.getMemoryItem().getContent());
         }
         String previousSearchId = searchResponse.getSearchId();
         assertNotNull(previousSearchId);
@@ -281,23 +250,19 @@ public void advancedMemoryStores(HttpClient httpClient, AgentsServiceVersion ser
             Arrays.asList(agentMessage, followupQuery), previousSearchId, searchOptions);
         assertNotNull(followupSearch);
         assertNotNull(followupSearch.getMemories());
-        System.out.println("Found " + followupSearch.getMemories().size() + " memories");
+
         for (MemorySearchItem memory : followupSearch.getMemories()) {
             assertNotNull(memory.getMemoryItem().getMemoryId());
             assertNotNull(memory.getMemoryItem().getContent());
-            System.out.println("  - Memory ID: " + memory.getMemoryItem().getMemoryId() + ", Content: "
-                + memory.getMemoryItem().getContent());
         }
 
         // Delete memories for the current scope
         memoryStoreClient.deleteScope(memoryStoreName, scope);
-        System.out.println("Deleted memories for scope '" + scope + "'");
 
         // Delete memory store
         DeleteMemoryStoreResult deleteResponse = memoryStoreClient.deleteMemoryStore(memoryStoreName);
         assertNotNull(deleteResponse);
         assertTrue(deleteResponse.isDeleted());
-        System.out.println("Deleted memory store `" + memoryStoreName + "`");
     }
 
     private void cleanupBeforeTest(MemoryStoresClient memoryStoreClient, String memoryStoreName) {
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
index 3f31ebf97773..31b1410e4458 100644
--- 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
@@ -105,4 +105,26 @@ void withFoundryFeaturesPreservesExistingContext() {
         assertNotNull(context.getData(AddHeadersFromContextPolicy.AZURE_REQUEST_HTTP_HEADERS_KEY).orElse(null),
             "Foundry-Features header should also be present");
     }
+
+    @Test
+    void withFoundryFeaturesMergesWithExistingHeaders() {
+        HttpHeaders existingHeaders = new HttpHeaders();
+        existingHeaders.set(com.azure.core.http.HttpHeaderName.fromString("X-Custom"), "custom-value");
+        Context contextWithHeaders
+            = new Context(AddHeadersFromContextPolicy.AZURE_REQUEST_HTTP_HEADERS_KEY, existingHeaders);
+        PollingStrategyOptions options
+            = new PollingStrategyOptions(new HttpPipelineBuilder().build()).setContext(contextWithHeaders);
+
+        PollingStrategyOptions result = AgentsServicePollUtils.withFoundryFeatures(options);
+
+        HttpHeaders merged = (HttpHeaders) result.getContext()
+            .getData(AddHeadersFromContextPolicy.AZURE_REQUEST_HTTP_HEADERS_KEY)
+            .orElse(null);
+        assertNotNull(merged);
+        assertEquals("custom-value", merged.getValue(com.azure.core.http.HttpHeaderName.fromString("X-Custom")),
+            "Pre-existing header should be preserved");
+        assertEquals("MemoryStores=V1Preview",
+            merged.getValue(com.azure.core.http.HttpHeaderName.fromString("Foundry-Features")),
+            "Foundry-Features header should be added");
+    }
 }

From 998166b5637d0dac7a49b726c28e6e5eb5eaa043 Mon Sep 17 00:00:00 2001
From: Jose Alvarez <josealvar@microsoft.com>
Date: Fri, 27 Feb 2026 15:17:37 +0100
Subject: [PATCH 04/10] New recordings

---
 sdk/ai/azure-ai-agents/assets.json | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/sdk/ai/azure-ai-agents/assets.json b/sdk/ai/azure-ai-agents/assets.json
index 485dd874b5b0..08d8cb4c7c20 100644
--- 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_8815e6a59a"
+  "Tag": "java/ai/azure-ai-agents_34d0d1c5d4"
 }
\ No newline at end of file

From e7dabcb7e417ac030baf83ff0ea4edfd7f055080 Mon Sep 17 00:00:00 2001
From: Jose Alvarez <josealvar@microsoft.com>
Date: Tue, 3 Mar 2026 10:28:44 +0100
Subject: [PATCH 05/10] Using TSP defined custom status for PollUtils

---
 .../implementation/AgentsServicePollUtils.java  | 17 +++++++----------
 1 file changed, 7 insertions(+), 10 deletions(-)

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
index 1890f53165c5..81d426ada37f 100644
--- 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
@@ -3,6 +3,7 @@
 
 package com.azure.ai.agents.implementation;
 
+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;
@@ -30,8 +31,6 @@ final class AgentsServicePollUtils {
     private AgentsServicePollUtils() {
     }
 
-    // ---- header injection -------------------------------------------------
-
     /**
      * Adds the {@code Foundry-Features} header to the given {@link PollingStrategyOptions}'s
      * {@link Context}.  If the context already carries {@link HttpHeaders} under the
@@ -52,17 +51,15 @@ static PollingStrategyOptions withFoundryFeatures(PollingStrategyOptions options
         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.
+     * cannot recognize.  If no remapping is needed the original response is returned as-is.
      *
-     * <p>The Memory Stores TypeSpec defines:</p>
+     * <p>The Memory Stores Azure core defines:</p>
      * <ul>
-     *   <li>{@code "completed"} &rarr; {@link LongRunningOperationStatus#SUCCESSFULLY_COMPLETED}</li>
-     *   <li>{@code "superseded"} &rarr; {@link LongRunningOperationStatus#USER_CANCELLED}</li>
+     *   <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) {
@@ -78,9 +75,9 @@ private static LongRunningOperationStatus mapCustomStatus(LongRunningOperationSt
         // 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)) {
+        if (MemoryStoreUpdateStatus.COMPLETED.toString().equalsIgnoreCase(name)) {
             return LongRunningOperationStatus.SUCCESSFULLY_COMPLETED;
-        } else if ("superseded".equalsIgnoreCase(name)) {
+        } else if (MemoryStoreUpdateStatus.SUPERSEDED.toString().equalsIgnoreCase(name)) {
             return LongRunningOperationStatus.USER_CANCELLED;
         }
         return status;

From ad70deb9234f107d41e6c2386093d1c8d4bdaa72 Mon Sep 17 00:00:00 2001
From: Jose Alvarez <josealvar@microsoft.com>
Date: Tue, 3 Mar 2026 10:44:02 +0100
Subject: [PATCH 06/10] Removed comments

---
 .../ai/agents/implementation/AgentsServicePollUtilsTest.java  | 4 ----
 1 file changed, 4 deletions(-)

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
index 31b1410e4458..56d744c43499 100644
--- 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
@@ -23,8 +23,6 @@
 
 class AgentsServicePollUtilsTest {
 
-    // ---- remapStatus tests ------------------------------------------------
-
     static Stream<Arguments> remapStatusCases() {
         return Stream.of(
             // Custom statuses that need remapping
@@ -76,8 +74,6 @@ void remapStatusPreservesRetryAfter() {
         assertEquals(retryAfter, remapped.getRetryAfter());
     }
 
-    // ---- withFoundryFeatures tests ----------------------------------------
-
     @Test
     void withFoundryFeaturesAddsHeaderToContext() {
         PollingStrategyOptions options = new PollingStrategyOptions(new HttpPipelineBuilder().build());

From 780e2ffab3e464f8e8aa2880186d28402178bfeb Mon Sep 17 00:00:00 2001
From: Jose Alvarez <josealvar@microsoft.com>
Date: Tue, 3 Mar 2026 10:49:31 +0100
Subject: [PATCH 07/10] Using right value for header from enum

---
 .../ai/agents/implementation/AgentsServicePollUtils.java      | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

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
index 81d426ada37f..b6dc84725be1 100644
--- 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
@@ -3,6 +3,7 @@
 
 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;
@@ -25,8 +26,7 @@ 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 static final String FOUNDRY_FEATURES_VALUE = FoundryFeaturesOptInKeys.MEMORY_STORES_V1_PREVIEW.toString();
 
     private AgentsServicePollUtils() {
     }

From d84dc59b67e5550cbbaea9525524b66349f6e537 Mon Sep 17 00:00:00 2001
From: Jose Alvarez <josealvar@microsoft.com>
Date: Tue, 3 Mar 2026 11:49:41 +0100
Subject: [PATCH 08/10] Added comment with warning

---
 .../implementation/OperationLocationPollingStrategy.java       | 3 +++
 .../implementation/SyncOperationLocationPollingStrategy.java   | 3 +++
 2 files changed, 6 insertions(+)

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
index 9d055893d9f9..842c024ee67a 100644
--- 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
@@ -23,6 +23,9 @@
 
 // DO NOT modify this helper class
 
+// Re-apply after regeneration: constructor wraps options with AgentsServicePollUtils.withFoundryFeatures(),
+// poll() override delegates to super then remaps statuses via AgentsServicePollUtils.remapStatus().
+
 /**
  * Implements an operation location polling strategy, from Operation-Location.
  *
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
index c8c2c46834f0..30ccd721d576 100644
--- 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
@@ -24,6 +24,9 @@
 
 // DO NOT modify this helper class
 
+// Re-apply after regeneration: constructor wraps options with AgentsServicePollUtils.withFoundryFeatures(),
+// poll() override delegates to super then remaps statuses via AgentsServicePollUtils.remapStatus().
+
 /**
  * Implements a synchronous operation location polling strategy, from Operation-Location.
  *

From 4a77c1c8e9e8d08b8dd4624e9b980afab06779f1 Mon Sep 17 00:00:00 2001
From: Weidong Xu <weidxu@microsoft.com>
Date: Wed, 4 Mar 2026 01:58:44 +0800
Subject: [PATCH 09/10] custmoize the lro strategy

---
 .../src/main/java/AgentsCustomizations.java   | 23 +++++++++++++++++++
 1 file changed, 23 insertions(+)

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
index d72807757c79..592020d63a72 100644
--- 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)); }"));
+                }));
+    }
 }

From 055f5d0f868f08def1d8b740636c43ce1d1e11db Mon Sep 17 00:00:00 2001
From: Weidong Xu <weidxu@microsoft.com>
Date: Wed, 4 Mar 2026 01:58:49 +0800
Subject: [PATCH 10/10] regen

---
 .../OperationLocationPollingStrategy.java     | 27 +++++-------------
 .../SyncOperationLocationPollingStrategy.java | 28 +++++--------------
 2 files changed, 14 insertions(+), 41 deletions(-)

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
index 842c024ee67a..6dc846c0de9b 100644
--- 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,10 +21,6 @@
 import reactor.core.publisher.Mono;
 
 // DO NOT modify this helper class
-
-// Re-apply after regeneration: constructor wraps options with AgentsServicePollUtils.withFoundryFeatures(),
-// poll() override delegates to super then remaps statuses via AgentsServicePollUtils.remapStatus().
-
 /**
  * Implements an operation location polling strategy, from Operation-Location.
  *
@@ -38,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;
 
     /**
@@ -75,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(),
@@ -84,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
@@ -111,19 +106,6 @@ 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}
      */
@@ -154,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
index 30ccd721d576..5d8c3582180f 100644
--- 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,10 +22,6 @@
 import java.util.Map;
 
 // DO NOT modify this helper class
-
-// Re-apply after regeneration: constructor wraps options with AgentsServicePollUtils.withFoundryFeatures(),
-// poll() override delegates to super then remaps statuses via AgentsServicePollUtils.remapStatus().
-
 /**
  * Implements a synchronous operation location polling strategy, from Operation-Location.
  *
@@ -39,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;
 
     /**
@@ -76,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(),
@@ -85,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
@@ -101,26 +96,12 @@ 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,
                 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}
      */
@@ -147,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));
+    }
 }