161 lines
3.1 KiB
Go
161 lines
3.1 KiB
Go
package tools
|
|
|
|
import (
|
|
"fmt"
|
|
"strings"
|
|
"sort"
|
|
"os"
|
|
"path/filepath"
|
|
"sync"
|
|
|
|
|
|
)
|
|
|
|
type Registry struct {
|
|
mu sync.RWMutex
|
|
global map[string]ToolInfo
|
|
sessions map[string]map[string]ToolInfo
|
|
revisions map[string]uint64
|
|
}
|
|
|
|
func NewRegistry() (*Registry, error) {
|
|
r := &Registry{
|
|
global: make(map[string]ToolInfo),
|
|
sessions: make(map[string]map[string]ToolInfo),
|
|
revisions: make(map[string]uint64),
|
|
}
|
|
if err := r.Discover(); err != nil {
|
|
return nil, err
|
|
}
|
|
return r, nil
|
|
}
|
|
|
|
func (r *Registry) Discover() error {
|
|
dir := ToolsPath()
|
|
entries, err := os.ReadDir(dir)
|
|
if err != nil {
|
|
return fmt.Errorf("read tools dir %s: %w", dir, err)
|
|
}
|
|
global := make(map[string]ToolInfo)
|
|
for _, e := range entries {
|
|
if e.IsDir() || e.Name() == "idx" || e.Name()[0] == '.' {
|
|
continue
|
|
}
|
|
data, err := os.ReadFile(filepath.Join(dir, e.Name()))
|
|
if err != nil {
|
|
continue
|
|
}
|
|
script := string(data)
|
|
if !strings.Contains(script, "ollie:prompt") {
|
|
continue
|
|
}
|
|
global[e.Name()] = ParseToolInfo(e.Name(), script)
|
|
}
|
|
r.mu.Lock()
|
|
r.global = global
|
|
r.mu.Unlock()
|
|
return nil
|
|
}
|
|
|
|
func (r *Registry) Summaries() []ToolInfo {
|
|
r.mu.RLock()
|
|
defer r.mu.RUnlock()
|
|
var summaries []ToolInfo
|
|
for _, info := range r.global {
|
|
summaries = append(summaries, ToolInfo{
|
|
Name: info.Name,
|
|
Description: info.Description,
|
|
})
|
|
}
|
|
sort.Slice(summaries, func(i, j int) bool { return summaries[i].Name < summaries[j].Name })
|
|
return summaries
|
|
}
|
|
|
|
func (r *Registry) Load(sessionID, name string) error {
|
|
r.mu.RLock()
|
|
tool, exists := r.global[name]
|
|
r.mu.RUnlock()
|
|
if !exists {
|
|
return fmt.Errorf("tool not found: %s", name)
|
|
}
|
|
|
|
r.mu.Lock()
|
|
defer r.mu.Unlock()
|
|
|
|
if r.sessions[sessionID] == nil {
|
|
r.sessions[sessionID] = make(map[string]ToolInfo)
|
|
}
|
|
|
|
if _, already := r.sessions[sessionID][name]; already {
|
|
return nil
|
|
}
|
|
|
|
r.sessions[sessionID][name] = tool
|
|
r.revisions[sessionID]++
|
|
return nil
|
|
}
|
|
|
|
func (r *Registry) Unload(sessionID, name string) error {
|
|
r.mu.Lock()
|
|
defer r.mu.Unlock()
|
|
|
|
sessionTools, ok := r.sessions[sessionID]
|
|
if !ok {
|
|
return nil
|
|
}
|
|
|
|
if _, exists := sessionTools[name]; !exists {
|
|
return nil
|
|
}
|
|
|
|
delete(sessionTools, name)
|
|
r.revisions[sessionID]++
|
|
return nil
|
|
}
|
|
|
|
func (r *Registry) Loaded(sessionID string) []ToolInfo {
|
|
r.mu.RLock()
|
|
defer r.mu.RUnlock()
|
|
|
|
sessionTools, ok := r.sessions[sessionID]
|
|
if !ok {
|
|
return nil
|
|
}
|
|
|
|
var loaded []ToolInfo
|
|
for _, info := range sessionTools {
|
|
loaded = append(loaded, info)
|
|
}
|
|
sort.Slice(loaded, func(i, j int) bool { return loaded[i].Name < loaded[j].Name })
|
|
return loaded
|
|
}
|
|
|
|
func (r *Registry) Lookup(sessionID, name string) (ToolInfo, bool) {
|
|
r.mu.RLock()
|
|
defer r.mu.RUnlock()
|
|
|
|
sessionTools, ok := r.sessions[sessionID]
|
|
if !ok {
|
|
return ToolInfo{}, false
|
|
}
|
|
|
|
tool, exists := sessionTools[name]
|
|
return tool, exists
|
|
}
|
|
|
|
func (r *Registry) Revision(sessionID string) uint64 {
|
|
r.mu.RLock()
|
|
defer r.mu.RUnlock()
|
|
return r.revisions[sessionID]
|
|
}
|
|
|
|
func (r *Registry) GlobalToolNames() []string {
|
|
r.mu.RLock()
|
|
defer r.mu.RUnlock()
|
|
names := make([]string, 0, len(r.global))
|
|
for name := range r.global {
|
|
names = append(names, name)
|
|
}
|
|
return names
|
|
}
|