agentu/main.go

package main

import (
	"bufio"
	"context"
	"errors"
	"flag"
	"fmt"
	"net/http"
	"os"
	"os/signal"
	"path/filepath"
	"strings"

	"agentu/internal/agent"
	"agentu/internal/session"
	"agentu/internal/tools"
	"agentu/internal/tui"
	"agentu/pkg/config"
	"agentu/pkg/llm"
)

const defaultConfigContent = `# agentu configuration
# Replace the placeholder values below with your actual provider settings.
# See agentu.example.yaml for the full reference.
providers:
  default:
    base_url: https://api.openai.com
    api_key: YOUR_API_KEY_HERE
    models:
      - id: gpt-4o
        context: 128000
      - id: gpt-4o-mini
        context: 128000
`

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

func run() error {
	configPath := flag.String("config", config.DefaultPath, "path to YAML config")
	yolo := flag.Bool("yolo", false, "also enable automatic file writes and shell execution")
	plain := flag.Bool("plain", false, "use the simple line-based REPL instead of the TUI")
	themeName := flag.String("theme", "light", "TUI theme: light or dark")
	resumeID := flag.String("resume", "", "resume a specific session by ID")
	flag.Parse()

	themeMode, err := tui.ParseThemeMode(*themeName)
	if err != nil {
		return err
	}

	cfg, err := config.Load(*configPath)
	if err != nil {
		if errors.Is(err, os.ErrNotExist) {
			if bootstrapErr := bootstrapConfig(*configPath); bootstrapErr != nil {
				return bootstrapErr
			}
			return fmt.Errorf("created default config at %s; edit it to set your API key and provider, then restart agentu", *configPath)
		}
		return err
	}

	providerName := cfg.DefaultProviderName()
	providerConfig := cfg.Providers[providerName]
	provider := llm.NewOpenAICompatibleClient(providerConfig.BaseURL, providerConfig.APIKey, http.DefaultClient)

	var registry *tools.Registry
	if *yolo {
		registry = tools.Builtins(cfg.Agent.WorkingDir)
	} else {
		registry = tools.ReadOnlyBuiltins(cfg.Agent.WorkingDir)
	}
	registry.Register(tools.NewBuiltinSearchTool(http.DefaultClient))
	registry.Register(tools.NewBuiltinFetchTool(http.DefaultClient))
	systemPrompt := strings.TrimSpace(cfg.Agent.SystemPrompt + "\n\n" + tools.AgentInstructions(cfg.Agent.WorkingDir, *yolo))
	maxContextTokens := cfg.Agent.MaxContextTokens
	if maxContextTokens == 0 {
		maxContextTokens = providerConfig.ContextTokens
	}

	assistant := agent.New(agent.Options{
		Provider:         provider,
		ProviderName:     providerName,
		Model:            providerConfig.Model,
		Thinking:         providerConfig.Thinking,
		ThinkingParam:    providerConfig.ThinkingParam,
		ThinkingEnabled:  providerConfig.ThinkingConfigured,
		SystemPrompt:     systemPrompt,
		ToolRegistry:     registry,
		ToolTimeout:      cfg.Agent.ToolTimeout,
		MaxContextTokens: maxContextTokens,
	})

	// Initialize session store
	home, err := os.UserHomeDir()
	if err != nil {
		return fmt.Errorf("resolve home dir: %w", err)
	}
	sessionsDir := filepath.Join(home, ".agentu", "sessions")
	store, err := session.NewStore(sessionsDir)
	if err != nil {
		return fmt.Errorf("init session store: %w", err)
	}

	modelManager, err := session.NewManager(cfg, assistant, http.DefaultClient, store)
	if err != nil {
		return err
	}

	if err := initializeSession(modelManager, *resumeID); err != nil {
		return err
	}

	ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt)
	defer stop()

	if !*plain {
		err := tui.Run(ctx, assistant, tui.Options{
			ModelName:    providerConfig.Model,
			Yolo:         *yolo,
			ThemeMode:    themeMode,
			ModelManager: modelManager,
			WorkingDir:   cfg.Agent.WorkingDir,
		})
		printExitSession(modelManager)
		return err
	}

	return repl(ctx, assistant, modelManager)
}

func initializeSession(modelManager *session.Manager, resumeID string) error {
	if resumeID != "" {
		if err := modelManager.Resume(resumeID); err != nil {
			return fmt.Errorf("resume session: %w", err)
		}
		return nil
	}
	if _, err := modelManager.NewSession(); err != nil {
		return fmt.Errorf("start new session: %w", err)
	}
	return nil
}

func bootstrapConfig(configPath string) error {
	resolved, err := config.ResolvePath(configPath)
	if err != nil {
		return fmt.Errorf("resolve config path: %w", err)
	}
	dir := filepath.Dir(resolved)
	if err := os.MkdirAll(dir, 0o755); err != nil {
		return fmt.Errorf("create config directory: %w", err)
	}
	if err := os.WriteFile(resolved, []byte(defaultConfigContent), 0o644); err != nil {
		return fmt.Errorf("write default config: %w", err)
	}
	return nil
}

func repl(ctx context.Context, assistant *agent.Agent, modelManager *session.Manager) error {
	scanner := bufio.NewScanner(os.Stdin)
	scanner.Buffer(make([]byte, 0, 64*1024), 1024*1024)

	fmt.Println("agentu")
	fmt.Println("Type /exit to quit, /clear to reset context, /compact to compress context.")
	for {
		fmt.Print("> ")
		if !scanner.Scan() {
			if err := scanner.Err(); err != nil {
				return err
			}
			fmt.Println()
			printExitSession(modelManager)
			return nil
		}

		input := strings.TrimSpace(scanner.Text())
		switch input {
		case "":
			continue
		case "/exit", "/quit":
			printExitSession(modelManager)
			return nil
		case "/clear":
			assistant.Clear()
			fmt.Println("context cleared")
			continue
		case "/compact":
			if err := assistant.Compact(ctx); err != nil {
				fmt.Fprintln(os.Stderr, "compact error:", err)
			} else {
				fmt.Println("context compacted")
				if modelManager != nil {
					_ = modelManager.Save()
				}
			}
			continue
		}

		if err := assistant.RunTurn(ctx, input, os.Stdout, os.Stderr); err != nil {
			fmt.Fprintln(os.Stderr, "\nerror:", err)
		}
		_ = modelManager.Save()
		fmt.Println()
	}
}

func printExitSession(m *session.Manager) {
	if m == nil {
		return
	}
	_ = m.Save()
	id := m.CurrentSessionID()
	name := m.CurrentSessionName()
	if id != "" {
		if name != "" && name != id {
			fmt.Fprintf(os.Stderr, "session: %s (%s)\n", id, name)
		} else {
			fmt.Fprintf(os.Stderr, "session: %s\n", id)
		}
		fmt.Fprintf(os.Stderr, "resume:  agentu --resume %s\n", id)
	}
}