Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
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
102 changes: 78 additions & 24 deletions provider/foundryprovider/memory.go
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,11 @@ const (
defaultMaxMemories = 5
defaultSourceID = "foundrymemory"

// Session-state keys for the static (item-less) memory cache, fetched once
// per session and reused on later turns.
memoryStaticInitStateKey = "foundrymemory.staticInitialized"
memoryStaticMemoriesStateKey = "foundrymemory.staticMemories"
Comment on lines +28 to +29

// memorySearchIDStateKey stores the last search id on the session so that
// successive turns issue incremental memory searches.
memorySearchIDStateKey = "foundrymemory.previousSearchID"
Expand Down Expand Up @@ -203,34 +208,44 @@ func (p *MemoryProvider) EnsureStoredMemoriesDeleted(ctx context.Context, sessio

func (p *MemoryProvider) provide(ctx context.Context, invoking agent.InvokingContext) ([]*message.Message, []agent.Option, error) {
telemetry.MarkUsed(telemetry.FeatureFoundryMemory)
items := searchMemoryItems(invoking.Messages)
if len(items) == 0 {
return nil, nil, nil
}
session, _ := agent.GetOption(invoking.Options, agent.WithSession)
scope := p.scope(session)
searchOptions := &azaiprojects.MemoryStoresClientSearchMemoriesOptions{
Items: items,
Options: &azaiprojects.MemorySearchResultOptions{MaxMemories: p.config.MaxMemories},
}
// Thread the previous search id so successive turns perform incremental
// searches, matching the Python provider.
var previousSearchID string
if ok, _ := session.Get(memorySearchIDStateKey, &previousSearchID); ok && previousSearchID != "" {
searchOptions.PreviousSearchID = &previousSearchID
}
result, err := p.client.SearchMemories(ctx, p.memoryStoreName, scope, searchOptions)
if err != nil {
if isCallerCancellation(ctx, err) {
return nil, nil, ctx.Err()

// Static (e.g. user-profile) memories are fetched once per session with an
// item-less search and prepended to the contextual memories on every turn,
// matching the Python provider.
staticMemories := p.staticMemories(ctx, session, scope)

var contextual []string
if items := searchMemoryItems(invoking.Messages); len(items) > 0 {
searchOptions := &azaiprojects.MemoryStoresClientSearchMemoriesOptions{
Items: items,
Options: &azaiprojects.MemorySearchResultOptions{MaxMemories: p.config.MaxMemories},
}
// Thread the previous search id so successive turns perform incremental
// searches, matching the Python provider.
var previousSearchID string
if ok, _ := session.Get(memorySearchIDStateKey, &previousSearchID); ok && previousSearchID != "" {
searchOptions.PreviousSearchID = &previousSearchID
}
result, err := p.client.SearchMemories(ctx, p.memoryStoreName, scope, searchOptions)
if err != nil {
if isCallerCancellation(ctx, err) {
return nil, nil, ctx.Err()
}
p.log(ctx, slog.LevelError, "foundrymemory: failed to search memories", "memory_store", p.memoryStoreName, "error", err)
} else {
if result.SearchID != nil && *result.SearchID != "" {
session.Set(memorySearchIDStateKey, *result.SearchID)
}
contextual = memoryContents(result.Memories)
}
p.log(ctx, slog.LevelError, "foundrymemory: failed to search memories", "memory_store", p.memoryStoreName, "error", err)
return nil, nil, nil
}
if result.SearchID != nil && *result.SearchID != "" {
session.Set(memorySearchIDStateKey, *result.SearchID)
}
memories := memoryContents(result.Memories)

// Copy into a fresh slice so appending never mutates the cached static list.
memories := make([]string, 0, len(staticMemories)+len(contextual))
memories = append(memories, staticMemories...)
memories = append(memories, contextual...)
p.log(ctx, slog.LevelInfo, "foundrymemory: retrieved memories", "memory_store", p.memoryStoreName, "count", len(memories))
if len(memories) == 0 {
return nil, nil, nil
Expand All @@ -239,6 +254,45 @@ func (p *MemoryProvider) provide(ctx context.Context, invoking agent.InvokingCon
return []*message.Message{contextMessage}, nil, nil
}

// staticMemories returns the scope's static (item-less) memories, fetching them
// once per session and caching them on the session for reuse on later turns.
// Without a session there is nowhere to cache, so it returns nothing rather
// than issuing an item-less search on every turn.
func (p *MemoryProvider) staticMemories(ctx context.Context, session *agent.Session, scope string) []string {
if session == nil {
return nil
}
// Key the cache by store and scope so a session shared across providers, or
// a scope callback whose result changes between turns, does not reuse the
// wrong scope's static memories.
initKey := p.staticStateKey(memoryStaticInitStateKey, scope)
memKey := p.staticStateKey(memoryStaticMemoriesStateKey, scope)
var initialized bool
if ok, _ := session.Get(initKey, &initialized); ok && initialized {
var cached []string
session.Get(memKey, &cached)
return cached
}
result, err := p.client.SearchMemories(ctx, p.memoryStoreName, scope, &azaiprojects.MemoryStoresClientSearchMemoriesOptions{
Options: &azaiprojects.MemorySearchResultOptions{MaxMemories: p.config.MaxMemories},
})
if err != nil {

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.

The retry-on-failure behavior diverges from the upstream Python provider.

Here, when the item-less static search fails, initKey is left unset so every subsequent turn re-issues the static search until it succeeds (per the "Leave the session uninitialized so the next turn retries" comment).

Upstream python/packages/foundry/agent_framework_foundry/_memory_provider.py (FoundryMemoryProvider.before_run) does the opposite: the try/except/finally block sets state["static_memories"] = [] on failure and unconditionally sets state["initialized"] = True in the finally clause, so a failed static search is attempted only once per session and the provider silently proceeds with no static memories thereafter — it never retries on later turns.

This is a meaningful parity gap: a persistently failing/unavailable memory store causes the Go provider to retry the item-less search on every turn indefinitely (extra latency/cost per turn), while Python fails once and moves on. Consider caching an "attempted" flag (with empty static memories) on failure, matching the Python finally semantics, unless the indefinite-retry behavior is an intentional Go-specific choice.

// Leave the session uninitialized so the next turn retries.

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.

This staticMemories helper leaves the session uninitialized on a failed item-less search so the next turn retries. The Python reference this PR states it's porting (python/packages/foundry/agent_framework_foundry/_memory_provider.py, before_run) always sets state["initialized"] = True in a finally block regardless of success, caching an empty static-memory list and never retrying the item-less search again for that session (test_static_memories_only_retrieved_once covers this).

So Go will keep re-issuing the item-less static search every subsequent turn after a transient failure, while Python treats the static fetch as one-shot best-effort per session. Since the PR explicitly aims to match the Python provider's caching semantics, consider setting the initialized flag (with an empty cached list) even on failure to match Python, or note this retry-on-failure behavior as an intentional Go-specific divergence.

p.log(ctx, slog.LevelError, "foundrymemory: failed to search static memories", "memory_store", p.memoryStoreName, "error", err)
return nil
}
Comment on lines +279 to +283

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 issue: static-memory failure retry policy diverges from the Python provider.

On a failed item-less static search, this leaves the session un-initialized ("Leave the session uninitialized so the next turn retries"), so every subsequent turn re-issues the static search until it succeeds.

The upstream Python provider does the opposite and does so deliberately. In _memory_provider.py, before_run wraps the static search in try/except/finally and sets state["initialized"] = True in the finally block — i.e. regardless of success or failure — with the explicit comment "Mark as initialized regardless of success to avoid repeated attempts". state["static_memories"] is set to [] on failure. This is also asserted by the upstream test test_static_memories_only_retrieved_once (and the failure-path test around line 293 in test_foundry_memory_provider.py), which expect exactly one static search per session even when it errors.

Go's current behavior (retry on every turn after a failure) means a persistently-failing or slow memory store is hit with a static search on every single agent turn for the life of the session, which is a different reliability/cost profile than upstream intentionally chose.

Suggested fix: set initKey to true (and cache an empty/nil static-memories slice) in the err != nil branch as well, so a failed static fetch is treated as "initialized" and not retried, matching the Python provider's documented behavior.

static := memoryContents(result.Memories)
session.Set(memKey, static)
session.Set(initKey, true)
return static
}

// staticStateKey qualifies a static-memory session-state key with the store and
// scope so cached state is not shared across providers or scopes on one session.
func (p *MemoryProvider) staticStateKey(prefix, scope string) string {
return prefix + ":" + p.memoryStoreName + ":" + scope
}

func (p *MemoryProvider) store(ctx context.Context, invoked agent.InvokedContext) error {
telemetry.MarkUsed(telemetry.FeatureFoundryMemory)
items := updateMemoryItems(invoked.RequestMessages)
Expand Down
88 changes: 81 additions & 7 deletions provider/foundryprovider/memory_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -171,6 +171,72 @@ func TestNewMemoryProviderStoreFiltersDefaultToExternalOnlyRequestAndPassThrough
}
}

// Static (user-profile) memories are fetched once per session with an item-less
// search and prepended to the contextual memories on every turn, matching the
// Python provider.
func TestMemoryProviderInjectsStaticMemories(t *testing.T) {
transport := &recordingTransport{}
transport.handle = func(req *http.Request, body string) (*http.Response, error) {
// A populated items array marks the contextual search; the item-less
// static search sends no items.
if strings.Contains(body, `"items":[{`) {
return jsonResponse(req, http.StatusOK, `{"memories":[{"memory_item":{"content":"contextual memory","kind":"procedural","memory_id":"mem_c","scope":"user-456","updated_at":1}}]}`), nil
}
return jsonResponse(req, http.StatusOK, `{"memories":[{"memory_item":{"content":"static profile","kind":"user_profile","memory_id":"mem_s","scope":"user-456","updated_at":1}}]}`), nil
}
provider := foundryprovider.NewMemoryProvider(validEndpoint, validCredential, "memory", validScope, foundryprovider.MemoryProviderConfig{
ClientOptions: azcore.ClientOptions{Transport: transport},
ContextPrompt: new("Memories:"),
})
session := &agent.Session{}

// First turn: item-less static search + contextual search; both surfaced.
messages, _, err := provider.Invoking(t.Context(), agent.InvokingContext{
Messages: []*message.Message{message.NewText("what do you remember?")},
Options: []agent.Option{agent.WithSession(session)},
})
if err != nil {
t.Fatalf("Invoking error = %v", err)
}
// Static memories are prepended before the contextual ones, so assert order,
// not just presence.
if ctx := messagesString(messages); !strings.Contains(ctx, "static profile\ncontextual memory") {
t.Fatalf("first-turn context = %q, want static memory prepended before contextual", ctx)
}

// Second turn: static is cached, so no second item-less search is issued,
// but the static memory is still prepended.
messages, _, err = provider.Invoking(t.Context(), agent.InvokingContext{
Messages: []*message.Message{message.NewText("again")},
Options: []agent.Option{agent.WithSession(session)},
})
if err != nil {
t.Fatalf("Invoking error = %v", err)
}
if ctx := messagesString(messages); !strings.Contains(ctx, "static profile") {
t.Fatalf("second-turn context = %q, want static memory prepended", ctx)
}

// One item-less static search total (turn 1) + one contextual search per turn.
itemless := 0
for _, r := range transport.Requests() {
if !strings.Contains(r.Body, `"items":[{`) {
itemless++
}
}
if itemless != 1 {
t.Fatalf("item-less static searches = %d, want 1 (cached after first turn)", itemless)
}
}

func messagesString(messages []*message.Message) string {
var b strings.Builder
for _, m := range messages {
b.WriteString(m.String())
}
return b.String()
}

func TestMemoryProviderPanicsWhenScopeIsEmptyOnUse(t *testing.T) {
provider := foundryprovider.NewMemoryProvider(validEndpoint, validCredential, "memory", func(*agent.Session) string { return " " }, foundryprovider.MemoryProviderConfig{})
assertPanics(t, func() {
Expand Down Expand Up @@ -270,15 +336,23 @@ func TestMemoryProviderThreadsPreviousSearchID(t *testing.T) {
}
}

requests := transport.Requests()
if len(requests) != 2 {
t.Fatalf("request count = %d, want 2", len(requests))
// Filter to the contextual (item-bearing) searches; the static
// (user-profile) search issued once per session is item-less and does not
// participate in incremental search-id threading.
var contextual []string
for _, r := range transport.Requests() {
if strings.Contains(r.Body, `"items":[{`) {
contextual = append(contextual, r.Body)
}
}
if len(contextual) != 2 {
t.Fatalf("contextual request count = %d, want 2", len(contextual))
}
if got := jsonMap(t, requests[0].Body)["previous_search_id"]; got != nil {
t.Fatalf("first request previous_search_id = %#v, want absent", got)
if got := jsonMap(t, contextual[0])["previous_search_id"]; got != nil {
t.Fatalf("first contextual request previous_search_id = %#v, want absent", got)
}
if got := jsonMap(t, requests[1].Body)["previous_search_id"]; got != "search_1" {
t.Fatalf("second request previous_search_id = %#v, want search_1", got)
if got := jsonMap(t, contextual[1])["previous_search_id"]; got != "search_1" {
t.Fatalf("second contextual request previous_search_id = %#v, want search_1", got)
}
}

Expand Down
Loading