root.go

  1package cmd
  2
  3import (
  4	"context"
  5	"log/slog"
  6	"os"
  7	"sync"
  8
  9	tea "github.com/charmbracelet/bubbletea"
 10	"github.com/kujtimiihoxha/termai/internal/app"
 11	"github.com/kujtimiihoxha/termai/internal/assets"
 12	"github.com/kujtimiihoxha/termai/internal/config"
 13	"github.com/kujtimiihoxha/termai/internal/db"
 14	"github.com/kujtimiihoxha/termai/internal/llm/agent"
 15	"github.com/kujtimiihoxha/termai/internal/logging"
 16	"github.com/kujtimiihoxha/termai/internal/tui"
 17	zone "github.com/lrstanley/bubblezone"
 18	"github.com/spf13/cobra"
 19)
 20
 21var rootCmd = &cobra.Command{
 22	Use:   "termai",
 23	Short: "A terminal ai assistant",
 24	Long:  `A terminal ai assistant`,
 25	RunE: func(cmd *cobra.Command, args []string) error {
 26		if cmd.Flag("help").Changed {
 27			cmd.Help()
 28			return nil
 29		}
 30		debug, _ := cmd.Flags().GetBool("debug")
 31		err := config.Load(debug)
 32		if err != nil {
 33			return err
 34		}
 35		cfg := config.Get()
 36		defaultLevel := slog.LevelInfo
 37		if cfg.Debug {
 38			defaultLevel = slog.LevelDebug
 39		}
 40		logger := slog.New(slog.NewTextHandler(logging.NewWriter(), &slog.HandlerOptions{
 41			Level: defaultLevel,
 42		}))
 43		slog.SetDefault(logger)
 44
 45		err = assets.WriteAssets()
 46		if err != nil {
 47			return err
 48		}
 49
 50		conn, err := db.Connect()
 51		if err != nil {
 52			return err
 53		}
 54		ctx := context.Background()
 55
 56		app := app.New(ctx, conn)
 57		logging.Info("Starting termai...")
 58		zone.NewGlobal()
 59		tui := tea.NewProgram(
 60			tui.New(app),
 61			tea.WithAltScreen(),
 62			tea.WithMouseCellMotion(),
 63		)
 64		logging.Info("Setting up subscriptions...")
 65		ch, unsub := setupSubscriptions(app)
 66		defer unsub()
 67
 68		go func() {
 69			// Set this up once
 70			agent.GetMcpTools(ctx, app.Permissions)
 71			for msg := range ch {
 72				tui.Send(msg)
 73			}
 74		}()
 75		if _, err := tui.Run(); err != nil {
 76			return err
 77		}
 78		return nil
 79	},
 80}
 81
 82func setupSubscriptions(app *app.App) (chan tea.Msg, func()) {
 83	ch := make(chan tea.Msg)
 84	wg := sync.WaitGroup{}
 85	ctx, cancel := context.WithCancel(app.Context)
 86	{
 87		sub := logging.Subscribe(ctx)
 88		wg.Add(1)
 89		go func() {
 90			for ev := range sub {
 91				ch <- ev
 92			}
 93			wg.Done()
 94		}()
 95	}
 96	{
 97		sub := app.Sessions.Subscribe(ctx)
 98		wg.Add(1)
 99		go func() {
100			for ev := range sub {
101				ch <- ev
102			}
103			wg.Done()
104		}()
105	}
106	{
107		sub := app.Messages.Subscribe(ctx)
108		wg.Add(1)
109		go func() {
110			for ev := range sub {
111				ch <- ev
112			}
113			wg.Done()
114		}()
115	}
116	{
117		sub := app.Permissions.Subscribe(ctx)
118		wg.Add(1)
119		go func() {
120			for ev := range sub {
121				ch <- ev
122			}
123			wg.Done()
124		}()
125	}
126	return ch, func() {
127		cancel()
128		wg.Wait()
129		close(ch)
130	}
131}
132
133func Execute() {
134	err := rootCmd.Execute()
135	if err != nil {
136		os.Exit(1)
137	}
138}
139
140func init() {
141	rootCmd.Flags().BoolP("help", "h", false, "Help")
142	rootCmd.Flags().BoolP("debug", "d", false, "Help")
143}