agentu/internal/agent/agent.go

package agent

import (
	"context"
	"encoding/json"
	"fmt"
	"io"
	"sort"
	"strings"
	"sync"
	"time"

	"agentu/internal/tools"
	"agentu/pkg/llm"
)

const (
	defaultMaxToolRounds      = 50
	doomLoopThreshold         = 3
	compactMinRecentMessages      = 6
	compactRecentRetentionPercent = 20
	summaryMessageLimit       = 8000
	contextSummaryPrefix      = "[context summary]"
	toolLogArgumentLimit      = 400
	toolLogOutputPreviewBytes = 4096
	toolLogOutputMarker       = "[output]"

	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 toolExecutionResult struct {
	call   llm.ToolCall
	output string
}

type Agent struct {
	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 {
	Provider          llm.Provider
	ProviderName      string
	Model             string
	Thinking          string
	ThinkingParam     string
	ThinkingEnabled   bool
	SystemPrompt      string
	ToolRegistry      *tools.Registry
	ToolTimeout       time.Duration
	MaxToolIterations int
	MaxContextTokens  int
}

func New(opts Options) *Agent {
	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,
		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
	MaxContextTokens int
}

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

func (a *Agent) SetRuntime(opts RuntimeOptions) {
	if opts.Provider != nil {
		a.provider = opts.Provider
	}
	if opts.ProviderName != "" {
		a.providerName = opts.ProviderName
	}
	if opts.Model != "" {
		a.model = opts.Model
	}
	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,
		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{
			Role:    llm.RoleSystem,
			Content: a.systemPrompt,
		})
	}
}

// Messages returns a copy of the current conversation messages.
func (a *Agent) Messages() []llm.Message {
	out := make([]llm.Message, len(a.messages))
	copy(out, a.messages)
	return out
}

// SetMessages replaces the conversation messages (used for session restore).
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
}
// 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.
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
	}
	recentStart := compactRecentStart(a.messages, prefixEnd)
	if recentStart <= prefixEnd || recentStart >= len(a.messages) {
		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)
			}
		}
	}

	var recentSignatures []string

	for round := 0; round < a.maxToolRounds; round++ {
		content, calls, err := a.streamAssistant(ctx, out)
		if err != nil {
			return err
		}
		if len(calls) == 0 {
			a.messages = append(a.messages, llm.Message{
				Role:    llm.RoleAssistant,
				Content: content,
			})
			return nil
		}

		a.messages = append(a.messages, llm.Message{
			Role:      llm.RoleAssistant,
			Content:   content,
			ToolCalls: calls,
		})
		if content != "" {
			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)
			}
		}

		results := a.executeTools(ctx, calls, logs)
		for _, result := range results {
			a.messages = append(a.messages, llm.Message{
				Role:       llm.RoleTool,
				ToolCallID: result.call.ID,
				Content:    result.output,
			})
		}

	}

	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) {
	var content strings.Builder
	calls := make(map[int]*toolCallBuilder)

	req := llm.ChatRequest{
		Model:    a.model,
		Messages: append([]llm.Message(nil), a.messages...),
		Tools:    a.toolDefinitions(),
	}
	if a.thinkingEnabled {
		thinkingParam := a.thinkingParam
		if thinkingParam == "" {
			thinkingParam = "thinking"
		}
		req.Extra = map[string]any{thinkingParam: a.thinking}
	}
	if len(req.Tools) > 0 {
		req.ToolChoice = "auto"
	}

	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 {
				if _, err := io.WriteString(out, event.Content); err != nil {
					return err
				}
			}
		}
		for _, delta := range event.ToolCalls {
			builder := calls[delta.Index]
			if builder == nil {
				builder = &toolCallBuilder{index: delta.Index}
				calls[delta.Index] = builder
			}
			builder.apply(delta)
		}
		return nil
	})
	if err != nil {
		return "", nil, err
	}

	return content.String(), finalizeToolCalls(calls), nil
}

func (a *Agent) toolDefinitions() []llm.Tool {
	if a.toolRegistry == nil {
		return nil
	}
	return a.toolRegistry.Definitions()
}

