Skip to content

Add WhenUpdatesCompleted to Foundry MemoryProvider to await async memory extraction #653

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
PratikDhanave (PratikDhanave) wants to merge 5 commits into
base: main
Choose a base branch
from
+120 −1
Open

Add WhenUpdatesCompleted to Foundry MemoryProvider to await async memory extraction #653

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
442ca39
Add WhenUpdatesCompleted to Foundry MemoryProvider
PratikDhanave Jul 23, 2026
ad48af7
Merge remote-tracking branch 'upstream/main' into pr653
PratikDhanave Aug 4, 2026
a98b2d6
Merge remote-tracking branch 'upstream/main' into m653
PratikDhanave Aug 6, 2026
2cb4693
Merge remote-tracking branch 'upstream/main' into x653
PratikDhanave Aug 20, 2026
2a3c877
Merge remote-tracking branch 'upstream/main' into x653
PratikDhanave Aug 22, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
40 changes: 39 additions & 1 deletion provider/foundryprovider/memory.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,8 +9,11 @@ import (
"log/slog"
"net/http"
"strings"
"sync"
"time"

"github.com/Azure/azure-sdk-for-go/sdk/azcore"
"github.com/Azure/azure-sdk-for-go/sdk/azcore/runtime"
"github.com/microsoft/agent-framework-go/agent"
"github.com/microsoft/agent-framework-go/internal/azaiprojects"
"github.com/microsoft/agent-framework-go/message"
Expand Down Expand Up @@ -55,6 +58,9 @@ type MemoryProvider struct {
config MemoryProviderConfig
providerConfig agent.ContextProviderConfig
provider agent.ContextProvider

mu sync.Mutex
pendingUpdate *runtime.Poller[azaiprojects.MemoryStoresClientUpdateMemoriesResponse]
}

// NewMemoryProvider creates a Foundry memory provider backed by a Microsoft Foundry Projects memory store.
Expand Down Expand Up @@ -223,14 +229,46 @@ func (p *MemoryProvider) store(ctx context.Context, invoked agent.InvokedContext
Items: items,
UpdateDelay: &p.config.UpdateDelay,
}
if _, err := p.client.BeginUpdateMemories(ctx, p.memoryStoreName, scope, updateOptions); err != nil {
poller, err := p.client.BeginUpdateMemories(ctx, p.memoryStoreName, scope, updateOptions)
if err != nil {
p.log(ctx, slog.LevelError, "foundrymemory: failed to update memories", "memory_store", p.memoryStoreName, "count", len(items), "error", err)
return nil
}
p.mu.Lock()
p.pendingUpdate = poller
p.mu.Unlock()
p.log(ctx, slog.LevelInfo, "foundrymemory: submitted memory update", "memory_store", p.memoryStoreName, "count", len(items))
return nil
}

// WhenUpdatesCompleted blocks until the memory update most recently submitted by
// the provider reaches a terminal state, polling the update status at
// pollingInterval. It returns nil when no update is pending or when the update
// completed successfully, and a non-nil error when the update failed.
//
// Foundry extracts memories asynchronously after a run, so [MemoryProvider.store]
// is fire-and-forget and returns before memories are persisted. Callers that must
// observe stored memories before the next search—typically tests—can await
// completion with this method. Normal runs need not call it and do not pay the
// polling cost.
func (p *MemoryProvider) WhenUpdatesCompleted(ctx context.Context, pollingInterval time.Duration) error {
p.mu.Lock()
poller := p.pendingUpdate
p.mu.Unlock()
if poller == nil {
return nil
}
if _, err := poller.PollUntilDone(ctx, &runtime.PollUntilDoneOptions{Frequency: pollingInterval}); err != nil {

@github-actions github-actions Bot Aug 20, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Parity gap — superseded terminal state not handled

The .NET WhenUpdatesCompletedAsync (FoundryMemoryProvider.cs) explicitly treats superseded as a success terminal state alongside completed:

if (status == MemoryStoreUpdateStatus.Completed || status == MemoryStoreUpdateStatus.Superseded)
{
    return;
}

This Go implementation delegates to runtime.Poller.PollUntilDone, which only knows about generic Azure LRO states (Succeeded, Failed, Canceled). The Foundry-specific superseded state is not a standard Azure LRO terminal state, so the poller will likely loop until ctx is cancelled if an update is superseded by a newer one.

The PR description flags this as a known follow-up, but callers awaiting in a test or workflow may hang silently. Consider adding an explicit doc note to the godoc comment warning that a superseded update will block until the context deadline, and linking to a tracking issue.

@github-actions github-actions Bot Aug 22, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Parity gap — required vs optional pollingInterval

In the .NET implementation, WhenUpdatesCompletedAsync makes pollingInterval optional with a 5-second default:

public async Task WhenUpdatesCompletedAsync(
    TimeSpan? pollingInterval = null,     // defaults to 5 seconds
    CancellationToken cancellationToken = default)

The Go signature requires callers to pass an explicit pollingInterval:

func (p *MemoryProvider) WhenUpdatesCompleted(ctx context.Context, pollingInterval time.Duration) error

This is a minor ergonomics gap: callers must always choose a value rather than relying on a sensible default. Consider treating a zero pollingInterval as "use the default 5 s interval" to stay aligned with the .NET UX.

return err

@github-actions github-actions Bot Aug 4, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Parity gap — superseded terminal state

The .NET WhenUpdatesCompletedAsync explicitly treats MemoryStoreUpdateStatus.Superseded as a successful terminal state (returns without error), mirroring Foundry service semantics where a superseded update is implicitly replaced by a newer one.

The Go implementation delegates to runtime.Poller.PollUntilDone, which only knows standard Azure LRO states and does not treat superseded as terminal. As a result, WhenUpdatesCompleted will keep polling past a superseded status until the service returns a different state or the context times out — diverging from the .NET behaviour.

The PR acknowledges this in the open design questions, but it is worth flagging explicitly since .NET has already shipped the fix:

// .NET — FoundryMemoryProvider.cs
if (status == MemoryStoreUpdateStatus.Completed || status == MemoryStoreUpdateStatus.Superseded)
    return; // treat superseded as success

Suggestion: before merging, consider either (a) using a custom PollingHandler / response classifier in the Azure Go SDK so the poller treats superseded as done, or (b) replacing PollUntilDone with a manual polling loop that mirrors .NET's WaitForUpdateAsync and explicitly handles superseded → success.

}
p.mu.Lock()
if p.pendingUpdate == poller {
p.pendingUpdate = nil
}
p.mu.Unlock()

@github-actions github-actions Bot Aug 22, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Parity gap — superseded terminal state not handled

The .NET implementation explicitly treats superseded as a success terminal state:

if (status == MemoryStoreUpdateStatus.Completed || status == MemoryStoreUpdateStatus.Superseded)
{
    return;
}

The Go implementation delegates polling to runtime.Poller.PollUntilDone from the Azure SDK, which only treats Succeeded / Completed as terminal success. If Foundry ever returns a superseded status for a polled LRO, PollUntilDone may spin indefinitely or return an error rather than treating it as success.

The PR description already flags this as an open question. For parity with .NET, superseded should be treated as a success outcome (the update was replaced by a newer one, so the net effect is the same). Consider wrapping the poller or post-processing the error to handle superseded as a no-op success.

return nil
}

func (p *MemoryProvider) log(ctx context.Context, level slog.Level, msg string, args ...any) {
if p.config.Logger == nil || !p.config.Logger.Enabled(ctx, level) {
return
Expand Down
81 changes: 81 additions & 0 deletions provider/foundryprovider/memory_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@ import (
"net/http"
"strings"
"testing"
"time"

"github.com/Azure/azure-sdk-for-go/sdk/azcore"
"github.com/Azure/azure-sdk-for-go/sdk/azcore/policy"
Expand Down Expand Up @@ -294,6 +295,86 @@ func TestMemoryProviderInvokedLogsUpdateFailureAndDoesNotReturnError(t *testing.
}
}

func TestMemoryProviderWhenUpdatesCompletedAwaitsPendingUpdate(t *testing.T) {
pollResponses := []string{
`{"update_id":"update_1","status":"queued"}`,
`{"update_id":"update_1","status":"completed"}`,
}
polls := 0
transport := &recordingTransport{handle: func(req *http.Request, _ string) (*http.Response, error) {
if req.Method == http.MethodPost && strings.HasSuffix(req.URL.Path, ":update_memories") {
resp := jsonResponse(req, http.StatusAccepted, `{"update_id":"update_1","status":"queued"}`)
resp.Header.Set("Operation-Location", validEndpoint+"/memory_stores/memory/updates/update_1?api-version=v1")
return resp, nil
}
body := pollResponses[len(pollResponses)-1]
if polls < len(pollResponses) {
body = pollResponses[polls]
}
polls++
return jsonResponse(req, http.StatusOK, body), nil
}}
provider := foundryprovider.NewMemoryProvider(validEndpoint, validCredential, "memory", validScope, foundryprovider.MemoryProviderConfig{
ClientOptions: azcore.ClientOptions{Transport: transport},
})

if err := provider.Invoked(t.Context(), agent.InvokedContext{RequestMessages: []*message.Message{message.NewText("remember me")}}); err != nil {
t.Fatalf("Invoked error = %v", err)
}

if err := provider.WhenUpdatesCompleted(t.Context(), time.Millisecond); err != nil {
t.Fatalf("WhenUpdatesCompleted error = %v", err)
}
if polls < 2 {
t.Fatalf("poll count = %d, want at least 2", polls)
}

// The pending update is cleared, so a second call is a no-op and issues no further requests.
before := len(transport.Requests())
if err := provider.WhenUpdatesCompleted(t.Context(), time.Millisecond); err != nil {
t.Fatalf("second WhenUpdatesCompleted error = %v", err)
}
if after := len(transport.Requests()); after != before {
t.Fatalf("request count after second call = %d, want %d", after, before)
}
}

func TestMemoryProviderWhenUpdatesCompletedReturnsErrorOnFailedUpdate(t *testing.T) {
transport := &recordingTransport{handle: func(req *http.Request, _ string) (*http.Response, error) {
if req.Method == http.MethodPost && strings.HasSuffix(req.URL.Path, ":update_memories") {
resp := jsonResponse(req, http.StatusAccepted, `{"update_id":"update_1","status":"queued"}`)
resp.Header.Set("Operation-Location", validEndpoint+"/memory_stores/memory/updates/update_1?api-version=v1")
return resp, nil
}
return jsonResponse(req, http.StatusOK, `{"update_id":"update_1","status":"failed","error":{"code":"extraction_failed","message":"boom"}}`), nil
}}
provider := foundryprovider.NewMemoryProvider(validEndpoint, validCredential, "memory", validScope, foundryprovider.MemoryProviderConfig{
ClientOptions: azcore.ClientOptions{Transport: transport},
})

if err := provider.Invoked(t.Context(), agent.InvokedContext{RequestMessages: []*message.Message{message.NewText("remember me")}}); err != nil {
t.Fatalf("Invoked error = %v", err)
}

if err := provider.WhenUpdatesCompleted(t.Context(), time.Millisecond); err == nil {
t.Fatal("WhenUpdatesCompleted error = nil, want non-nil for failed update")
}
}

func TestMemoryProviderWhenUpdatesCompletedReturnsNilWhenNoPendingUpdate(t *testing.T) {
transport := &recordingTransport{}
provider := foundryprovider.NewMemoryProvider(validEndpoint, validCredential, "memory", validScope, foundryprovider.MemoryProviderConfig{
ClientOptions: azcore.ClientOptions{Transport: transport},
})

if err := provider.WhenUpdatesCompleted(t.Context(), time.Millisecond); err != nil {
t.Fatalf("WhenUpdatesCompleted error = %v", err)
}
if got := len(transport.Requests()); got != 0 {
t.Fatalf("request count = %d, want 0", got)
}
}

func TestMemoryProviderEnsureMemoryStoreCreated(t *testing.T) {
description := "team memory"
tests := []struct {
Expand Down
Loading