62 lines
1.5 KiB
Go
62 lines
1.5 KiB
Go
package model
|
|
|
|
import "sync"
|
|
|
|
// InMemoryAgentRegistry 基于内存的 Agent 注册表
|
|
type InMemoryAgentRegistry struct {
|
|
mu sync.RWMutex
|
|
agents map[string]RegisteredAgent
|
|
}
|
|
|
|
func NewInMemoryAgentRegistry() *InMemoryAgentRegistry {
|
|
return &InMemoryAgentRegistry{agents: make(map[string]RegisteredAgent)}
|
|
}
|
|
|
|
func (r *InMemoryAgentRegistry) Register(agent RegisteredAgent) error {
|
|
r.mu.Lock()
|
|
defer r.mu.Unlock()
|
|
r.agents[agent.AgentID] = agent
|
|
return nil
|
|
}
|
|
|
|
func (r *InMemoryAgentRegistry) Get(agentID string) (RegisteredAgent, bool) {
|
|
r.mu.RLock()
|
|
defer r.mu.RUnlock()
|
|
agent, ok := r.agents[agentID]
|
|
return agent, ok
|
|
}
|
|
|
|
func (r *InMemoryAgentRegistry) List() []RegisteredAgent {
|
|
r.mu.RLock()
|
|
defer r.mu.RUnlock()
|
|
agents := make([]RegisteredAgent, 0, len(r.agents))
|
|
for _, agent := range r.agents {
|
|
agents = append(agents, agent)
|
|
}
|
|
return agents
|
|
}
|
|
|
|
// InMemorySessionStore 基于内存的会话存储
|
|
type InMemorySessionStore struct {
|
|
mu sync.RWMutex
|
|
sessions map[string]string
|
|
}
|
|
|
|
func NewInMemorySessionStore() *InMemorySessionStore {
|
|
return &InMemorySessionStore{sessions: make(map[string]string)}
|
|
}
|
|
|
|
func (s *InMemorySessionStore) Get(userID, agentID string) (string, bool) {
|
|
s.mu.RLock()
|
|
defer s.mu.RUnlock()
|
|
sessionID, ok := s.sessions[userID+":"+agentID]
|
|
return sessionID, ok
|
|
}
|
|
|
|
func (s *InMemorySessionStore) Set(userID, agentID, sessionID string) error {
|
|
s.mu.Lock()
|
|
defer s.mu.Unlock()
|
|
s.sessions[userID+":"+agentID] = sessionID
|
|
return nil
|
|
}
|