func (a *Agent) executeTool(ctx context.Context, call llm.ToolCall) string {
	name := call.Function.Name

	if a.toolRegistry == nil {
		return tools.FormatError(fmt.Errorf("tool registry is disabled"))
	}
	tool, ok := a.toolRegistry.Get(name)
	if !ok {
		return tools.FormatError(fmt.Errorf("unknown tool: %s", name))
	}

	args := strings.TrimSpace(call.Function.Arguments)
	if args == "" {
		args = "{}"
	}

	toolCtx, cancel := context.WithTimeout(ctx, a.toolTimeout)
	defer cancel()
	output, err := tool.Execute(toolCtx, json.RawMessage(args))
	if err != nil {
		if output != "" {
			return output + "\n" + tools.FormatError(err)
		}
		return tools.FormatError(err)
	}
	return output
}

func (a *Agent) hasMutatingCall(calls []llm.ToolCall) bool {
	if a.toolRegistry == nil {
		return false
	}
	for _, call := range calls {
		if tool, ok := a.toolRegistry.Get(call.Function.Name); ok {
			if tools.IsMutating(tool) {
				return true
			}
		}
	}
	return false
}

func (a *Agent) executeTools(ctx context.Context, calls []llm.ToolCall, logs io.Writer) []toolExecutionResult {
	results := make([]toolExecutionResult, len(calls))

	// Log all tool starts sequentially so TUI gets ordered placeholders.
	if logs != nil {
		for _, call := range calls {
			logToolStart(logs, call)
		}
	}

	if a.hasMutatingCall(calls) {
		// Sequential: avoid racing mutating tools.
		for i, call := range calls {
			output := a.executeTool(ctx, call)
			results[i] = toolExecutionResult{call: call, output: output}
			if logs != nil {
				logToolOutput(logs, call, output)
			}
		}
		return results
	}

	// Concurrent: independent read-only tools.
	var mu sync.Mutex
	var wg sync.WaitGroup
	for i, call := range calls {
		wg.Add(1)
		go func(idx int, c llm.ToolCall) {
			defer wg.Done()
			output := a.executeTool(ctx, c)
			results[idx] = toolExecutionResult{call: c, output: output}
			if logs != nil {
				mu.Lock()
				logToolOutput(logs, c, output)
				mu.Unlock()
			}
		}(i, call)
	}
	wg.Wait()
	return results
}

func logToolStart(logs io.Writer, call llm.ToolCall) {
	fmt.Fprintf(logs, "\n[tool] %s %s\n", call.Function.Name, compact(call.Function.Arguments, toolLogArgumentLimit))
}

func logToolOutput(logs io.Writer, call llm.ToolCall, output string) {
	fmt.Fprintf(logs, "\n[tool] %s %s\n%s\n%s\n", call.Function.Name, compact(call.Function.Arguments, toolLogArgumentLimit), toolLogOutputMarker, toolOutputPreview(output))
}

func toolOutputPreview(output string) string {
	output = strings.TrimSpace(output)
	if output == "" {
		return "(no output)"
	}
	if len(output) <= toolLogOutputPreviewBytes {
		return output
	}
	return output[:toolLogOutputPreviewBytes] + fmt.Sprintf("\n[truncated: showing first %d bytes; omitted %d bytes]", toolLogOutputPreviewBytes, len(output)-toolLogOutputPreviewBytes)
}

type toolCallBuilder struct {
	index     int
	id        string
	callType  string
	name      string
	arguments strings.Builder
}

func (b *toolCallBuilder) apply(delta llm.ToolCallDelta) {
	if delta.ID != "" {
		b.id = delta.ID
	}
	if delta.Type != "" {
		b.callType = delta.Type
	}
	if delta.Name != "" {
		b.name = delta.Name
	}
	if delta.Arguments != "" {
		b.arguments.WriteString(delta.Arguments)
	}
}

func finalizeToolCalls(builders map[int]*toolCallBuilder) []llm.ToolCall {
	if len(builders) == 0 {
		return nil
	}
	indexes := make([]int, 0, len(builders))
	for index := range builders {
		indexes = append(indexes, index)
	}
	sort.Ints(indexes)

	calls := make([]llm.ToolCall, 0, len(indexes))
	for _, index := range indexes {
		builder := builders[index]
		id := builder.id
		if id == "" {
			id = fmt.Sprintf("call_%d", index)
		}
		callType := builder.callType
		if callType == "" {
			callType = "function"
		}
		calls = append(calls, llm.ToolCall{
			ID:   id,
			Type: callType,
			Function: llm.FunctionCall{
				Name:      builder.name,
				Arguments: builder.arguments.String(),
			},
		})
	}
	return calls
}

func compact(value string, limit int) string {
	value = strings.TrimSpace(value)
	if len(value) <= limit {
		return value
	}
	return value[:limit] + "...[truncated]"
}