ollie-core/toolsrv/registry.go

package toolsrv

import (
	"fmt"
	"os"
	"path/filepath"
	"sort"
	"strings"
	"sync"
)

type Registry struct {
	mu        sync.RWMutex
	sessions  map[string]map[string]ToolInfo
	revisions map[string]uint64
}

func NewRegistry() (*Registry, error) {
	r := &Registry{
		sessions:  make(map[string]map[string]ToolInfo),
		revisions: make(map[string]uint64),
	}
	// Validate that the tools directory is readable.
	dir := ToolsPath()
	if _, err := os.ReadDir(dir); err != nil {
		return nil, fmt.Errorf("read tools dir %s: %w", dir, err)
	}
	return r, nil
}

// discover scans the tools directory and returns all valid tool infos.
func (r *Registry) discover() map[string]ToolInfo {
	dir := ToolsPath()
	entries, err := os.ReadDir(dir)
	if err != nil {
		return nil
	}
	tools := 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
		}
		tools[e.Name()] = ParseToolInfo(e.Name(), script)
	}
	return tools
}

func (r *Registry) Summaries() []ToolInfo {
	tools := r.discover()
	var summaries []ToolInfo
	for _, info := range tools {
		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 {
	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.sessions[sessionID] == nil {
		r.sessions[sessionID] = make(map[string]ToolInfo)
	}

	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
}

// RefreshLoaded re-reads all loaded tools from disk, updating their schemas
// and metadata in place. Called from tool_list to keep schemas current.
func (r *Registry) RefreshLoaded(sessionID string) {
	r.mu.RLock()
	sessionTools, ok := r.sessions[sessionID]
	if !ok || len(sessionTools) == 0 {
		r.mu.RUnlock()
		return
	}
	names := make([]string, 0, len(sessionTools))
	for name := range sessionTools {
		names = append(names, name)
	}
	r.mu.RUnlock()

	fresh := r.discover()

	r.mu.Lock()
	defer r.mu.Unlock()
	changed := false
	for _, name := range names {
		if tool, exists := fresh[name]; exists {
			r.sessions[sessionID][name] = tool
			changed = true
		} else {
			// Tool was removed from disk — unload it.
			delete(r.sessions[sessionID], name)
			changed = true
		}
	}
	if changed {
		r.revisions[sessionID]++
	}
}

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]
}