ollie-9p/main.go

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

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

	"9fans.net/go/plan9/client"
	"ollie/agent"
	"ollie/backend"
	"ollie/elevate"
	"ollie/env"
	olog "ollie/log"
	"ollie/paths"
	"ollie/toolsrv"
	"ollie/skills"
	fs "olliesrv/fs"
	"olliesrv/mount"
	"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 shell")
var nodbus = flag.Bool("nodbus", false, "disable D-Bus adapter (for headless/container use)")
var no9p = flag.Bool("no9p", false, "disable 9P listener (D-Bus only mode)")

// 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() {
	flag.CommandLine.Parse(os.Args[1:]) //nolint:errcheck

	// Handle mount subcommand (kept as a convenience)
	if flag.NArg() > 0 && flag.Arg(0) == "mount" {
		cmdMount()
		return
	}

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

	sockPath := filepath.Join(ns, serviceName)
	runServer(sockPath)
}

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"
	}
	fmt.Printf("mounting %s at %s\n", addr, mnt)
	if err := mount.Run(mount.Config{
		Network:    network,
		Address:    addr,
		Mountpoint: mnt,
	}); err != nil {
		fmt.Fprintf(os.Stderr, "mount: %v\n", err)
		os.Exit(1)
	}
}

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

	// -no9p requires D-Bus (otherwise there's no interface at all)
	if *no9p && *nodbus {
		fmt.Fprintln(os.Stderr, "error: -no9p and -nodbus are mutually exclusive")
		os.Exit(1)
	}

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

	// Write PID file

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

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

	// Create the tool registry
	toolRegistry, regErr := toolsrv.NewRegistry()
	if regErr != nil {
		fmt.Fprintf(os.Stderr, "warning: tool registry: %v\n", regErr)
	}

	// Create the skills registry
	skillsRegistry := skills.NewRegistry()

	// D-Bus adapter (initialized after manager so callbacks can reference it).
	var dbusAdapter *DBusAdapter

	// Elevate broker (initialized after manager; closures capture the pointer).
	var elevateBroker *elevate.Broker

	daemonCtx, daemonCancel := context.WithCancel(context.Background())
	defer daemonCancel()

	sessionTree := session.NewRoot(session.Config{
		Ctx:              daemonCtx,
		ToolRegistry:     toolRegistry,
		SkillsRegistry:  skillsRegistry,
		AgentsDir:        agentsDirs[0],
		SessionsDir:      sessionsDir,
		Log:              sink.NewLogger("9p"),
		Sink:             sink,
		Strict:           *strict,
		Yolo:             *yolo,
		NoMount:          *no9p || *tcpAddr != "",
		Enable9P:         !*no9p,
		EnableDBus:       !*nodbus && (*no9p || *tcpAddr == ""),
		InvalidateModels: modelCache.Invalidate,
		OnSessionCreated: func(id string, sess *session.Session) {
			if dbusAdapter != nil {
				dbusAdapter.OnSessionCreated(id, sess)
			}
		},
		OnSessionKilled: func(id string) {
			if dbusAdapter != nil {
				dbusAdapter.OnSessionKilled(id)
			}
		},
		OnSessionRenamed: func(oldID, newID string) {
			if dbusAdapter != nil {
				dbusAdapter.OnSessionRenamed(oldID, newID)
			}
		},
	})

	var srv *server.Server
	// Server creation deferred until after elevate broker is ready (see below)

	// Start D-Bus adapter.
	// -tcp implies -nodbus (remote/headless deployments have no session bus).
	// -no9p forces D-Bus on (it's the only interface).
	if !*nodbus && (*no9p || *tcpAddr == "") {
		dbusAdapter = startDBus(sessionTree)
		if dbusAdapter != nil {
			dbusAdapter.WatchExisting()
		}
		if *no9p && dbusAdapter == nil {
			fmt.Fprintln(os.Stderr, "error: -no9p mode requires D-Bus but session bus is unavailable")
			os.Exit(1)
		}
	}

	// Start elevation broker
	elevateSocket := os.Getenv("OLLIE_ELEVATE_SOCKET")
	if elevateSocket == "" {
		xdg := os.Getenv("XDG_RUNTIME_DIR")
		if xdg == "" {
			xdg = fmt.Sprintf("/run/user/%d", os.Getuid())
		}
		elevateSocket = filepath.Join(xdg, "ollie", "elevate.sock")
	}
	policyPath := filepath.Join(paths.DataDir(), "elevate-policy.yaml")

	{
		notifyFn := func(req *elevate.Request) {
			if dbusAdapter != nil {
				dbusAdapter.NotifyElevation(req)
			}
		}
		var err error
		elevateBroker, err = elevate.NewBroker(elevate.BrokerConfig{
			SocketPath:   elevateSocket,
			PolicyPath:   policyPath,
			Notify:       notifyFn,
			Logf:         func(f string, a ...any) { fmt.Printf(f+"\n", a...) },
			SessionValid: func(id string) bool { return session.Lookup(sessionTree, id) != nil },
		})
		if err != nil {
			fmt.Fprintf(os.Stderr, "warning: %v\n", err)
		} else {
			defer elevateBroker.Close()
			if dbusAdapter != nil {
				initElevateNotifier(dbusAdapter.conn, elevateBroker)
			}
		}

	}

	// Create 9P server (after broker is ready so it can be passed in)
	if !*no9p {
		srv = server.New(server.Config{
			Sink:             sink,
			SessionTree:      sessionTree,
			RootStore:        NewRootStore(toolRegistry),
			ElevateBroker:    elevateBroker,
			InvalidateModels: modelCache.Invalidate,
			ToolPrompt:       toolsrv.ToolPrompt,
		})
	}

	// Prefetch models asynchronously at startup
	go modelCache.Refresh()

	var listener net.Listener
	if !*no9p {
		var err error
		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 !*no9p && *tcpAddr != "" {
		var err error
		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)
	}

	if *no9p {
		fmt.Println("olliesrv running in D-Bus only mode")
	}

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

	fmt.Println("shutting down")
	daemonCancel() // signal all sessions via context propagation
	if dbusAdapter != nil {
		dbusAdapter.Close()
	}
	if srv != nil {
		srv.Shutdown()
	}
	if listener != nil {
		listener.Close() //nolint:errcheck
	}
	if tcpListener != nil {
		tcpListener.Close() //nolint:errcheck
	}
	if !*no9p {
		os.Remove(sockPath)
	}
	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(reg *toolsrv.Registry) *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" || name == "generate" || name == "route" {
				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.DirEntry("elevate", 0755),
				fs.FileEntry("generate", 0666),
				fs.FileEntry("models", fs.Perms[fs.PathRoot].FileMode),
				fs.FileEntry("help", fs.Perms[fs.PathRoot].FileMode),
				fs.FileEntry("route", 0666),
			}, 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 == "generate" {
				return &fs.FileConfig{
					StatFn:         func() (os.FileInfo, error) { return &fs.SyntheticFileInfo{Name_: "generate", Mode_: 0666}, nil },
					ReadFn:         func() ([]byte, error) { return nil, nil },
					WriteFn:        func([]byte) error { return nil },
					BlockingReadFn: notBlocking,
				}, nil
			}
			if name == "route" {
				return &fs.FileConfig{
					StatFn:         func() (os.FileInfo, error) { return &fs.SyntheticFileInfo{Name_: "route", 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
		}),
	)
}