ollie-9p/main.go

// olliesrv - 9P server for ollie sessions
package main

import (
	"context"
	"flag"
	"fmt"
	"net"
	"os"
	"os/exec"
	"os/signal"
	"path/filepath"
	"strings"
	"sync"
	"syscall"

	"9fans.net/go/plan9/client"
	"ollie/pkg/agent"
	"ollie/pkg/backend"
	"ollie/pkg/env"
	olog "ollie/pkg/log"
	"ollie/pkg/paths"
	fs "olliesrv/fs"
	"olliesrv/server"
	"olliesrv/session"
)

const serviceName = "ollie"

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 <start|fgstart|stop|status|mount>")
		os.Exit(1)
	}
	subcmd := os.Args[1]
	flag.CommandLine.Parse(os.Args[2:]) //nolint:errcheck

	switch subcmd {
	case "mount":
		cmdMount()
		return
	}

	ns := client.Namespace()
	if ns == "" {
		fmt.Fprintln(os.Stderr, "no namespace")
		os.Exit(1)
	}

	sockPath := filepath.Join(ns, serviceName)
	pidPath := filepath.Join(ns, serviceName+".pid")

	switch subcmd {
	case "start":
		if isRunning(sockPath) {
			fmt.Println("olliesrv already running")
			os.Exit(0)
		}
		daemonize(pidPath)
	case "fgstart":
		if isRunning(sockPath) {
			fmt.Println("olliesrv already running")
			os.Exit(0)
		}
		runServer(sockPath, pidPath)
	case "stop":
		stopServer(sockPath, pidPath)
	case "status":
		if isRunning(sockPath) {
			fmt.Println("olliesrv running")
		} else {
			fmt.Println("olliesrv not running")
			os.Exit(1)
		}
	default:
		fmt.Fprintln(os.Stderr, "usage: olliesrv <start|fgstart|stop|status|mount>")
		os.Exit(1)
	}
}

func cmdMount() {
	if flag.NArg() < 1 {
		fmt.Fprintln(os.Stderr, "usage: olliesrv mount <address> [mountpoint]")
		os.Exit(1)
	}
	addr := flag.Arg(0)
	mnt := flag.Arg(1)
	if mnt == "" {
		home, _ := os.UserHomeDir()
		mnt = filepath.Join(home, "mnt", addr)
	}
	if err := os.MkdirAll(mnt, fs.PermMkdir); err != nil {
		fmt.Fprintf(os.Stderr, "cannot create mount dir: %v\n", err)
		os.Exit(1)
	}
	network := "unix"
	if strings.Contains(addr, ":") {
		network = "tcp"
	}
	exe, _ := os.Executable()
	cmd := exec.Command(filepath.Join(filepath.Dir(exe), "ollie-9p-mount"),
		"-u", os.Getenv("USER"), "-net", network, "-addr", addr, mnt)
	cmd.SysProcAttr = &syscall.SysProcAttr{Setsid: true}
	if err := cmd.Start(); err != nil {
		fmt.Fprintf(os.Stderr, "ollie-9p-mount: %v\n", err)
		os.Exit(1)
	}
	fmt.Printf("mounted %s at %s (pid %d)\n", addr, mnt, cmd.Process.Pid)
}

func isRunning(sockPath string) bool {
	conn, err := net.Dial("unix", sockPath)
	if err == nil {
		conn.Close()
		return true
	}
	return false
}

func daemonize(pidPath string) {
	exe, _ := os.Executable()
	args := []string{"fgstart"}
	if *tcpAddr != "" {
		args = append(args, "-tcp", *tcpAddr)
	}
	if *strict {
		args = append(args, "-strict")
	}
	if *yolo {
		args = append(args, "-yolo")
	}
	cmd := exec.Command(exe, args...)
	cmd.SysProcAttr = &syscall.SysProcAttr{Setsid: true}
	if err := cmd.Start(); err != nil {
		fmt.Fprintf(os.Stderr, "failed to start: %v\n", err)
		os.Exit(1)
	}
	fmt.Printf("olliesrv started (pid %d)\n", cmd.Process.Pid)
}

func stopServer(sockPath, pidPath string) {
	data, err := os.ReadFile(pidPath)
	if err != nil {
		fmt.Println("olliesrv not running")
		return
	}
	var pid int
	fmt.Sscanf(string(data), "%d", &pid)
	if pid > 0 {
		syscall.Kill(pid, syscall.SIGTERM) //nolint:errcheck
	}
	os.Remove(sockPath) //nolint:errcheck
	os.Remove(pidPath)  //nolint:errcheck
	fmt.Println("olliesrv stopped")
}

