refactor: reset

Carlos Alexandro Becker created 1 month ago

Signed-off-by: Carlos Alexandro Becker <[e-mail]>

Change summary

internal/csync/maps.go      |  7 +++++++
internal/csync/maps_test.go | 16 ++++++++++++++++
internal/llm/agent/agent.go |  3 +--
3 files changed, 24 insertions(+), 2 deletions(-)

Detailed changes

internal/csync/maps.go

@@ -39,6 +39,13 @@ func NewLazyMap[K comparable, V any](load func() map[K]V) *Map[K, V] {
 	return m
 }
 
+// Reset replaces the inner map with the new one.
+func (m *Map[K, V]) Reset(input map[K]V) {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+	m.inner = input
+}
+
 // Set sets the value for the specified key in the map.
 func (m *Map[K, V]) Set(key K, value V) {
 	m.mu.Lock()

internal/csync/maps_test.go

@@ -72,6 +72,22 @@ func TestNewLazyMap(t *testing.T) {
 	})
 }
 
+func TestMap_Reset(t *testing.T) {
+	t.Parallel()
+
+	m := NewMapFrom(map[string]int{
+		"a": 10,
+	})
+
+	m.Reset(map[string]int{
+		"b": 20,
+	})
+	value, ok := m.Get("b")
+	require.True(t, ok)
+	require.Equal(t, 20, value)
+	require.Equal(t, 1, m.Len())
+}
+
 func TestMap_Set(t *testing.T) {
 	t.Parallel()
 

internal/llm/agent/agent.go

@@ -1129,8 +1129,7 @@ func (a *agent) setupEvents(ctx context.Context) {
 						continue
 					}
 					updateMcpTools(name, tools)
-					// Update the lazy map with the new tools
-					a.mcpTools = csync.NewMapFrom(maps.Collect(mcpTools.Seq2()))
+					a.mcpTools.Reset(maps.Collect(mcpTools.Seq2()))
 					updateMCPState(name, MCPStateConnected, nil, c, a.mcpTools.Len())
 				default:
 					continue