From dabf5bfeccda907dd5738ab907ceae5f0fafd73f Mon Sep 17 00:00:00 2001 From: loveuer Date: Wed, 24 Jun 2026 23:59:21 -0700 Subject: [PATCH] compact: retain newest ~20% of context; refactor startup into internal/app - Replace fixed 6-message compact retention with ratio-based logic: retain ~20% of compactable estimated tokens, with a 6-message floor and user-turn boundary alignment - Add compactRetainedTokenBudget, compactRecentStart, compactTurnBoundary helpers in internal/agent/agent.go - Add TestAgentCompactKeepsTwentyPercentRecentContext - Move startup/session/repl/bootstrap logic from root into internal/app so main.go is a thin binary entry point (15 lines) - Update /compact docs in README.md --- README.md | 2 +- internal/agent/agent.go | 66 ++++++- internal/agent/agent_test.go | 45 +++++ internal/app/app.go | 164 ++++++++++++++++ internal/app/bootstrap.go | 38 ++++ main_test.go => internal/app/main_test.go | 2 +- internal/app/repl.go | 60 ++++++ internal/app/session_startup.go | 38 ++++ main.go | 217 +--------------------- 9 files changed, 408 insertions(+), 224 deletions(-) create mode 100644 internal/app/app.go create mode 100644 internal/app/bootstrap.go rename main_test.go => internal/app/main_test.go (99%) create mode 100644 internal/app/repl.go create mode 100644 internal/app/session_startup.go diff --git a/README.md b/README.md index ffb836d..c4cafbf 100644 --- a/README.md +++ b/README.md @@ -109,7 +109,7 @@ go run . [flags] Slash commands: - `/clear` resets conversation history. -- `/compact` summarizes older conversation history and keeps recent context. +- `/compact` summarizes older conversation history and keeps the newest ~20% of context, with a small recent-message floor. - `/status` shows changed files with `git status --short`. - `/diff [--stat] [path...]` shows the unstaged git diff. - `/test [command...]` runs a project test command; defaults to `go test ./...` for Go modules or `npm test` for Node projects. diff --git a/internal/agent/agent.go b/internal/agent/agent.go index c282966..f295cc1 100644 --- a/internal/agent/agent.go +++ b/internal/agent/agent.go @@ -17,7 +17,8 @@ import ( const ( defaultMaxToolRounds = 50 doomLoopThreshold = 3 - compactRecentMessages = 6 + compactMinRecentMessages = 6 + compactRecentRetentionPercent = 20 summaryMessageLimit = 8000 contextSummaryPrefix = "[context summary]" toolLogArgumentLimit = 400 @@ -181,6 +182,61 @@ func (a *Agent) SetLastUsage(usage *llm.Usage) { copy := *usage a.lastUsage = © } +// compactRetainedTokenBudget returns the number of estimated tokens to retain +// based on compactRecentRetentionPercent, using integer ceiling division. +func compactRetainedTokenBudget(totalTokens int) int { + if totalTokens <= 0 { + return 1 + } + return (totalTokens*compactRecentRetentionPercent + 99) / 100 +} + +// compactTurnBoundary walks backward from start until it finds a user message +// or reaches prefixEnd, ensuring the retained suffix starts at a user turn. +func compactTurnBoundary(messages []llm.Message, prefixEnd int, start int) int { + for start > prefixEnd && messages[start].Role != llm.RoleUser { + start-- + } + return start +} + +// compactRecentStart returns the index where the retained recent suffix begins. +// It balances the 20% token-budget ratio with the compactMinRecentMessages floor +// and aligns to a user-turn boundary. +func compactRecentStart(messages []llm.Message, prefixEnd int) int { + compactableCount := len(messages) - prefixEnd + if compactableCount <= compactMinRecentMessages { + return len(messages) + } + + compactable := messages[prefixEnd:] + totalTokens := llm.EstimateTokens(compactable, nil) + retainTokens := compactRetainedTokenBudget(totalTokens) + + // Walk backward from the end to find the earliest suffix fitting the budget. + ratioStart := len(messages) - 1 + for i := len(messages) - 1; i >= prefixEnd; i-- { + if llm.EstimateTokens(messages[i:], nil) <= retainTokens { + ratioStart = i + } else { + break + } + } + + // Enforce the minimum recent-message floor. + minStart := len(messages) - compactMinRecentMessages + if minStart < prefixEnd { + minStart = prefixEnd + } + // Choose the earlier start so the floor is never weakened. + start := minStart + if ratioStart < start { + start = ratioStart + } + + return compactTurnBoundary(messages, prefixEnd, start) +} + // Compact summarizes older conversation messages while preserving the system // prompt and recent messages intact. @@ -196,12 +252,8 @@ func (a *Agent) Compact(ctx context.Context) error { if a.messages[0].Role == llm.RoleSystem { prefixEnd = 1 } - if len(a.messages)-prefixEnd <= compactRecentMessages { - return nil - } - - recentStart := len(a.messages) - compactRecentMessages - if recentStart <= prefixEnd { + recentStart := compactRecentStart(a.messages, prefixEnd) + if recentStart <= prefixEnd || recentStart >= len(a.messages) { return nil } middle := append([]llm.Message(nil), a.messages[prefixEnd:recentStart]...) diff --git a/internal/agent/agent_test.go b/internal/agent/agent_test.go index c51180c..5183e90 100644 --- a/internal/agent/agent_test.go +++ b/internal/agent/agent_test.go @@ -145,6 +145,9 @@ func TestAgentCompact(t *testing.T) { if messages[len(messages)-1].Content != "recent-answer-2" { t.Fatalf("last recent message = %#v", messages[len(messages)-1]) } + if messages[2].Content != "recent-user-0" { + t.Fatalf("first retained recent message = %#v", messages[2]) + } if usage := a.LastUsage(); usage == nil || usage.TotalTokens != 12 { t.Fatalf("usage = %#v", usage) } @@ -217,6 +220,48 @@ func longConversation() []llm.Message { return messages } +func manyTurnConversation(turns int) []llm.Message { + messages := []llm.Message{{Role: llm.RoleSystem, Content: "system"}} + for i := 0; i < turns; i++ { + messages = append(messages, + llm.Message{Role: llm.RoleUser, Content: fmt.Sprintf("user-%02d", i)}, + llm.Message{Role: llm.RoleAssistant, Content: fmt.Sprintf("assistant-%02d", i)}, + ) + } + return messages +} + +func TestAgentCompactKeepsTwentyPercentRecentContext(t *testing.T) { + provider := &compactProvider{} + a := New(Options{Provider: provider, Model: "test", SystemPrompt: "system"}) + a.SetMessages(manyTurnConversation(25)) + + if err := a.Compact(context.Background()); err != nil { + t.Fatal(err) + } + + messages := a.Messages() + // 1 system + 1 summary + 10 retained recent messages (5 user/assistant turns) + if len(messages) != 12 { + t.Fatalf("messages = %d, want 12; contents: %v", len(messages), joinMessageContents(messages)) + } + if provider.summaryCalls != 1 { + t.Fatalf("summaryCalls = %d, want 1", provider.summaryCalls) + } + if messages[2].Content != "user-20" { + t.Fatalf("first retained message = %#v, want user-20", messages[2]) + } + if messages[len(messages)-1].Content != "assistant-24" { + t.Fatalf("last retained message = %#v, want assistant-24", messages[len(messages)-1]) + } + if !strings.Contains(provider.lastSummary, "user-19") { + t.Fatalf("summary should contain user-19 (part of summarized 80%%): %q", provider.lastSummary) + } + if strings.Contains(provider.lastSummary, "user-20") { + t.Fatalf("summary should not contain user-20 (part of retained 20%%): %q", provider.lastSummary) + } +} + func joinMessageContents(messages []llm.Message) string { var b strings.Builder for _, msg := range messages { diff --git a/internal/app/app.go b/internal/app/app.go new file mode 100644 index 0000000..37e3e99 --- /dev/null +++ b/internal/app/app.go @@ -0,0 +1,164 @@ +package app + +import ( + "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" +) + +type cliOptions struct { + configPath string + yolo bool + plain bool + themeName string + resumeID string +} + +func Run() error { + opts := parseCLIOptions() + + themeMode, err := tui.ParseThemeMode(opts.themeName) + if err != nil { + return err + } + + cfg, err := loadConfig(opts.configPath) + if err != nil { + return err + } + + providerName, providerConfig, provider := defaultProvider(cfg, http.DefaultClient) + assistant := newAssistant(cfg, providerName, providerConfig, provider, opts.yolo, http.DefaultClient) + modelManager, err := newSessionManager(cfg, assistant, http.DefaultClient) + if err != nil { + return err + } + + if err := initializeSession(modelManager, opts.resumeID); err != nil { + return err + } + + ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt) + defer stop() + + return runInterface(ctx, assistant, modelManager, providerConfig, cfg.Agent.WorkingDir, opts, themeMode) +} + +func parseCLIOptions() cliOptions { + 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() + + return cliOptions{ + configPath: *configPath, + yolo: *yolo, + plain: *plain, + themeName: *themeName, + resumeID: *resumeID, + } +} + +func loadConfig(configPath string) (*config.Config, error) { + cfg, err := config.Load(configPath) + if err == nil { + return cfg, nil + } + if !errors.Is(err, os.ErrNotExist) { + return nil, err + } + if bootstrapErr := bootstrapConfig(configPath); bootstrapErr != nil { + return nil, bootstrapErr + } + return nil, fmt.Errorf("created default config at %s; edit it to set your API key and provider, then restart agentu", configPath) +} + +func defaultProvider(cfg *config.Config, httpClient *http.Client) (string, config.ProviderConfig, llm.Provider) { + providerName := cfg.DefaultProviderName() + providerConfig := cfg.Providers[providerName] + provider := llm.NewOpenAICompatibleClient(providerConfig.BaseURL, providerConfig.APIKey, httpClient) + return providerName, providerConfig, provider +} + +func newAssistant(cfg *config.Config, providerName string, providerConfig config.ProviderConfig, provider llm.Provider, yolo bool, httpClient *http.Client) *agent.Agent { + return agent.New(agent.Options{ + Provider: provider, + ProviderName: providerName, + Model: providerConfig.Model, + Thinking: providerConfig.Thinking, + ThinkingParam: providerConfig.ThinkingParam, + ThinkingEnabled: providerConfig.ThinkingConfigured, + SystemPrompt: systemPrompt(cfg, yolo), + ToolRegistry: newToolRegistry(cfg.Agent.WorkingDir, yolo, httpClient), + ToolTimeout: cfg.Agent.ToolTimeout, + MaxContextTokens: maxContextTokens(cfg, providerConfig), + }) +} + +func systemPrompt(cfg *config.Config, yolo bool) string { + instructions := tools.AgentInstructions(cfg.Agent.WorkingDir, yolo) + return strings.TrimSpace(cfg.Agent.SystemPrompt + "\n\n" + instructions) +} + +func maxContextTokens(cfg *config.Config, providerConfig config.ProviderConfig) int { + if cfg.Agent.MaxContextTokens != 0 { + return cfg.Agent.MaxContextTokens + } + return providerConfig.ContextTokens +} + +func newToolRegistry(workingDir string, yolo bool, httpClient *http.Client) *tools.Registry { + var registry *tools.Registry + if yolo { + registry = tools.Builtins(workingDir) + } else { + registry = tools.ReadOnlyBuiltins(workingDir) + } + registry.Register(tools.NewBuiltinSearchTool(httpClient)) + registry.Register(tools.NewBuiltinFetchTool(httpClient)) + return registry +} + +func newSessionManager(cfg *config.Config, assistant *agent.Agent, httpClient *http.Client) (*session.Manager, error) { + home, err := os.UserHomeDir() + if err != nil { + return nil, fmt.Errorf("resolve home dir: %w", err) + } + sessionsDir := filepath.Join(home, ".agentu", "sessions") + store, err := session.NewStore(sessionsDir) + if err != nil { + return nil, fmt.Errorf("init session store: %w", err) + } + return session.NewManager(cfg, assistant, httpClient, store) +} + +func runInterface(ctx context.Context, assistant *agent.Agent, modelManager *session.Manager, providerConfig config.ProviderConfig, workingDir string, opts cliOptions, themeMode tui.ThemeMode) error { + if opts.plain { + return repl(ctx, assistant, modelManager) + } + + err := tui.Run(ctx, assistant, tui.Options{ + ModelName: providerConfig.Model, + Yolo: opts.yolo, + ThemeMode: themeMode, + ModelManager: modelManager, + WorkingDir: workingDir, + }) + printExitSession(modelManager) + return err +} diff --git a/internal/app/bootstrap.go b/internal/app/bootstrap.go new file mode 100644 index 0000000..ffe760b --- /dev/null +++ b/internal/app/bootstrap.go @@ -0,0 +1,38 @@ +package app + +import ( + "fmt" + "os" + "path/filepath" + + "agentu/pkg/config" +) + +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 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 +} diff --git a/main_test.go b/internal/app/main_test.go similarity index 99% rename from main_test.go rename to internal/app/main_test.go index e69e746..7545900 100644 --- a/main_test.go +++ b/internal/app/main_test.go @@ -1,4 +1,4 @@ -package main +package app import ( "context" diff --git a/internal/app/repl.go b/internal/app/repl.go new file mode 100644 index 0000000..c85cdaf --- /dev/null +++ b/internal/app/repl.go @@ -0,0 +1,60 @@ +package app + +import ( + "bufio" + "context" + "fmt" + "os" + "strings" + + "agentu/internal/agent" + "agentu/internal/session" +) + +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() + } +} diff --git a/internal/app/session_startup.go b/internal/app/session_startup.go new file mode 100644 index 0000000..f03b2cb --- /dev/null +++ b/internal/app/session_startup.go @@ -0,0 +1,38 @@ +package app + +import ( + "fmt" + "os" + + "agentu/internal/session" +) + +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 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) + } +} diff --git a/main.go b/main.go index 47741d9..4976a42 100644 --- a/main.go +++ b/main.go @@ -1,228 +1,15 @@ 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" + "agentu/internal/app" ) -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 { + if err := app.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) - } -}