func runServer(sockPath, pidPath string) {
	env.EnsureDefaults()

	// Remove stale socket
	if _, err := os.Stat(sockPath); err == nil {
		os.Remove(sockPath) //nolint:errcheck
	}

	// Write PID file
	os.WriteFile(pidPath, []byte(fmt.Sprintf("%d", os.Getpid())), fs.Perms[fs.PathRoot].FileMode) //nolint:errcheck

	sink := olog.NewSink(os.Stdout, os.Stderr, olog.ParseLevel(os.Getenv("OLLIE_LOG"), olog.LevelWarn))

	agentsDirs := agent.AgentsDirs()
	sessionsDir := paths.CfgDir() + "/sessions"

	mgr := session.NewManager(session.ManagerConfig{
		AgentsDir:   agentsDirs[0],
		SessionsDir: sessionsDir,
		Log:         sink.NewLogger("9p"),
		Sink:        sink,
		Strict:      *strict,
		Yolo:        *yolo,
		NoMount:     *tcpAddr != "",
		InvalidateModels: modelCache.Invalidate,
	})

	srv := server.New(server.Config{
		Sink:            sink,
		SessionMgr:      mgr,
		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)
		os.Exit(1)
	}

	go func() {
		for {
			conn, err := listener.Accept()
			if err != nil {
				return
			}
			go srv.Serve(conn)
		}
	}()

	fmt.Printf("olliesrv listening on %s\n", sockPath)

	var tcpListener net.Listener
	if *tcpAddr != "" {
		tcpListener, err = net.Listen("tcp", *tcpAddr)
		if err != nil {
			fmt.Fprintf(os.Stderr, "tcp listen: %v\n", err)
			os.Exit(1)
		}
		go func() {
			for {
				conn, err := tcpListener.Accept()
				if err != nil {
					return
				}
				go srv.Serve(conn)
			}
		}()
		fmt.Printf("olliesrv listening on tcp %s\n", *tcpAddr)
	}

	sigChan := make(chan os.Signal, 1)
	signal.Notify(sigChan, syscall.SIGINT, syscall.SIGTERM)
	<-sigChan

	srv.InterruptAll()
	fmt.Println("shutting down")
	srv.Shutdown()
	listener.Close() //nolint:errcheck
	if tcpListener != nil {
		tcpListener.Close() //nolint:errcheck
	}
	os.Remove(sockPath)
	os.Remove(pidPath)
	sink.Flush()
}


// ToolIndex generates a tool index from a tree's file listing.
// NewRootStore returns a read-only FileTree for synthetic root-level files.
func NewRootStore() *fs.Tree {
	helpPath := paths.CfgDir() + "/help.md"
	readEntries := map[string]func() ([]byte, error){
		"backends": func() ([]byte, error) {
			return []byte(strings.Join(backend.Backends(), "\n") + "\n"), nil
		},
		"models": func() ([]byte, error) {
			return modelCache.Get(), nil
		},
		"help": func() ([]byte, error) {
			return os.ReadFile(helpPath)
		},
		"agents": func() ([]byte, error) {
			var sb strings.Builder
			for _, dir := range agent.AgentsDirs() {
				entries, err := os.ReadDir(dir)
				if err != nil {
					continue
				}
				for _, e := range entries {
					if !e.IsDir() && strings.HasSuffix(e.Name(), ".json") {
						sb.WriteString(strings.TrimSuffix(e.Name(), ".json"))
						sb.WriteByte('\n')
					}
				}
			}
			return []byte(sb.String()), nil
		},
	}
	notBlocking := func(context.Context, string) ([]byte, string, error) {
		return nil, "", fmt.Errorf("blocking read not supported")
	}
	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.WithResolver(func(_ []string, name string) (string, error) {
			if _, ok := readEntries[name]; ok {
				return name, nil
			}
			if name == "ctl" || name == "complete" {
				return name, nil
			}
			return "", fmt.Errorf("%s: not found", name)
		}),
		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) {
			if name == "complete" {
				return &fs.FileConfig{
					StatFn:         func() (os.FileInfo, error) { return &fs.SyntheticFileInfo{Name_: "complete", Mode_: 0666}, nil },
					ReadFn:         func() ([]byte, error) { return nil, nil },
					WriteFn:        func([]byte) error { return nil },
					BlockingReadFn: notBlocking,
				}, nil
			}
			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) {
					data, _ := readFn()
					return &fs.SyntheticFileInfo{Name_: name, Mode_: fs.Perms[fs.PathRoot].FileMode, Size_: int64(len(data))}, nil
				},
				ReadFn:         readFn,
				WriteFn:        readOnly,
				BlockingReadFn: notBlocking,
			}, nil
		}),
	)
}