140 lines
3.0 KiB
Go
140 lines
3.0 KiB
Go
// Package registry provides the tool registry for toolsrvclient.
|
|
// Tools are scoped per-agent within a session (one toolsrv per session).
|
|
package registry
|
|
|
|
import (
|
|
"fmt"
|
|
"sort"
|
|
"sync"
|
|
|
|
"ollie/toolsrv/metadata"
|
|
"ollie/toolsrv/protocol"
|
|
)
|
|
|
|
// Registry tracks loaded tools per agent.
|
|
type Registry struct {
|
|
mu sync.RWMutex
|
|
agents map[string]map[string]protocol.ToolInfo
|
|
revisions map[string]uint64
|
|
}
|
|
|
|
// New creates a new tool registry.
|
|
func New() (*Registry, error) {
|
|
return &Registry{
|
|
agents: make(map[string]map[string]protocol.ToolInfo),
|
|
revisions: make(map[string]uint64),
|
|
}, nil
|
|
}
|
|
|
|
// discover scans the tools directory and returns all valid tool infos.
|
|
func (r *Registry) discover() map[string]protocol.ToolInfo {
|
|
infos := metadata.DiscoverTools()
|
|
tools := make(map[string]protocol.ToolInfo, len(infos))
|
|
for _, info := range infos {
|
|
tools[info.Name] = info
|
|
}
|
|
return tools
|
|
}
|
|
|
|
// Load loads a tool for the given agent.
|
|
func (r *Registry) Load(agentID, name string) error {
|
|
tools := r.discover()
|
|
tool, exists := tools[name]
|
|
if !exists {
|
|
return fmt.Errorf("tool not found: %s", name)
|
|
}
|
|
|
|
r.mu.Lock()
|
|
defer r.mu.Unlock()
|
|
|
|
if r.agents[agentID] == nil {
|
|
r.agents[agentID] = make(map[string]protocol.ToolInfo)
|
|
}
|
|
|
|
r.agents[agentID][name] = tool
|
|
r.revisions[agentID]++
|
|
return nil
|
|
}
|
|
|
|
// Unload removes a tool from an agent's registry.
|
|
func (r *Registry) Unload(agentID, name string) error {
|
|
r.mu.Lock()
|
|
defer r.mu.Unlock()
|
|
|
|
agentTools, ok := r.agents[agentID]
|
|
if !ok {
|
|
return nil
|
|
}
|
|
|
|
if _, exists := agentTools[name]; !exists {
|
|
return nil
|
|
}
|
|
|
|
delete(agentTools, name)
|
|
if len(agentTools) == 0 {
|
|
delete(r.agents, agentID)
|
|
delete(r.revisions, agentID)
|
|
return nil
|
|
}
|
|
r.revisions[agentID]++
|
|
return nil
|
|
}
|
|
|
|
// ClearAgent removes all tools for an agent.
|
|
func (r *Registry) ClearAgent(agentID string) {
|
|
r.mu.Lock()
|
|
defer r.mu.Unlock()
|
|
delete(r.agents, agentID)
|
|
delete(r.revisions, agentID)
|
|
}
|
|
|
|
// Loaded returns all loaded tools for an agent.
|
|
func (r *Registry) Loaded(agentID string) []protocol.ToolInfo {
|
|
r.mu.RLock()
|
|
defer r.mu.RUnlock()
|
|
|
|
agentTools, ok := r.agents[agentID]
|
|
if !ok {
|
|
return nil
|
|
}
|
|
|
|
var loaded []protocol.ToolInfo
|
|
for _, info := range agentTools {
|
|
loaded = append(loaded, info)
|
|
}
|
|
sort.Slice(loaded, func(i, j int) bool { return loaded[i].Name < loaded[j].Name })
|
|
return loaded
|
|
}
|
|
|
|
// Lookup returns a tool by name for an agent.
|
|
func (r *Registry) Lookup(agentID, name string) (protocol.ToolInfo, bool) {
|
|
r.mu.RLock()
|
|
defer r.mu.RUnlock()
|
|
|
|
agentTools, ok := r.agents[agentID]
|
|
if !ok {
|
|
return protocol.ToolInfo{}, false
|
|
}
|
|
|
|
tool, exists := agentTools[name]
|
|
return tool, exists
|
|
}
|
|
|
|
// Revision returns the revision counter for an agent's tools.
|
|
func (r *Registry) Revision(agentID string) uint64 {
|
|
r.mu.RLock()
|
|
defer r.mu.RUnlock()
|
|
return r.revisions[agentID]
|
|
}
|
|
|
|
// All returns all available tool names (from the tools directory).
|
|
func (r *Registry) All() []string {
|
|
tools := r.discover()
|
|
names := make([]string, 0, len(tools))
|
|
for name := range tools {
|
|
names = append(names, name)
|
|
}
|
|
sort.Strings(names)
|
|
return names
|
|
}
|