ollie/cmd/toolsrv/main.go

// toolsrv — 9P tool execution server.
//
// Executes tool calls in a sandboxed environment. Used both locally
// (spawned by olliesrv) and remotely (deployed via SSH bootstrap).
//
// Usage:
//
//	toolsrv serve --cwd /path/to/project --listen /path/to/sock [--yolo]
//
// Listens on a Unix socket and serves a 9P filesystem for tool execution.
// Authentication happens via Tauth: first client to connect sets the secret,
// subsequent clients must provide the same secret.
package main

import (
	"context"
	"flag"
	"fmt"
	"log/slog"
	"net"
	"os"
	"os/signal"
	"path/filepath"
	"sync/atomic"
	"syscall"
	"time"

	"ollie/cmd/toolsrv/internal/fs"
	"ollie/cmd/toolsrv/internal/registry"
	"ollie/env"
	"ollie/fsedsl"
)

var (
	cwd        = flag.String("cwd", ".", "working directory for execution")
	listenPath = flag.String("listen", "", "Unix socket path to listen on (required)")
	sessionID  = flag.String("session-id", "", "session ID for tool registry")
	yolo       = flag.Bool("yolo", false, "skip sandbox enforcement")
	noAuth     = flag.Bool("no-auth", false, "disable Tauth authentication (for debugging)")
	idleTimeout = flag.Duration("idle-timeout", 30*time.Second, "exit after this duration with no connections (0 to disable)")
)

// logger is the toolsrv structured logger.
var logger *slog.Logger

// activeConns tracks the number of active 9P connections.
var activeConns atomic.Int64

// hadConnection is set to true once the first client connects.
var hadConnection atomic.Bool

func main() {
	if len(os.Args) < 2 || os.Args[1] != "serve" {
		fmt.Fprintln(os.Stderr, "usage: toolsrv serve --cwd <path> --listen <socket> [--yolo]")
		os.Exit(1)
	}
	flag.CommandLine.Parse(os.Args[2:])

	if *listenPath == "" {
		fmt.Fprintln(os.Stderr, "error: --listen is required")
		os.Exit(1)
	}

	// Set up structured logger
	logger = slog.New(slog.NewTextHandler(os.Stderr, &slog.HandlerOptions{
		Level: slog.LevelInfo,
	})).With("svc", "toolsrv")

	env.EnsureDefaults()

	// Prepend our bin dir to PATH so landrun is found (deployed alongside us).
	home, _ := os.UserHomeDir()
	binDir := filepath.Join(home, ".config", "ollie", "bin")
	os.Setenv("PATH", binDir+":"+os.Getenv("PATH"))

	// Create tool registry
	toolReg, _ := registry.New()
	sid := *sessionID
	if sid == "" {
		sid = os.Getenv("OLLIE_SESSION_ID") // fallback to env var
	}

	// Create server state
	srv := fs.NewServer()
	srv.SetYolo(*yolo)
	if toolReg != nil && sid != "" {
		srv.SetRegistry(toolReg, sid)
	}

	// Build the filesystem tree
	ctx := fs.Ctx{Server: srv}
	tree := fsedsl.BuildTree(fs.Spec(), ctx)
	if tree == nil {
		logger.Error("failed to build filesystem tree")
		os.Exit(1)
	}

	// Set up signal handling
	runCtx, cancel := context.WithCancel(context.Background())
	defer cancel()
	sigCh := make(chan os.Signal, 1)
	signal.Notify(sigCh, syscall.SIGINT, syscall.SIGTERM)
	go func() {
		<-sigCh
		cancel()
	}()

	// Kill all procs on shutdown
	defer srv.Fs.KillAll()

	// Start idle timeout monitor
	if *idleTimeout > 0 {
		go idleMonitor(runCtx, cancel, *idleTimeout)
	}

	// Start proc GC
	srv.Fs.StartProcGC(runCtx)

	// Remove stale socket if it exists
	os.Remove(*listenPath)

	// Listen on Unix socket
	ln, err := net.Listen("unix", *listenPath)
	if err != nil {
		logger.Error("listen failed", "err", err)
		os.Exit(1)
	}
	defer ln.Close()
	defer os.Remove(*listenPath)

	logger.Info("listening on " + *listenPath)

	// Accept connections
	go func() {
		<-runCtx.Done()
		ln.Close()
	}()

	for {
		conn, err := ln.Accept()
		if err != nil {
			select {
			case <-runCtx.Done():
				return
			default:
				logger.Warn("accept error", "err", err)
				continue
			}
		}
		activeConns.Add(1)
		hadConnection.Store(true)
		go func() {
			defer activeConns.Add(-1)
			serve9P(runCtx, conn, tree, srv, *noAuth)
		}()
	}
}

// idleMonitor exits the process if there are no active connections for the
// given duration. It waits for the first connection before monitoring (with
// a startup grace period of 60s).
func idleMonitor(ctx context.Context, cancel context.CancelFunc, timeout time.Duration) {
	const startupGrace = 60 * time.Second

	// Wait for first connection (or startup grace period)
	deadline := time.After(startupGrace)
	for {
		select {
		case <-ctx.Done():
			return
		case <-deadline:
			if !hadConnection.Load() {
				logger.Info("no connections within startup grace period, exiting")
				cancel()
				return
			}
		case <-time.After(1 * time.Second):
			if hadConnection.Load() {
				goto monitor
			}
		}
	}

monitor:
	// Monitor for idle (zero connections after having had at least one)
	idleSince := time.Time{}
	for {
		select {
		case <-ctx.Done():
			return
		case <-time.After(5 * time.Second):
			if activeConns.Load() > 0 {
				idleSince = time.Time{}
				continue
			}
			// No active connections
			if idleSince.IsZero() {
				idleSince = time.Now()
				logger.Info("all connections closed, idle timer started", "timeout", timeout)
			} else if time.Since(idleSince) >= timeout {
				logger.Info("idle timeout reached, exiting")
				cancel()
				return
			}
		}
	}
}