localchat/cmd/localchat/main.go

// Command localchat is a chat UI for a local, open-weight model.
//
// It runs in the foreground, or registers itself as a background service
// with the OS (systemd on Linux, launchd on macOS, the Service Control
// Manager on Windows).
package main

import (
	"context"
	"errors"
	"flag"
	"fmt"
	"log/slog"
	"net"
	"net/http"
	"os"
	"path/filepath"
	"time"

	"github.com/kardianos/service"

	"git.b0b.be/bdeb/localchat/internal/config"
	"git.b0b.be/bdeb/localchat/internal/web"
)

var version = "dev" // set with -ldflags "-X main.version=..."

const usage = `localchat %s: chat with a local AI model in your browser.

Usage:
  localchat [command] [flags]

Commands:
  serve       run in the foreground (default)
  install     register as a background service that starts on boot/login
  uninstall   remove the service
  start       start the installed service
  stop        stop the installed service
  restart     restart the installed service
  status      show whether the service is running
  healthcheck exit 0 if the server answers (for container health checks)
  version     print the version

Flags:
`

func main() {
	if err := run(os.Args[1:]); err != nil {
		fmt.Fprintln(os.Stderr, "localchat:", err)
		os.Exit(1)
	}
}

func run(args []string) error {
	cmd := "serve"
	if len(args) > 0 && args[0] != "" && args[0][0] != '-' {
		cmd, args = args[0], args[1:]
	}

	fs := flag.NewFlagSet("localchat", flag.ContinueOnError)
	envFile := fs.String("env-file", "", "load settings from this KEY=VALUE file (default: ./.env if present)")
	user := fs.Bool("user", false, "install as a per-user service (launchd agent / systemd --user) instead of system-wide")
	debug := fs.Bool("debug", false, "verbose logging")
	fs.Usage = func() {
		fmt.Fprintf(fs.Output(), usage, version)
		fs.PrintDefaults()
	}
	if err := fs.Parse(args); err != nil {
		if errors.Is(err, flag.ErrHelp) {
			return nil
		}
		return err
	}

	if cmd == "version" {
		fmt.Println(version)
		return nil
	}

	if *envFile == "" {
		if _, err := os.Stat(".env"); err == nil {
			*envFile = ".env"
		}
	}
	if *envFile != "" {
		abs, err := filepath.Abs(*envFile)
		if err != nil {
			return err
		}
		*envFile = abs
		if err := config.LoadEnvFile(abs); err != nil {
			return fmt.Errorf("env file: %w", err)
		}
	}

	level := slog.LevelInfo
	if *debug {
		level = slog.LevelDebug
	}
	log := slog.New(slog.NewTextHandler(os.Stderr, &slog.HandlerOptions{Level: level}))

	// The service manager re-runs us as "localchat serve --env-file <abs>",
	// so the installed service sees the same settings as this invocation.
	svcArgs := []string{"serve"}
	if *envFile != "" {
		svcArgs = append(svcArgs, "--env-file", *envFile)
	}
	wd, _ := os.Getwd()
	prg := &program{log: log}
	svc, err := service.New(prg, &service.Config{
		Name:             "localchat",
		DisplayName:      "localchat",
		Description:      "Chat with a local AI model in your browser.",
		Arguments:        svcArgs,
		WorkingDirectory: wd,
		Option:           service.KeyValue{"UserService": *user},
	})
	if err != nil {
		return err
	}

	switch cmd {
	case "healthcheck":
		return healthcheck()
	case "serve":
		return svc.Run() // handles Ctrl+C in a terminal and stop requests from the OS
	case "status":
		st, err := svc.Status()
		if err != nil {
			return err
		}
		fmt.Println(map[service.Status]string{
			service.StatusRunning: "running",
			service.StatusStopped: "stopped",
		}[st])
		return nil
	case "install", "uninstall", "start", "stop", "restart":
		if err := service.Control(svc, cmd); err != nil {
			return err
		}
		fmt.Printf("localchat: %s ok\n", cmd)
		return nil
	default:
		fs.Usage()
		return fmt.Errorf("unknown command %q", cmd)
	}
}

func healthcheck() error {
	cfg, err := config.Load()
	if err != nil {
		return err
	}
	host, port, err := net.SplitHostPort(cfg.Addr)
	if err != nil {
		return err
	}
	if host == "" || host == "0.0.0.0" || host == "::" {
		host = "127.0.0.1"
	}
	c := &http.Client{Timeout: 3 * time.Second}
	resp, err := c.Get("http://" + net.JoinHostPort(host, port) + "/healthz")
	if err != nil {
		return err
	}
	resp.Body.Close()
	if resp.StatusCode != http.StatusOK {
		return fmt.Errorf("healthz: %s", resp.Status)
	}
	return nil
}

// program adapts the web server to the service lifecycle.
type program struct {
	log    *slog.Logger
	cancel context.CancelFunc
	done   chan struct{}
}

func (p *program) Start(service.Service) error {
	cfg, err := config.Load()
	if err != nil {
		return err
	}
	ctx, cancel := context.WithCancel(context.Background())
	p.cancel, p.done = cancel, make(chan struct{})
	go func() {
		defer close(p.done)
		if err := web.New(cfg, p.log).Run(ctx); err != nil {
			p.log.Error("server stopped", "err", err)
			os.Exit(1)
		}
	}()
	return nil
}

func (p *program) Stop(service.Service) error {
	p.cancel()
	<-p.done
	return nil
}