Merge pull request #1392 from trheyi/main

Enhance Memory Namespace Isolation Tests and Functionality

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
+	t.Run("User isolation", func(t *testing.T) {
-	mem1, err := memory.New(nil, "user1", "", "", "")
+		// Create two memory instances with different user IDs
-	require.NoError(t, err)
+		mem1, err := memory.New(nil, "user1", "", "", "")
+		require.NoError(t, err)
 
-	mem2, err := memory.New(nil, "user2", "", "", "")
+		mem2, err := memory.New(nil, "user2", "", "", "")
-	require.NoError(t, err)
+		require.NoError(t, err)
 
-	// Set value in user1's namespace
+		// Set value in user1's namespace
-	err = mem1.GetUser().Set("key", "user1_value", 0)
+		err = mem1.GetUser().Set("key", "user1_value", 0)
-	require.NoError(t, err)
+		require.NoError(t, err)
 
-	// Set value in user2's namespace
+		// Set value in user2's namespace
-	err = mem2.GetUser().Set("key", "user2_value", 0)
+		err = mem2.GetUser().Set("key", "user2_value", 0)
-	require.NoError(t, err)
+		require.NoError(t, err)
 
-	// Verify isolation
+		// Verify isolation - each user sees their own value
-	val1, ok := mem1.GetUser().Get("key")
+		val1, ok := mem1.GetUser().Get("key")
-	assert.True(t, ok)
+		assert.True(t, ok)
-	assert.Equal(t, "user1_value", val1)
+		assert.Equal(t, "user1_value", val1)
 
-	val2, ok := mem2.GetUser().Get("key")
+		val2, ok := mem2.GetUser().Get("key")
-	assert.True(t, ok)
+		assert.True(t, ok)
-	assert.Equal(t, "user2_value", val2)
+		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) {

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 {
+// Uses pattern-based query for efficiency
-	allKeys := ns.Store.Keys()
+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
+	// Remove prefix from keys
-	var result []string
+	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 {
+// Uses pattern-based query for efficiency
-	return len(ns.Keys())
+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