Make the todo tool thread safe

Signed-off-by: David Gageot <david.gageot@docker.com>

Author: dgageot

pkg/concurrent/map.go

@@ -0,0 +1,47 @@
+package concurrent
+
+import "sync"
+
+type Map[K comparable, V any] struct {
+	mu     sync.RWMutex
+	values map[K]V
+}
+
+func NewMap[K comparable, V any]() *Map[K, V] {
+	return &Map[K, V]{
+		values: make(map[K]V),
+	}
+}
+
+func (m *Map[K, V]) Load(key K) (V, bool) {
+	m.mu.RLock()
+	defer m.mu.RUnlock()
+
+	val, ok := m.values[key]
+	return val, ok
+}
+
+func (m *Map[K, V]) Store(key K, value V) {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+
+	m.values[key] = value
+}
+
+func (m *Map[K, V]) Length() int {
+	m.mu.RLock()
+	defer m.mu.RUnlock()
+
+	return len(m.values)
+}
+
+func (m *Map[K, V]) Range(f func(key K, value V) bool) {
+	m.mu.RLock()
+	defer m.mu.RUnlock()
+
+	for k, v := range m.values {
+		if !f(k, v) {
+			break
+		}
+	}
+}

pkg/tools/builtin/todo.go

@@ -7,6 +7,7 @@ import (
 	"strings"
 	"sync"
 
+	"github.com/docker/cagent/pkg/concurrent"
 	"github.com/docker/cagent/pkg/tools"
 )
 
@@ -38,15 +39,15 @@ type UpdateTodoArgs struct {
 }
 
 type todoHandler struct {
-	todos map[string]Todo
+	todos *concurrent.Map[string, Todo]
 }
 
 var NewSharedTodoTool = sync.OnceValue(NewTodoTool)
 
 func NewTodoTool() *TodoTool {
 	return &TodoTool{
 		handler: &todoHandler{
-			todos: make(map[string]Todo),
+			todos: concurrent.NewMap[string, Todo](),
 		},
 	}
 }
@@ -78,12 +79,12 @@ func (h *todoHandler) createTodo(_ context.Context, toolCall tools.ToolCall) (*t
 		return nil, fmt.Errorf("invalid arguments: %w", err)
 	}
 
-	id := fmt.Sprintf("todo_%d", len(h.todos)+1)
-	h.todos[id] = Todo{
+	id := fmt.Sprintf("todo_%d", h.todos.Length()+1)
+	h.todos.Store(id, Todo{
 		ID:          id,
 		Description: params.Description,
 		Status:      "pending",
-	}
+	})
 
 	return &tools.ToolCallResult{
 		Output: fmt.Sprintf("Created todo [%s]: %s", id, params.Description),
@@ -97,14 +98,14 @@ func (h *todoHandler) createTodos(_ context.Context, toolCall tools.ToolCall) (*
 	}
 
 	ids := make([]string, len(params.Descriptions))
-	start := len(h.todos)
+	start := h.todos.Length()
 	for i, desc := range params.Descriptions {
 		id := fmt.Sprintf("todo_%d", start+i+1)
-		h.todos[id] = Todo{
+		h.todos.Store(id, Todo{
 			ID:          id,
 			Description: desc,
 			Status:      "pending",
-		}
+		})
 		ids[i] = id
 	}
 
@@ -127,13 +128,13 @@ func (h *todoHandler) updateTodo(_ context.Context, toolCall tools.ToolCall) (*t
 		return nil, fmt.Errorf("invalid arguments: %w", err)
 	}
 
-	todo, exists := h.todos[params.ID]
+	todo, exists := h.todos.Load(params.ID)
 	if !exists {
 		return nil, fmt.Errorf("todo [%s] not found", params.ID)
 	}
 
 	todo.Status = params.Status
-	h.todos[params.ID] = todo
+	h.todos.Store(params.ID, todo)
 
 	return &tools.ToolCallResult{
 		Output: fmt.Sprintf("Updated todo [%s] to status: [%s]", params.ID, params.Status),
@@ -144,10 +145,11 @@ func (h *todoHandler) listTodos(context.Context, tools.ToolCall) (*tools.ToolCal
 	var output strings.Builder
 	output.WriteString("Current todos:\n")
 
-	for _, todo := range h.todos {
+	h.todos.Range(func(_ string, todo Todo) bool {
 		output.WriteString(fmt.Sprintf("- [%s] %s (Status: %s)\n",
 			todo.ID, todo.Description, todo.Status))
-	}
+		return true
+	})
 
 	return &tools.ToolCallResult{
 		Output: output.String(),

pkg/tools/builtin/todo_test.go

@@ -15,7 +15,7 @@ func TestNewTodoTool(t *testing.T) {
 
 	assert.NotNil(t, tool)
 	assert.NotNil(t, tool.handler)
-	assert.Empty(t, tool.handler.todos)
+	assert.Zero(t, tool.handler.todos.Length())
 }
 
 func TestTodoTool_Instructions(t *testing.T) {
@@ -152,8 +152,8 @@ func TestTodoTool_CreateTodo(t *testing.T) {
 	assert.Contains(t, result.Output, "Created todo [todo_1]: Test todo item")
 
 	// Verify todo was added to the handler's todos map
-	assert.Len(t, tool.handler.todos, 1)
-	todo, exists := tool.handler.todos["todo_1"]
+	assert.Equal(t, 1, tool.handler.todos.Length())
+	todo, exists := tool.handler.todos.Load("todo_1")
 	assert.True(t, exists)
 	assert.Equal(t, "Test todo item", todo.Description)
 	assert.Equal(t, "pending", todo.Status)
@@ -198,7 +198,7 @@ func TestTodoTool_CreateTodos(t *testing.T) {
 	assert.Contains(t, result.Output, "todo_3")
 
 	// Verify todos were added to the handler's todos map
-	assert.Len(t, tool.handler.todos, 3)
+	assert.Equal(t, 3, tool.handler.todos.Length())
 
 	// Create multiple todos
 	args = CreateTodosArgs{
@@ -222,7 +222,7 @@ func TestTodoTool_CreateTodos(t *testing.T) {
 	require.NoError(t, err)
 	assert.Contains(t, result.Output, "Created 1 todos:")
 	assert.Contains(t, result.Output, "todo_4")
-	assert.Len(t, tool.handler.todos, 4)
+	assert.Equal(t, 4, tool.handler.todos.Length())
 }
 
 func TestTodoTool_UpdateTodo(t *testing.T) {
@@ -276,7 +276,7 @@ func TestTodoTool_UpdateTodo(t *testing.T) {
 	assert.Contains(t, result.Output, "Updated todo [todo_1] to status: [completed]")
 
 	// Verify todo status was updated
-	todo, exists := tool.handler.todos["todo_1"]
+	todo, exists := tool.handler.todos.Load("todo_1")
 	assert.True(t, exists)
 	assert.Equal(t, "completed", todo.Status)
 }