v0.0.5: context compaction, minimal config, retry/search/code symbols/diff/status/test commands

2026-06-24 02:20:19 -07:00
parent e4c75c7c0e
commit 950c63adaa
24 changed files with 1939 additions and 220 deletions
@@ -13,20 +13,36 @@ import (
 	"agentu/internal/tools"
 )

-const defaultMaxToolIterations = 8
+const (
+	defaultMaxToolRounds  = 50
+	doomLoopThreshold     = 3
+	compactRecentMessages = 6
+	summaryMessageLimit   = 8000
+	contextSummaryPrefix  = "[context summary]"
+
+	summarizerSystemPrompt = "You are a conversation summarizer. Summarize the following conversation history into a concise but detailed summary. Preserve:\n" +
+		"- Key facts, decisions, and conclusions\n" +
+		"- File paths and code changes made\n" +
+		"- Tool call results (especially errors)\n" +
+		"- Current task state and next steps\n\n" +
+		"Do NOT include: tool call syntax, raw JSON, or verbatim file contents.\n" +
+		"Output only the summary, no preamble."
+)

 type Agent struct {
-	provider          llm.Provider
-	providerName      string
-	model             string
-	thinking          string
-	thinkingParam     string
-	thinkingEnabled   bool
-	systemPrompt      string
-	toolRegistry      *tools.Registry
-	toolTimeout       time.Duration
-	maxToolIterations int
-	messages          []llm.Message
+	provider         llm.Provider
+	providerName     string
+	model            string
+	thinking         string
+	thinkingParam    string
+	thinkingEnabled  bool
+	systemPrompt     string
+	toolRegistry     *tools.Registry
+	toolTimeout      time.Duration
+	maxToolRounds    int
+	maxContextTokens int
+	lastUsage        *llm.Usage
+	messages         []llm.Message
 }

 type Options struct {
@@ -40,47 +56,52 @@ type Options struct {
 	ToolRegistry      *tools.Registry
 	ToolTimeout       time.Duration
 	MaxToolIterations int
+	MaxContextTokens  int
 }

 func New(opts Options) *Agent {
-	maxToolIterations := opts.MaxToolIterations
-	if maxToolIterations <= 0 {
-		maxToolIterations = defaultMaxToolIterations
+	maxToolRounds := opts.MaxToolIterations
+	if maxToolRounds <= 0 {
+		maxToolRounds = defaultMaxToolRounds
 	}
 	toolTimeout := opts.ToolTimeout
 	if toolTimeout <= 0 {
 		toolTimeout = 30 * time.Second
 	}
 	a := &Agent{
-		provider:          opts.Provider,
-		providerName:      opts.ProviderName,
-		model:             opts.Model,
-		thinking:          opts.Thinking,
-		thinkingParam:     opts.ThinkingParam,
-		thinkingEnabled:   opts.ThinkingEnabled,
-		systemPrompt:      opts.SystemPrompt,
-		toolRegistry:      opts.ToolRegistry,
-		toolTimeout:       toolTimeout,
-		maxToolIterations: maxToolIterations,
+		provider:         opts.Provider,
+		providerName:     opts.ProviderName,
+		model:            opts.Model,
+		thinking:         opts.Thinking,
+		thinkingParam:    opts.ThinkingParam,
+		thinkingEnabled:  opts.ThinkingEnabled,
+		systemPrompt:     opts.SystemPrompt,
+		toolRegistry:     opts.ToolRegistry,
+		toolTimeout:      toolTimeout,
+		maxToolRounds:    maxToolRounds,
+		maxContextTokens: opts.MaxContextTokens,
 	}
 	a.Clear()
 	return a
 }

 type RuntimeOptions struct {
-	Provider        llm.Provider
-	ProviderName    string
-	Model           string
-	Thinking        string
-	ThinkingParam   string
-	ThinkingEnabled bool
+	Provider         llm.Provider
+	ProviderName     string
+	Model            string
+	Thinking         string
+	ThinkingParam    string
+	ThinkingEnabled  bool
+	MaxContextTokens int
 }

 type RuntimeInfo struct {
-	ProviderName    string
-	Model           string
-	Thinking        string
-	ThinkingEnabled bool
+	ProviderName     string
+	Model            string
+	Thinking         string
+	ThinkingEnabled  bool
+	ContextTokens    int
+	MaxContextTokens int
 }

 func (a *Agent) SetRuntime(opts RuntimeOptions) {
@@ -96,18 +117,22 @@ func (a *Agent) SetRuntime(opts RuntimeOptions) {
 	a.thinking = opts.Thinking
 	a.thinkingParam = opts.ThinkingParam
 	a.thinkingEnabled = opts.ThinkingEnabled
+	a.maxContextTokens = opts.MaxContextTokens
 }

 func (a *Agent) RuntimeInfo() RuntimeInfo {
 	return RuntimeInfo{
-		ProviderName:    a.providerName,
-		Model:           a.model,
-		Thinking:        a.thinking,
-		ThinkingEnabled: a.thinkingEnabled,
+		ProviderName:     a.providerName,
+		Model:            a.model,
+		Thinking:         a.thinking,
+		ThinkingEnabled:  a.thinkingEnabled,
+		ContextTokens:    llm.EstimateTokens(a.messages, a.toolDefinitions()),
+		MaxContextTokens: a.maxContextTokens,
 	}
 }

 func (a *Agent) Clear() {
+	a.lastUsage = nil
 	a.messages = nil
 	if strings.TrimSpace(a.systemPrompt) != "" {
 		a.messages = append(a.messages, llm.Message{
@@ -129,14 +154,161 @@ func (a *Agent) SetMessages(messages []llm.Message) {
 	a.messages = append([]llm.Message(nil), messages...)
 }

+// LastUsage returns a copy of the most recent provider token usage data.
+func (a *Agent) LastUsage() *llm.Usage {
+	if a.lastUsage == nil {
+		return nil
+	}
+	usage := *a.lastUsage
+	return &usage
+}
+
+// SetLastUsage restores the most recent provider token usage data.
+func (a *Agent) SetLastUsage(usage *llm.Usage) {
+	if usage == nil {
+		a.lastUsage = nil
+		return
+	}
+	copy := *usage
+	a.lastUsage = &copy
+}
+
+// Compact summarizes older conversation messages while preserving the system
+// prompt and recent messages intact.
+func (a *Agent) Compact(ctx context.Context) error {
+	if len(a.messages) < 4 {
+		return nil
+	}
+	if a.provider == nil {
+		return fmt.Errorf("provider is not configured")
+	}
+
+	prefixEnd := 0
+	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 {
+		return nil
+	}
+	middle := append([]llm.Message(nil), a.messages[prefixEnd:recentStart]...)
+	if len(middle) == 0 {
+		return nil
+	}
+
+	summary, usage, err := a.summarizeMessages(ctx, middle)
+	if err != nil {
+		return fmt.Errorf("compact context: %w", err)
+	}
+	summary = strings.TrimSpace(summary)
+	if summary == "" {
+		return fmt.Errorf("compact context: summary was empty")
+	}
+
+	compressed := make([]llm.Message, 0, prefixEnd+1+len(a.messages[recentStart:]))
+	compressed = append(compressed, a.messages[:prefixEnd]...)
+	compressed = append(compressed, llm.Message{
+		Role:    llm.RoleAssistant,
+		Content: contextSummaryPrefix + "\n" + summary,
+	})
+	compressed = append(compressed, a.messages[recentStart:]...)
+	a.messages = compressed
+	if usage != nil {
+		a.lastUsage = usage
+	}
+	return nil
+}
+
+func (a *Agent) shouldCompact() bool {
+	if a.maxContextTokens <= 0 {
+		return false
+	}
+	if a.lastUsage != nil && a.lastUsage.TotalTokens >= a.maxContextTokens {
+		return true
+	}
+	return llm.EstimateTokens(a.messages, a.toolDefinitions()) >= a.maxContextTokens
+}
+
+func (a *Agent) summarizeMessages(ctx context.Context, messages []llm.Message) (string, *llm.Usage, error) {
+	req := llm.ChatRequest{
+		Model: a.model,
+		Messages: []llm.Message{
+			{Role: llm.RoleSystem, Content: summarizerSystemPrompt},
+			{Role: llm.RoleUser, Content: formatMessagesForSummary(messages)},
+		},
+	}
+
+	var summary strings.Builder
+	var usage *llm.Usage
+	if err := a.provider.ChatStream(ctx, req, func(event llm.StreamEvent) error {
+		if event.Content != "" {
+			summary.WriteString(event.Content)
+		}
+		if event.Usage != nil {
+			copy := *event.Usage
+			usage = &copy
+		}
+		return nil
+	}); err != nil {
+		return "", nil, err
+	}
+
+	return strings.TrimSpace(summary.String()), usage, nil
+}
+
+func formatMessagesForSummary(messages []llm.Message) string {
+	var b strings.Builder
+	b.WriteString("Conversation history to summarize:\n\n")
+	for index, msg := range messages {
+		fmt.Fprintf(&b, "Message %d (%s):\n", index+1, msg.Role)
+		if msg.ToolCallID != "" {
+			fmt.Fprintf(&b, "tool_call_id: %s\n", msg.ToolCallID)
+		}
+		if msg.Content != "" {
+			fmt.Fprintf(&b, "%s\n", compact(msg.Content, summaryMessageLimit))
+		}
+		for _, call := range msg.ToolCalls {
+			fmt.Fprintf(&b, "tool_call: %s %s\n", call.Function.Name, compact(call.Function.Arguments, summaryMessageLimit))
+		}
+		b.WriteByte('\n')
+	}
+	return b.String()
+}
+
+// roundSignature returns a deterministic signature for a set of tool calls
+// in a single round, used for doom loop detection.
+func roundSignature(calls []llm.ToolCall) string {
+	var b strings.Builder
+	for _, c := range calls {
+		b.WriteString(c.Function.Name)
+		b.WriteByte('\x00')
+		b.WriteString(strings.TrimSpace(c.Function.Arguments))
+		b.WriteByte('\x00')
+	}
+	return b.String()
+}
+
 func (a *Agent) RunTurn(ctx context.Context, input string, out io.Writer, logs io.Writer) error {
 	input = strings.TrimSpace(input)
 	if input == "" {
 		return nil
 	}
 	a.messages = append(a.messages, llm.Message{Role: llm.RoleUser, Content: input})
+	if a.shouldCompact() {
+		if err := a.Compact(ctx); err != nil {
+			if logs != nil {
+				fmt.Fprintf(logs, "[compact] warning: %v\n", err)
+			}
+		}
+	}

-	for iteration := 0; iteration <= a.maxToolIterations; iteration++ {
+	var recentSignatures []string
+
+	for round := 0; round < a.maxToolRounds; round++ {
 		content, calls, err := a.streamAssistant(ctx, out)
 		if err != nil {
 			return err
@@ -158,6 +330,40 @@ func (a *Agent) RunTurn(ctx context.Context, input string, out io.Writer, logs i
 			fmt.Fprintln(out)
 		}

+		// Doom loop detection: check if the last N rounds all called
+		// the exact same tools with the exact same arguments.
+		sig := roundSignature(calls)
+		recentSignatures = append(recentSignatures, sig)
+		if len(recentSignatures) > doomLoopThreshold {
+			recentSignatures = recentSignatures[len(recentSignatures)-doomLoopThreshold:]
+		}
+
+		if len(recentSignatures) == doomLoopThreshold {
+			allSame := true
+			for i := 1; i < len(recentSignatures); i++ {
+				if recentSignatures[i] != recentSignatures[0] {
+					allSame = false
+					break
+				}
+			}
+			if allSame {
+				toolNames := make([]string, 0, len(calls))
+				for _, c := range calls {
+					toolNames = append(toolNames, c.Function.Name)
+				}
+				msg := fmt.Sprintf(
+					"error: doom loop detected — the same tool call (%s) with identical arguments has been repeated %d times. Stopping. Try a different approach or break the task into smaller steps.",
+					strings.Join(toolNames, ", "),
+					doomLoopThreshold,
+				)
+				a.messages = append(a.messages, llm.Message{
+					Role:    llm.RoleAssistant,
+					Content: msg,
+				})
+				return fmt.Errorf("%s", msg)
+			}
+		}
+
 		for _, call := range calls {
 			result := a.executeTool(ctx, call, logs)
 			a.messages = append(a.messages, llm.Message{
@@ -168,7 +374,7 @@ func (a *Agent) RunTurn(ctx context.Context, input string, out io.Writer, logs i
 		}
 	}

-	return fmt.Errorf("tool iteration limit reached: %d", a.maxToolIterations)
+	return fmt.Errorf("tool round limit reached after %d rounds; the task may be too complex for a single turn", a.maxToolRounds)
 }

 func (a *Agent) streamAssistant(ctx context.Context, out io.Writer) (string, []llm.ToolCall, error) {
@@ -192,6 +398,9 @@ func (a *Agent) streamAssistant(ctx context.Context, out io.Writer) (string, []l
 	}

 	err := a.provider.ChatStream(ctx, req, func(event llm.StreamEvent) error {
+		if event.Usage != nil {
+			a.SetLastUsage(event.Usage)
+		}
 		if event.Content != "" {
 			content.WriteString(event.Content)
 			if out != nil {
@@ -3,6 +3,7 @@ package agent
 import (
 	"context"
 	"encoding/json"
+	"fmt"
 	"strings"
 	"testing"

@@ -96,6 +97,134 @@ func (p *contentProvider) ChatStream(ctx context.Context, req llm.ChatRequest, e
 	return emit(llm.StreamEvent{Content: p.text})
 }

+type compactProvider struct {
+	summaryCalls int
+	answerCalls  int
+	lastSummary  string
+}
+
+func (p *compactProvider) ChatStream(ctx context.Context, req llm.ChatRequest, emit func(llm.StreamEvent) error) error {
+	if len(req.Messages) == 2 && req.Messages[0].Role == llm.RoleSystem && req.Messages[0].Content == summarizerSystemPrompt {
+		p.summaryCalls++
+		p.lastSummary = req.Messages[1].Content
+		if err := emit(llm.StreamEvent{Content: "old work summarized"}); err != nil {
+			return err
+		}
+		return emit(llm.StreamEvent{Usage: &llm.Usage{PromptTokens: 10, CompletionTokens: 2, TotalTokens: 12}})
+	}
+
+	p.answerCalls++
+	return emit(llm.StreamEvent{Content: "done"})
+}
+
+func TestAgentCompact(t *testing.T) {
+	provider := &compactProvider{}
+	a := New(Options{Provider: provider, Model: "test", SystemPrompt: "system"})
+	a.SetMessages(longConversation())
+
+	if err := a.Compact(context.Background()); err != nil {
+		t.Fatal(err)
+	}
+
+	messages := a.Messages()
+	if len(messages) != 8 {
+		t.Fatalf("messages = %d, want 8", len(messages))
+	}
+	if messages[0].Role != llm.RoleSystem || messages[0].Content != "system" {
+		t.Fatalf("system message = %#v", messages[0])
+	}
+	if messages[1].Role != llm.RoleAssistant || !strings.HasPrefix(messages[1].Content, contextSummaryPrefix) {
+		t.Fatalf("summary message = %#v", messages[1])
+	}
+	if !strings.Contains(messages[1].Content, "old work summarized") {
+		t.Fatalf("summary content = %q", messages[1].Content)
+	}
+	if !strings.Contains(provider.lastSummary, "old-user-0") {
+		t.Fatalf("summary input did not include old history: %q", provider.lastSummary)
+	}
+	if messages[len(messages)-1].Content != "recent-answer-2" {
+		t.Fatalf("last recent message = %#v", messages[len(messages)-1])
+	}
+	if usage := a.LastUsage(); usage == nil || usage.TotalTokens != 12 {
+		t.Fatalf("usage = %#v", usage)
+	}
+}
+
+func TestAgentAutoCompact(t *testing.T) {
+	provider := &compactProvider{}
+	a := New(Options{
+		Provider:          provider,
+		Model:             "test",
+		SystemPrompt:      "system",
+		MaxContextTokens:  1,
+		ToolRegistry:      tools.NewRegistry(),
+		MaxToolIterations: 1,
+	})
+	a.SetMessages(longConversation())
+
+	var out strings.Builder
+	if err := a.RunTurn(context.Background(), "new question", &out, nil); err != nil {
+		t.Fatal(err)
+	}
+	if provider.summaryCalls != 1 {
+		t.Fatalf("summaryCalls = %d, want 1", provider.summaryCalls)
+	}
+	if provider.answerCalls != 1 {
+		t.Fatalf("answerCalls = %d, want 1", provider.answerCalls)
+	}
+	if out.String() != "done" {
+		t.Fatalf("out = %q", out.String())
+	}
+	if !strings.Contains(joinMessageContents(a.Messages()), contextSummaryPrefix) {
+		t.Fatalf("messages missing summary: %#v", a.Messages())
+	}
+}
+
+func TestAgentCompactSkipsShortHistory(t *testing.T) {
+	provider := &contentProvider{text: "unused"}
+	a := New(Options{Provider: provider, Model: "test", SystemPrompt: "system"})
+	a.SetMessages([]llm.Message{
+		{Role: llm.RoleSystem, Content: "system"},
+		{Role: llm.RoleUser, Content: "question"},
+		{Role: llm.RoleAssistant, Content: "answer"},
+	})
+
+	if err := a.Compact(context.Background()); err != nil {
+		t.Fatal(err)
+	}
+	if provider.called {
+		t.Fatal("provider should not be called for short history")
+	}
+	if len(a.Messages()) != 3 {
+		t.Fatalf("messages = %d, want 3", len(a.Messages()))
+	}
+}
+
+func longConversation() []llm.Message {
+	messages := []llm.Message{{Role: llm.RoleSystem, Content: "system"}}
+	for i := 0; i < 3; i++ {
+		messages = append(messages,
+			llm.Message{Role: llm.RoleUser, Content: fmt.Sprintf("old-user-%d", i)},
+			llm.Message{Role: llm.RoleAssistant, Content: fmt.Sprintf("old-answer-%d", i)},
+		)
+	}
+	for i := 0; i < 3; i++ {
+		messages = append(messages,
+			llm.Message{Role: llm.RoleUser, Content: fmt.Sprintf("recent-user-%d", i)},
+			llm.Message{Role: llm.RoleAssistant, Content: fmt.Sprintf("recent-answer-%d", i)},
+		)
+	}
+	return messages
+}
+
+func joinMessageContents(messages []llm.Message) string {
+	var b strings.Builder
+	for _, msg := range messages {
+		b.WriteString(msg.Content)
+		b.WriteByte('\n')
+	}
+	return b.String()
+}
 func TestAgentLetsProviderDecideShellRequestWhenShellDisabled(t *testing.T) {
 	provider := &contentProvider{text: "shell is disabled"}
 	a := New(Options{
@@ -137,3 +266,77 @@ func TestAgentDoesNotBlockFileRequest(t *testing.T) {
 		t.Fatalf("out = %q", out.String())
 	}
 }
+
+// alwaysToolProvider always returns a tool call with the same args on every request.
+type alwaysToolProvider struct {
+	calls int
+}
+
+func (p *alwaysToolProvider) ChatStream(ctx context.Context, req llm.ChatRequest, emit func(llm.StreamEvent) error) error {
+	p.calls++
+	return emit(llm.StreamEvent{ToolCalls: []llm.ToolCallDelta{
+		{Index: 0, ID: fmt.Sprintf("call_%d", p.calls), Type: "function", Name: "test_echo", Arguments: `{"text":"loop"}`},
+	}})
+}
+
+func TestAgentStopsOnDoomLoop(t *testing.T) {
+	provider := &alwaysToolProvider{}
+	a := New(Options{
+		Provider:          provider,
+		Model:             "test",
+		SystemPrompt:      "system",
+		ToolRegistry:      tools.NewRegistry(echoTool{}),
+		MaxToolIterations: 50, // High hard limit — doom loop should trigger first
+	})
+
+	var out strings.Builder
+	err := a.RunTurn(context.Background(), "loop forever", &out, nil)
+	if err == nil {
+		t.Fatal("expected error when doom loop is detected")
+	}
+	if !strings.Contains(err.Error(), "doom loop detected") {
+		t.Fatalf("unexpected error: %v", err)
+	}
+	// Should have stopped at doomLoopThreshold (3) rounds, not reached hard limit.
+	if provider.calls != doomLoopThreshold {
+		t.Fatalf("calls = %d, want %d (doom loop threshold)", provider.calls, doomLoopThreshold)
+	}
+}
+
+// varyingToolProvider alternates between different tool calls to avoid doom loop detection.
+type varyingToolProvider struct {
+	calls int
+}
+
+func (p *varyingToolProvider) ChatStream(ctx context.Context, req llm.ChatRequest, emit func(llm.StreamEvent) error) error {
+	p.calls++
+	// Vary the arguments each round to avoid doom loop detection.
+	return emit(llm.StreamEvent{ToolCalls: []llm.ToolCallDelta{
+		{Index: 0, ID: fmt.Sprintf("call_%d", p.calls), Type: "function", Name: "test_echo",
+			Arguments: fmt.Sprintf(`{"text":"round-%d"}`, p.calls)},
+	}})
+}
+
+func TestAgentHardLimitStopsVaryingCalls(t *testing.T) {
+	provider := &varyingToolProvider{}
+	const limit = 5
+	a := New(Options{
+		Provider:          provider,
+		Model:             "test",
+		SystemPrompt:      "system",
+		ToolRegistry:      tools.NewRegistry(echoTool{}),
+		MaxToolIterations: limit,
+	})
+
+	var out strings.Builder
+	err := a.RunTurn(context.Background(), "vary forever", &out, nil)
+	if err == nil {
+		t.Fatal("expected error when hard limit is reached")
+	}
+	if !strings.Contains(err.Error(), "tool round limit reached") {
+		t.Fatalf("unexpected error: %v", err)
+	}
+	if provider.calls != limit {
+		t.Fatalf("calls = %d, want %d", provider.calls, limit)
+	}
+}