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"
	"syscall"
	"time"

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

const serviceName = "ollie"

var mountPath = flag.String("mount", "", "FUSE mount path (default: $HOME/mnt/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")

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, 0755); err != nil {
		fmt.Fprintf(os.Stderr, "cannot create mount dir: %v\n", err)
		os.Exit(1)
	}
	cmd := exec.Command("9pfuse", addr, mnt)
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	if err := cmd.Run(); err != nil {
		fmt.Fprintf(os.Stderr, "9pfuse: %v\n", err)
		os.Exit(1)
	}
	fmt.Printf("mounted %s at %s\n", addr, mnt)
}

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 *mountPath != "" {
		args = append(args, "-mount", *mountPath)
	}
	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())), 0644) //nolint:errcheck

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

	memDir := defaultMemDir()
	os.MkdirAll(memDir, 0755) //nolint:errcheck
	transcriptDir := defaultTranscriptDir()
	os.MkdirAll(transcriptDir, 0755) //nolint:errcheck
	tmpDir := defaultTmpDir()
	os.MkdirAll(tmpDir, 0755) //nolint:errcheck

	agentsDirs := agent.AgentsDirs()
	sessionsDir := paths.CfgDir() + "/sessions"
	toolTree := fs.NewTree([]string{execute.ToolsPath()}, 0755, fs.WithIndex(ToolIndex))
	transcriptStore := fs.NewTree([]string{transcriptDir}, 0444, fs.WithReadOnly())

	mgr := session.NewManager(session.ManagerConfig{
		AgentsDir:   agentsDirs[0],
		SessionsDir: sessionsDir,
		Log:         sink.NewLogger("9p"),
		Sink:        sink,
		Strict:      *strict,
		Yolo:        *yolo,
		ToolTree:   toolTree,
		SaveTranscript: func(data []byte) error {
			name := time.Now().Format("20060102T150405") + "-chat.md"
			e, err := transcriptStore.Open(name)
			if err != nil {
				transcriptStore.Create(name) //nolint:errcheck
				e, err = transcriptStore.Open(name)
			}
			if err != nil {
				return err
			}
			return e.Write(data)
		},
	})

	srv := server.New(server.Config{
		Sink:            sink,
		AgentStore:      fs.NewTree(agentsDirs, 0644),
		PromptStore:     fs.NewTree(agent.PromptsDirs(), 0444, fs.WithReadOnly()),
		MemStore:        fs.NewTree([]string{memDir}, 0644),
		ToolTree:       toolTree,
		UtilStore:       fs.NewTree([]string{paths.CfgDir() + "/scripts/u"}, 0555, fs.WithReadOnly()),
		PluginStore:     fs.NewTree([]string{execute.PluginsPath()}, 0555, fs.WithReadOnly()),
		SkillTree:      fs.NewSkillTree(),
		SessionMgr:      mgr,
		TranscriptStore: transcriptStore,
		TmpStore:        fs.NewTree([]string{tmpDir}, 0600),
		RootStore:       NewRootStore(),
	})

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

	// Optional FUSE mount
	mnt := *mountPath
	if mnt == "" {
		mnt = os.Getenv("OLLIE")
	}
	if mnt == "" {
		home, _ := os.UserHomeDir()
		mnt = filepath.Join(home, "mnt", "ollie")
	}
	var fuseCmd *exec.Cmd
	if err := os.MkdirAll(mnt, 0755); err != nil {
		fmt.Fprintf(os.Stderr, "warning: cannot create mount dir: %v\n", err)
	} else {
		fuseCmd = exec.Command("9pfuse", sockPath, mnt)
		if err := fuseCmd.Start(); err != nil {
			fmt.Fprintf(os.Stderr, "warning: 9pfuse: %v\n", err)
			fuseCmd = nil
		} else {
			fmt.Printf("mounted at %s\n", mnt)
		}
	}

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

	srv.InterruptAll()
	fmt.Println("shutting down")
	srv.Shutdown()
	if fuseCmd != nil {
		exec.Command("fusermount", "-u", mnt).Run() //nolint:errcheck
		fuseCmd.Wait()                               //nolint:errcheck
	}
	listener.Close() //nolint:errcheck
	if tcpListener != nil {
		tcpListener.Close() //nolint:errcheck
	}
	os.Remove(sockPath)
	os.Remove(pidPath)
	sink.Flush()
}

func defaultMemDir() string {
	if p := os.Getenv("OLLIE_MEMORY_PATH"); p != "" {
		return p
	}
	return paths.CfgDir() + "/memory"
}

func defaultTranscriptDir() string {
	if p := os.Getenv("OLLIE_TRANSCRIPT_PATH"); p != "" {
		return p
	}
	return paths.CfgDir() + "/transcript"
}

func defaultTmpDir() string {
	if p := os.Getenv("OLLIE_TMP_PATH"); p != "" {
		return p
	}
	return paths.DataDir() + "/tmp"
}

// ToolIndex generates a tool index from a tree's file listing.
func ToolIndex(s fs.FileTree) ([]byte, error) {
	entries, err := s.List()
	if err != nil {
		return nil, err
	}
	var sb strings.Builder
	for _, e := range entries {
		if e.IsDir() || e.Name() == "idx" {
			continue
		}
		entry, err := s.Open(e.Name())
		if err != nil {
			continue
		}
		data, err := entry.Read()
		if err != nil {
			continue
		}
		var desc, args string
		for line := range strings.SplitSeq(string(data), "\n") {
			if d, ok := strings.CutPrefix(line, "# description:"); ok {
				desc = strings.TrimSpace(d)
			} else if a, ok := strings.CutPrefix(line, "# Args:"); ok {
				args = strings.TrimSpace(a)
			} else if a, ok := strings.CutPrefix(line, "# args:"); ok {
				args = strings.TrimSpace(a)
			}
			if !strings.HasPrefix(line, "#") && line != "" {
				break
			}
		}
		if desc == "" {
			continue
		}
		fmt.Fprintf(&sb, "## %s\n", e.Name())
		fmt.Fprintf(&sb, "description: %s\n", desc)
		if args != "" {
			fmt.Fprintf(&sb, "args: %s\n", args)
		}
		sb.WriteString("\n")
	}
	return []byte(sb.String()), nil
}

// 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){
		"backends": func() ([]byte, error) {
			return []byte(strings.Join(backend.Backends(), "\n") + "\n"), nil
		},
		"help": func() ([]byte, error) {
			return os.ReadFile(helpPath)
		},
	}
	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") }

	return fs.NewTree(nil, 0444,
		fs.WithReadOnly(),
		fs.WithResolver(func(_ []string, name string) (string, error) {
			if _, ok := entries[name]; ok {
				return name, nil
			}
			return "", fmt.Errorf("%s: not found", name)
		}),
		fs.WithLister(func(_ []string) ([]os.DirEntry, error) {
			return []os.DirEntry{
				fs.FileEntry("backends", 0444),
				fs.FileEntry("help", 0444),
			}, nil
		}),
		fs.WithOpener(func(_ []string, name string) (fs.File, error) {
			readFn, ok := entries[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_: 0444}, nil },
				ReadFn:         readFn,
				WriteFn:        readOnly,
				BlockingReadFn: notBlocking,
			}, nil
		}),
	)
}