From 6ce6b65567809a9eb49871ade5b55bb4ecfb7588 Mon Sep 17 00:00:00 2001 From: Max Date: Mon, 22 Dec 2025 14:54:42 +0800 Subject: [PATCH] Enhance Memory Namespace Isolation Tests and Functionality - Expanded the memory isolation tests to cover user, team, chat, context, and key operations, ensuring comprehensive validation of namespace behavior. - Implemented pattern-based key retrieval and length calculation in the Namespace struct for improved efficiency and flexibility. - Added assertions to verify that operations on one namespace do not affect others, reinforcing the integrity of memory isolation across different contexts. --- agent/memory/memory_test.go | 244 +++++++++++++++++++++++++++++++++--- agent/memory/namespace.go | 32 +++-- 2 files changed, 251 insertions(+), 25 deletions(-) diff --git a/agent/memory/memory_test.go b/agent/memory/memory_test.go index cf238c5c..bd7dcfd3 100644 --- a/agent/memory/memory_test.go +++ b/agent/memory/memory_test.go @@ -123,29 +123,237 @@ func TestNamespaceIsolation(t *testing.T) { test.Prepare(t, config.Conf) defer test.Clean() - // Create two memory instances with different user IDs - mem1, err := memory.New(nil, "user1", "", "", "") - require.NoError(t, err) + t.Run("User isolation", func(t *testing.T) { + // Create two memory instances with different user IDs + mem1, err := memory.New(nil, "user1", "", "", "") + require.NoError(t, err) - mem2, err := memory.New(nil, "user2", "", "", "") - require.NoError(t, err) + mem2, err := memory.New(nil, "user2", "", "", "") + require.NoError(t, err) - // Set value in user1's namespace - err = mem1.GetUser().Set("key", "user1_value", 0) - require.NoError(t, err) + // Set value in user1's namespace + err = mem1.GetUser().Set("key", "user1_value", 0) + require.NoError(t, err) - // Set value in user2's namespace - err = mem2.GetUser().Set("key", "user2_value", 0) - require.NoError(t, err) + // Set value in user2's namespace + err = mem2.GetUser().Set("key", "user2_value", 0) + require.NoError(t, err) - // Verify isolation - val1, ok := mem1.GetUser().Get("key") - assert.True(t, ok) - assert.Equal(t, "user1_value", val1) + // Verify isolation - each user sees their own value + val1, ok := mem1.GetUser().Get("key") + assert.True(t, ok) + assert.Equal(t, "user1_value", val1) - val2, ok := mem2.GetUser().Get("key") - assert.True(t, ok) - assert.Equal(t, "user2_value", val2) + val2, ok := mem2.GetUser().Get("key") + assert.True(t, ok) + assert.Equal(t, "user2_value", val2) + + // Delete from user1 should not affect user2 + err = mem1.GetUser().Del("key") + require.NoError(t, err) + + _, ok = mem1.GetUser().Get("key") + assert.False(t, ok, "user1's key should be deleted") + + val2, ok = mem2.GetUser().Get("key") + assert.True(t, ok, "user2's key should still exist") + assert.Equal(t, "user2_value", val2) + + // Clear user1 should not affect user2 + mem1.GetUser().Clear() + val2, ok = mem2.GetUser().Get("key") + assert.True(t, ok, "user2's key should still exist after user1 clear") + assert.Equal(t, "user2_value", val2) + }) + + t.Run("Team isolation", func(t *testing.T) { + memA, err := memory.New(nil, "", "teamA", "", "") + require.NoError(t, err) + + memB, err := memory.New(nil, "", "teamB", "", "") + require.NoError(t, err) + + // Set same key in different teams + memA.GetTeam().Set("config", "teamA_config", 0) + memB.GetTeam().Set("config", "teamB_config", 0) + + // Verify isolation + valA, ok := memA.GetTeam().Get("config") + assert.True(t, ok) + assert.Equal(t, "teamA_config", valA) + + valB, ok := memB.GetTeam().Get("config") + assert.True(t, ok) + assert.Equal(t, "teamB_config", valB) + }) + + t.Run("Chat isolation", func(t *testing.T) { + mem1, err := memory.New(nil, "", "", "chat1", "") + require.NoError(t, err) + + mem2, err := memory.New(nil, "", "", "chat2", "") + require.NoError(t, err) + + // Set same key in different chats + mem1.GetChat().Set("topic", "chat1_topic", 0) + mem2.GetChat().Set("topic", "chat2_topic", 0) + + // Verify isolation + val1, ok := mem1.GetChat().Get("topic") + assert.True(t, ok) + assert.Equal(t, "chat1_topic", val1) + + val2, ok := mem2.GetChat().Get("topic") + assert.True(t, ok) + assert.Equal(t, "chat2_topic", val2) + }) + + t.Run("Context isolation", func(t *testing.T) { + mem1, err := memory.New(nil, "", "", "", "ctx1") + require.NoError(t, err) + + mem2, err := memory.New(nil, "", "", "", "ctx2") + require.NoError(t, err) + + // Set same key in different contexts + mem1.GetContext().Set("temp", "ctx1_temp", 0) + mem2.GetContext().Set("temp", "ctx2_temp", 0) + + // Verify isolation + val1, ok := mem1.GetContext().Get("temp") + assert.True(t, ok) + assert.Equal(t, "ctx1_temp", val1) + + val2, ok := mem2.GetContext().Get("temp") + assert.True(t, ok) + assert.Equal(t, "ctx2_temp", val2) + }) + + t.Run("Keys and Len isolation", func(t *testing.T) { + mem1, err := memory.New(nil, "userA", "", "", "") + require.NoError(t, err) + + mem2, err := memory.New(nil, "userB", "", "", "") + require.NoError(t, err) + + // Clear first + mem1.GetUser().Clear() + mem2.GetUser().Clear() + + // Set keys in userA + mem1.GetUser().Set("a", 1, 0) + mem1.GetUser().Set("b", 2, 0) + mem1.GetUser().Set("c", 3, 0) + + // Set keys in userB + mem2.GetUser().Set("x", 10, 0) + mem2.GetUser().Set("y", 20, 0) + + // Verify Keys isolation + keys1 := mem1.GetUser().Keys() + assert.Equal(t, 3, len(keys1), "userA should have 3 keys") + + keys2 := mem2.GetUser().Keys() + assert.Equal(t, 2, len(keys2), "userB should have 2 keys") + + // Verify Len isolation + assert.Equal(t, 3, mem1.GetUser().Len(), "userA Len should be 3") + assert.Equal(t, 2, mem2.GetUser().Len(), "userB Len should be 2") + + // Keys should not contain prefix + for _, k := range keys1 { + assert.NotContains(t, k, "user:", "Key should not contain prefix") + } + }) + + t.Run("Incr/Decr isolation", func(t *testing.T) { + mem1, err := memory.New(nil, "userX", "", "", "") + require.NoError(t, err) + + mem2, err := memory.New(nil, "userY", "", "", "") + require.NoError(t, err) + + // Incr counter in userX + val1, err := mem1.GetUser().Incr("counter", 10) + require.NoError(t, err) + assert.Equal(t, int64(10), val1) + + // Incr counter in userY + val2, err := mem2.GetUser().Incr("counter", 5) + require.NoError(t, err) + assert.Equal(t, int64(5), val2) + + // Incr again - should be independent + val1, err = mem1.GetUser().Incr("counter", 1) + require.NoError(t, err) + assert.Equal(t, int64(11), val1) + + val2, err = mem2.GetUser().Incr("counter", 1) + require.NoError(t, err) + assert.Equal(t, int64(6), val2) + }) + + t.Run("List operations isolation", func(t *testing.T) { + mem1, err := memory.New(nil, "listUser1", "", "", "") + require.NoError(t, err) + + mem2, err := memory.New(nil, "listUser2", "", "", "") + require.NoError(t, err) + + // Push to user1's list + err = mem1.GetUser().Push("items", "a", "b", "c") + require.NoError(t, err) + + // Push to user2's list + err = mem2.GetUser().Push("items", "x", "y") + require.NoError(t, err) + + // Verify isolation + assert.Equal(t, 3, mem1.GetUser().ArrayLen("items")) + assert.Equal(t, 2, mem2.GetUser().ArrayLen("items")) + + all1, _ := mem1.GetUser().ArrayAll("items") + all2, _ := mem2.GetUser().ArrayAll("items") + + assert.Equal(t, 3, len(all1)) + assert.Equal(t, 2, len(all2)) + + // Pop from user1 should not affect user2 + mem1.GetUser().Pop("items", 1) + assert.Equal(t, 2, mem1.GetUser().ArrayLen("items")) + assert.Equal(t, 2, mem2.GetUser().ArrayLen("items")) + }) + + t.Run("Del pattern isolation", func(t *testing.T) { + mem1, err := memory.New(nil, "patternUser1", "", "", "") + require.NoError(t, err) + + mem2, err := memory.New(nil, "patternUser2", "", "", "") + require.NoError(t, err) + + // Set keys with pattern in both users + mem1.GetUser().Set("file:1", "data1", 0) + mem1.GetUser().Set("file:2", "data2", 0) + mem1.GetUser().Set("other", "other1", 0) + + mem2.GetUser().Set("file:1", "data1", 0) + mem2.GetUser().Set("file:2", "data2", 0) + mem2.GetUser().Set("other", "other2", 0) + + // Delete pattern from user1 + err = mem1.GetUser().Del("file:*") + require.NoError(t, err) + + // user1's file:* keys should be deleted + assert.False(t, mem1.GetUser().Has("file:1")) + assert.False(t, mem1.GetUser().Has("file:2")) + assert.True(t, mem1.GetUser().Has("other")) + + // user2's keys should be unaffected + assert.True(t, mem2.GetUser().Has("file:1")) + assert.True(t, mem2.GetUser().Has("file:2")) + assert.True(t, mem2.GetUser().Has("other")) + }) } func TestNamespaceIncrDecr(t *testing.T) { diff --git a/agent/memory/namespace.go b/agent/memory/namespace.go index 10cc843a..ee8afd86 100644 --- a/agent/memory/namespace.go +++ b/agent/memory/namespace.go @@ -46,14 +46,23 @@ func (ns *Namespace) Del(key string) error { } // Keys returns all keys in this namespace -func (ns *Namespace) Keys() []string { - allKeys := ns.Store.Keys() +// Uses pattern-based query for efficiency +func (ns *Namespace) Keys(pattern ...string) []string { + // Build pattern with namespace prefix + var storePattern string + if len(pattern) > 0 && pattern[0] != "" { + storePattern = ns.Prefix + pattern[0] + } else { + storePattern = ns.Prefix + "*" + } + + allKeys := ns.Store.Keys(storePattern) prefixLen := len(ns.Prefix) - // Filter keys that belong to this namespace - var result []string + // Remove prefix from keys + result := make([]string, 0, len(allKeys)) for _, key := range allKeys { - if len(key) >= prefixLen && key[:prefixLen] == ns.Prefix { + if len(key) >= prefixLen { result = append(result, key[prefixLen:]) } } @@ -61,8 +70,17 @@ func (ns *Namespace) Keys() []string { } // Len returns the number of keys in this namespace -func (ns *Namespace) Len() int { - return len(ns.Keys()) +// Uses pattern-based query for efficiency +func (ns *Namespace) Len(pattern ...string) int { + // Build pattern with namespace prefix + var storePattern string + if len(pattern) > 0 && pattern[0] != "" { + storePattern = ns.Prefix + pattern[0] + } else { + storePattern = ns.Prefix + "*" + } + + return ns.Store.Len(storePattern) } // Clear deletes all keys in this namespace