diff --git a/main.go b/main.go index 746d19b..2fa376d 100644 --- a/main.go +++ b/main.go @@ -11,6 +11,7 @@ import ( "os/signal" "path/filepath" "strings" + "sync" "syscall" "time" @@ -33,6 +34,58 @@ var tcpAddr = flag.String("tcp", "", "also listen on TCP address (e.g. :564)") var strict = flag.Bool("strict", false, "only allow tool steps; reject inline code") var yolo = flag.Bool("yolo", false, "skip landrun sandbox for execute_code") +// ModelCache caches the list of models from all backends. +type ModelCache struct { + mu sync.RWMutex + data []byte + fetched bool +} + +var modelCache = &ModelCache{} + +// Get returns the cached models, fetching if not yet cached. +func (c *ModelCache) Get() []byte { + c.mu.RLock() + if c.fetched { + data := c.data + c.mu.RUnlock() + return data + } + c.mu.RUnlock() + return c.Refresh() +} + +// Refresh fetches models from all backends and updates the cache. +func (c *ModelCache) Refresh() []byte { + var sb strings.Builder + for _, name := range backend.Backends() { + b, err := backend.NewWithName(name) + if err != nil { + continue + } + for _, m := range b.Models(context.Background()) { + sb.WriteString(name) + sb.WriteByte('\t') + sb.WriteString(m) + sb.WriteByte('\n') + } + } + data := []byte(sb.String()) + c.mu.Lock() + c.data = data + c.fetched = true + c.mu.Unlock() + return data +} + +// Invalidate clears the cache, forcing a refresh on next Get. +func (c *ModelCache) Invalidate() { + c.mu.Lock() + c.fetched = false + c.data = nil + c.mu.Unlock() +} + func main() { if len(os.Args) < 2 { fmt.Fprintln(os.Stderr, "usage: olliesrv ") @@ -191,6 +244,7 @@ func runServer(sockPath, pidPath string) { Strict: *strict, Yolo: *yolo, ToolTree: toolTree, + InvalidateModels: modelCache.Invalidate, SaveTranscript: func(data []byte) error { name := time.Now().Format("20060102T150405") + "-chat.md" e, err := transcriptStore.Open(name) @@ -218,8 +272,12 @@ func runServer(sockPath, pidPath string) { TranscriptStore: transcriptStore, TmpStore: fs.NewTree([]string{tmpDir}, fs.Perms[fs.PathTmp].FileMode), RootStore: NewRootStore(), + InvalidateModels: modelCache.Invalidate, }) + // Prefetch models asynchronously at startup + go modelCache.Refresh() + listener, err := net.Listen("unix", sockPath) if err != nil { fmt.Fprintf(os.Stderr, "listen: %v\n", err) @@ -368,25 +426,12 @@ func ToolIndex(s fs.FileTree) ([]byte, error) { // NewRootStore returns a read-only FileTree for synthetic root-level files. func NewRootStore() *fs.Tree { helpPath := paths.CfgDir() + "/help.md" - entries := map[string]func() ([]byte, error){ + readEntries := map[string]func() ([]byte, error){ "backends": func() ([]byte, error) { return []byte(strings.Join(backend.Backends(), "\n") + "\n"), nil }, "models": func() ([]byte, error) { - var sb strings.Builder - for _, name := range backend.Backends() { - b, err := backend.NewWithName(name) - if err != nil { - continue - } - for _, m := range b.Models(context.Background()) { - sb.WriteString(name) - sb.WriteByte('\t') - sb.WriteString(m) - sb.WriteByte('\n') - } - } - return []byte(sb.String()), nil + return modelCache.Get(), nil }, "help": func() ([]byte, error) { return os.ReadFile(helpPath) @@ -397,10 +442,24 @@ func NewRootStore() *fs.Tree { } readOnly := func([]byte) error { return fmt.Errorf("read-only") } + // ctl commands + ctlHandler := func(data []byte) error { + cmd := strings.TrimSpace(string(data)) + switch cmd { + case "invalidate": + modelCache.Invalidate() + return nil + default: + return fmt.Errorf("unknown ctl command: %s", cmd) + } + } + return fs.NewTree(nil, fs.Perms[fs.PathRoot].FileMode, - fs.WithReadOnly(), fs.WithResolver(func(_ []string, name string) (string, error) { - if _, ok := entries[name]; ok { + if _, ok := readEntries[name]; ok { + return name, nil + } + if name == "ctl" { return name, nil } return "", fmt.Errorf("%s: not found", name) @@ -408,17 +467,29 @@ func NewRootStore() *fs.Tree { fs.WithLister(func(_ []string) ([]os.DirEntry, error) { return []os.DirEntry{ fs.FileEntry("backends", fs.Perms[fs.PathRoot].FileMode), + fs.FileEntry("ctl", 0666), fs.FileEntry("models", fs.Perms[fs.PathRoot].FileMode), fs.FileEntry("help", fs.Perms[fs.PathRoot].FileMode), }, nil }), fs.WithOpener(func(_ []string, name string) (fs.File, error) { - readFn, ok := entries[name] + if name == "ctl" { + return &fs.FileConfig{ + StatFn: func() (os.FileInfo, error) { return &fs.SyntheticFileInfo{Name_: "ctl", Mode_: 0666}, nil }, + ReadFn: func() ([]byte, error) { return nil, nil }, + WriteFn: ctlHandler, + BlockingReadFn: notBlocking, + }, nil + } + readFn, ok := readEntries[name] if !ok { return nil, fmt.Errorf("%s: not found", name) } return &fs.FileConfig{ - StatFn: func() (os.FileInfo, error) { return &fs.SyntheticFileInfo{Name_: name, Mode_: fs.Perms[fs.PathRoot].FileMode}, nil }, + StatFn: func() (os.FileInfo, error) { + data, _ := readFn() + return &fs.SyntheticFileInfo{Name_: name, Mode_: fs.Perms[fs.PathRoot].FileMode, Size_: int64(len(data))}, nil + }, ReadFn: readFn, WriteFn: readOnly, BlockingReadFn: notBlocking, diff --git a/server/server.go b/server/server.go index 5571b6d..c7de354 100644 --- a/server/server.go +++ b/server/server.go @@ -57,38 +57,40 @@ type connState struct { // Server is the 9P server for ollie sessions. type Server struct { - mu sync.RWMutex - conns []*connState - log *olog.Logger - sink *olog.Sink - agentStore FileTree - promptStore FileTree - memStore FileTree - toolTree *fs.Tree - utilStore FileTree - pluginStore FileTree - skillTree FileTree - sessionMgr *session.Manager - transcriptStore FileTree - tmpStore FileTree - rootStore *fs.Tree - groups map[string]map[string]bool // group → set of members + mu sync.RWMutex + conns []*connState + log *olog.Logger + sink *olog.Sink + agentStore FileTree + promptStore FileTree + memStore FileTree + toolTree *fs.Tree + utilStore FileTree + pluginStore FileTree + skillTree FileTree + sessionMgr *session.Manager + transcriptStore FileTree + tmpStore FileTree + rootStore *fs.Tree + groups map[string]map[string]bool // group → set of members + invalidateModels func() } // Config holds the pre-built trees and manager for the server. type Config struct { - Sink *olog.Sink - AgentStore FileTree - PromptStore FileTree - MemStore FileTree - ToolTree *fs.Tree - UtilStore FileTree - PluginStore FileTree - SkillTree FileTree - SessionMgr *session.Manager - TranscriptStore FileTree - TmpStore FileTree - RootStore *fs.Tree + Sink *olog.Sink + AgentStore FileTree + PromptStore FileTree + MemStore FileTree + ToolTree *fs.Tree + UtilStore FileTree + PluginStore FileTree + SkillTree FileTree + SessionMgr *session.Manager + TranscriptStore FileTree + TmpStore FileTree + RootStore *fs.Tree + InvalidateModels func() } // Option configures the 9P server. @@ -97,20 +99,21 @@ type Option func(*Server) // New creates a new Server from a pre-built Config. func New(cfg Config) *Server { s := &Server{ - log: cfg.Sink.Logger("9p", olog.LevelDebug), - sink: cfg.Sink, - agentStore: cfg.AgentStore, - promptStore: cfg.PromptStore, - memStore: cfg.MemStore, - toolTree: cfg.ToolTree, - utilStore: cfg.UtilStore, - pluginStore: cfg.PluginStore, - skillTree: cfg.SkillTree, - sessionMgr: cfg.SessionMgr, - transcriptStore: cfg.TranscriptStore, - tmpStore: cfg.TmpStore, - rootStore: cfg.RootStore, - groups: make(map[string]map[string]bool), + log: cfg.Sink.Logger("9p", olog.LevelDebug), + sink: cfg.Sink, + agentStore: cfg.AgentStore, + promptStore: cfg.PromptStore, + memStore: cfg.MemStore, + toolTree: cfg.ToolTree, + utilStore: cfg.UtilStore, + pluginStore: cfg.PluginStore, + skillTree: cfg.SkillTree, + sessionMgr: cfg.SessionMgr, + transcriptStore: cfg.TranscriptStore, + tmpStore: cfg.TmpStore, + rootStore: cfg.RootStore, + groups: make(map[string]map[string]bool), + invalidateModels: cfg.InvalidateModels, } return s } @@ -217,6 +220,13 @@ func (s *Server) filePerm(path string) os.FileMode { } } } + // Root-level files: query the rootStore for actual permission + name := strings.TrimPrefix(path, "/") + if !strings.Contains(name, "/") { + if info, err := s.rootStore.Stat(name); err == nil { + return info.Mode() + } + } return fs.PermIdx } @@ -951,7 +961,7 @@ func (s *Server) readDir(path string, offset uint64, count uint32) []byte { var dirs []plan9.Dir if path == "/" { - rootEntries := []string{"a", "backends", "help", "m", "p", "s", "sk", "tmp", "u", "x", "tr"} + rootEntries := []string{"a", "backends", "ctl", "help", "m", "models", "p", "s", "sk", "tmp", "u", "x", "tr"} for _, name := range rootEntries { fpath := "/" + name st := s.makeStat(fpath) @@ -1088,7 +1098,9 @@ func (s *Server) makeStat(path string) plan9.Dir { case path == "/backends": dir.Length = uint64(len(strings.Join(backend.Backends(), "\n") + "\n")) case path == "/models": - // dynamic; length unknown until read + if content, err := readFile(s.rootStore, "models"); err == nil { + dir.Length = uint64(len(content)) + } case path == "/help": if info, err := os.Stat(s.helpPath()); err == nil { dir.Length = uint64(info.Size()) diff --git a/session/session.go b/session/session.go index 9796fd4..806d615 100644 --- a/session/session.go +++ b/session/session.go @@ -141,6 +141,8 @@ type ManagerConfig struct { Strict bool // Yolo skips the landrun sandbox. Yolo bool + // InvalidateModels clears the model cache, forcing a refresh. + InvalidateModels func() } // Manager manages session lifecycle and exposes sessions as a Tree. @@ -465,6 +467,7 @@ func (s *Manager) openStore(sess *Session) (*fs.Tree, error) { func() { s.KillSession(sess.id) }, func(newID string) error { return s.renameSession(sess.id, newID) }, s.cfg.SaveTranscript, + s.cfg.InvalidateModels, ), nil } diff --git a/session/sessionfile.go b/session/sessionfile.go index 63b87bf..5c65cb8 100644 --- a/session/sessionfile.go +++ b/session/sessionfile.go @@ -48,8 +48,8 @@ func sessionFilePerms() map[string]os.FileMode { } -func NewSessionTree(sess *Session, log *olog.Logger, kill func(), rename func(newID string) error, saveTranscript func([]byte) error) *fs.Tree { - h := &sessionHelper{sess: sess, log: log, kill: kill, rename: rename, saveTranscript: saveTranscript} +func NewSessionTree(sess *Session, log *olog.Logger, kill func(), rename func(newID string) error, saveTranscript func([]byte) error, invalidateModels func()) *fs.Tree { + h := &sessionHelper{sess: sess, log: log, kill: kill, rename: rename, saveTranscript: saveTranscript, invalidateModels: invalidateModels} perms := sessionFilePerms() specs := make([]fs.FileSpec, len(FileList)) for i, f := range FileList { @@ -62,11 +62,12 @@ func NewSessionTree(sess *Session, log *olog.Logger, kill func(), rename func(ne // sessionHelper holds the dependencies needed to build session FileSpecs. type sessionHelper struct { - sess *Session - log *olog.Logger - kill func() - rename func(newID string) error - saveTranscript func([]byte) error + sess *Session + log *olog.Logger + kill func() + rename func(newID string) error + saveTranscript func([]byte) error + invalidateModels func() } func (h *sessionHelper) fileSpec(name string, mode os.FileMode) fs.FileSpec { @@ -148,6 +149,13 @@ func (h *sessionHelper) fileSpec(name string, mode os.FileMode) fs.FileSpec { if input == "" { return nil } + // Handle /invalidate locally before passing to core + if input == "/invalidate" { + if h.invalidateModels != nil { + h.invalidateModels() + } + return nil + } h.sess.mu.Lock() h.sess.prevPrompt = []byte(input) h.sess.mu.Unlock() @@ -530,6 +538,10 @@ func (h *sessionHelper) handleCtl(input string) error { copy(data, h.sess.log) h.sess.mu.RUnlock() return h.saveTranscript(data) + case "invalidate": + if h.invalidateModels != nil { + h.invalidateModels() + } case "compact", "clear", "backend", "model", "models", "agents", "agent", "sessions", "cwd", "skills", "tools", "context", "usage", "cost", "history",