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