From a78cc120c22037bc1d5ebb2f0ce4630031277adb Mon Sep 17 00:00:00 2001
From: Carlos Alexandro Becker <caarlos0@users.noreply.github.com>
Date: Tue, 30 Sep 2025 17:18:41 -0300
Subject: [PATCH] refactor: reset

Signed-off-by: Carlos Alexandro Becker <caarlos0@users.noreply.github.com>
---
 internal/csync/maps.go      |  7 +++++++
 internal/csync/maps_test.go | 16 ++++++++++++++++
 internal/llm/agent/agent.go |  3 +--
 3 files changed, 24 insertions(+), 2 deletions(-)

diff --git a/internal/csync/maps.go b/internal/csync/maps.go
index 8151e4b6029c93196196dddfd501a21e5b2f2179..1fd2005790014b2ce4bd5a78dbb7931d54cbe66c 100644
--- a/internal/csync/maps.go
+++ b/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()
diff --git a/internal/csync/maps_test.go b/internal/csync/maps_test.go
index 2ce29d94e1e59644b2bde6c2de704c609038cbfb..4c590f008dad91e8dcbc40d1b90d87ef1b3e5750 100644
--- a/internal/csync/maps_test.go
+++ b/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()
 
diff --git a/internal/llm/agent/agent.go b/internal/llm/agent/agent.go
index 6888de73da0c44b1d24b141e2de980298b3c9bdd..91f02397bd7ad2610c950719b738965577649f6f 100644
--- a/internal/llm/agent/agent.go
+++ b/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