Merge pull request #566 from dgageot/token-thread-safe

Make Token Store thread safe

Author: dgageot

pkg/concurrent/map.go

@@ -28,6 +28,13 @@ func (m *Map[K, V]) Store(key K, value V) {
 	m.values[key] = value
 }
 
+func (m *Map[K, V]) Delete(key K) {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+
+	delete(m.values, key)
+}
+
 func (m *Map[K, V]) Length() int {
 	m.mu.RLock()
 	defer m.mu.RUnlock()

pkg/tools/mcp/tokenstore.go

@@ -3,6 +3,8 @@ package mcp
 import (
 	"fmt"
 	"time"
+
+	"github.com/docker/cagent/pkg/concurrent"
 )
 
 // OAuthTokenStore manages OAuth tokens
@@ -35,30 +37,30 @@ func (t *OAuthToken) IsExpired() bool {
 
 // InMemoryTokenStore implements OAuthTokenStore in memory
 type InMemoryTokenStore struct {
-	tokens map[string]*OAuthToken
+	tokens *concurrent.Map[string, *OAuthToken]
 }
 
 // NewInMemoryTokenStore creates a new in-memory token store
 func NewInMemoryTokenStore() OAuthTokenStore {
 	return &InMemoryTokenStore{
-		tokens: make(map[string]*OAuthToken),
+		tokens: concurrent.NewMap[string, *OAuthToken](),
 	}
 }
 
 func (s *InMemoryTokenStore) GetToken(resourceURL string) (*OAuthToken, error) {
-	token, ok := s.tokens[resourceURL]
+	token, ok := s.tokens.Load(resourceURL)
 	if !ok {
 		return nil, fmt.Errorf("no token found for resource: %s", resourceURL)
 	}
 	return token, nil
 }
 
 func (s *InMemoryTokenStore) StoreToken(resourceURL string, token *OAuthToken) error {
-	s.tokens[resourceURL] = token
+	s.tokens.Store(resourceURL, token)
 	return nil
 }
 
 func (s *InMemoryTokenStore) RemoveToken(resourceURL string) error {
-	delete(s.tokens, resourceURL)
+	s.tokens.Delete(resourceURL)
 	return nil
 }