rony-llm-agent/pkg/agent/loop.go

package agent

import (
	"context"
	"encoding/json"
	"fmt"
	"iter"
	"strings"
	"time"

	"github.com/VictorVargas/rony-llm-agent/pkg/llm"
	"github.com/VictorVargas/rony-llm-agent/pkg/persona"
	"github.com/VictorVargas/rony-llm-agent/pkg/tools"
)

// MaxIterations is the default maximum number of agent iterations.
const DefaultMaxIterations = 50

// MaxToolOutputBytes is the default maximum size for tool output in bytes.
const DefaultMaxToolOutputBytes = 50 * 1024 // 50KB

// Sandbox defines the interface for filesystem sandbox operations.
type Sandbox interface {
	ValidateToolCall(tool tools.Tool, call llm.ToolCall) error
}

// Approver is called before executing tools with Ask permission.
// Return true to allow, false to deny.
type Approver func(tool tools.Tool, call llm.ToolCall) bool

// OnIterationHook is called after each iteration.
type OnIterationHook func(Iteration)

// Config holds the dependencies and settings for the agent loop.
type Config struct {
	LLM                llm.LLMClient
	Persona            persona.Persona
	Tools              tools.Registry
	Sandbox            Sandbox
	MaxIters           int
	Approver           Approver
	OnIteration        OnIterationHook
	ToolTimeout        time.Duration
	ChatTemplateKwargs map[string]any // passed to the LLM provider (e.g. Qwen enable_thinking)
	AgentsMD           string         // discovered AGENTS.md content, folded into the system prompt
}

// Iteration represents a single cycle of the agent loop.
type Iteration struct {
	Number    int
	ToolCalls []llm.ToolCall
	ToolsUsed int
	Duration  time.Duration
}

// Response is the final output of the agent loop.
type Response struct {
	Content    string
	ToolCalls  []llm.ToolCall
	Iterations int
	Duration   time.Duration
	TokenUsage llm.TokenUsage
}

// Loop is the main agent loop that orchestrates LLM calls and tool execution.
type Loop struct {
	cfg Config
}

// New creates a new Loop with the given configuration.
func New(cfg Config) *Loop {
	if cfg.MaxIters == 0 {
		cfg.MaxIters = DefaultMaxIterations
	}
	return &Loop{cfg: cfg}
}

// Run executes the agent loop and returns the final response.
// Optional history messages are appended after the system prompt and before
// the new user input.
func (l *Loop) Run(ctx context.Context, input string, history ...llm.Message) (Response, error) {
	start := time.Now()

	messages := l.buildInitialMessages(input, history)
	// Tool schemas don't change between iterations, so build the JSON once
	// per Run instead of re-marshaling every tool on every loop pass.
	toolSchemas := l.getToolSchemas()
	var finalContent string
	var allToolCalls []llm.ToolCall
	var totalUsage llm.TokenUsage
	iterations := 0
	completed := false

	for iterations < l.cfg.MaxIters {
		iterations++

		resp, err := l.cfg.LLM.Generate(ctx, llm.CompletionRequest{
			Messages:           messages,
			Tools:              toolSchemas,
			ChatTemplateKwargs: l.cfg.ChatTemplateKwargs,
		})
		if err != nil {
			return Response{}, fmt.Errorf("LLM generate failed: %w", err)
		}

		totalUsage.InputTokens += resp.Usage.InputTokens
		totalUsage.OutputTokens += resp.Usage.OutputTokens
		totalUsage.TotalTokens += resp.Usage.TotalTokens

		if len(resp.ToolCalls) == 0 {
			finalContent = resp.Content
			completed = true
			break
		}

		// Record the assistant's own turn (including which tools it asked
		// for) before the results, so the next request has a coherent
		// assistant-tool_calls / tool-result pair instead of a dangling
		// tool message the model can't attribute to anything.
		messages = append(messages, llm.Message{
			Role:      llm.RoleAssistant,
			Content:   resp.Content,
			ToolCalls: resp.ToolCalls,
		})

		for _, call := range resp.ToolCalls {
			result, err := l.executeTool(ctx, call)
			if err != nil {
				messages = append(messages, llm.Message{
					Role:       llm.RoleTool,
					ToolCallID: call.ID,
					Content:    fmt.Sprintf("Error: %v", err),
				})
				continue
			}
			messages = append(messages, llm.Message{
				Role:       llm.RoleTool,
				ToolCallID: call.ID,
				Content:    result.Content,
			})
			allToolCalls = append(allToolCalls, call)
		}
	}

	duration := time.Since(start)
	// Only report the max-iterations failure when the loop actually ran out
	// of budget without producing a final answer — an answer that arrives
	// exactly on the last allowed iteration is still a success (the old
	// `iterations >= MaxIters` check threw that valid response away).
	if !completed {
		return Response{}, fmt.Errorf("max iterations (%d) reached", l.cfg.MaxIters)
	}

	return Response{
		Content:    finalContent,
		ToolCalls:  allToolCalls,
		Iterations: iterations,
		Duration:   duration,
		TokenUsage: totalUsage,
	}, nil
}

// RunStream executes the agent loop with streaming output.
// Optional history messages are appended after the system prompt and before
// the new user input.
func (l *Loop) RunStream(ctx context.Context, input string, history ...llm.Message) iter.Seq2[llm.StreamChunk, error] {
	return func(yield func(llm.StreamChunk, error) bool) {
		messages := l.buildInitialMessages(input, history)
		// Same as Run: the schemas are identical on every iteration.
		toolSchemas := l.getToolSchemas()
		iterations := 0

		for iterations < l.cfg.MaxIters {
			iterations++

			stream := l.cfg.LLM.Stream(ctx, llm.CompletionRequest{
				Messages:           messages,
				Tools:              toolSchemas,
				ChatTemplateKwargs: l.cfg.ChatTemplateKwargs,
			})

			var hasToolCalls bool
			var responseBuilder strings.Builder
			for chunk, err := range stream {
				if err != nil {
					yield(llm.StreamChunk{}, err)
					return
				}

				if len(chunk.ToolCalls) > 0 {
					hasToolCalls = true

					// Same reasoning as in Run: without recording the
					// assistant's own tool_calls turn first, the tool
					// results that follow have nothing for the model to
					// attribute them to on the next request.
					messages = append(messages, llm.Message{
						Role:      llm.RoleAssistant,
						Content:   responseBuilder.String(),
						ToolCalls: chunk.ToolCalls,
					})

					for _, tc := range chunk.ToolCalls {
						result, err := l.executeTool(ctx, tc)
						if err != nil {
							messages = append(messages, llm.Message{
								Role:       llm.RoleTool,
								ToolCallID: tc.ID,
								Content:    fmt.Sprintf("Error: %v", err),
							})
							continue
						}
						messages = append(messages, llm.Message{
							Role:       llm.RoleTool,
							ToolCallID: tc.ID,
							Content:    result.Content,
						})
					}

					// Surface which tools were actually called, and with
					// what arguments, to the caller — a dedicated chunk,
					// separate from the content-streaming gate below, since
					// that gate exists to hide raw provider deltas during a
					// tool-call round, not to hide the fact that a call
					// happened at all. Without this, callers (e.g. a UI
					// wanting to show "used tool X" or track which files a
					// write/edit touched) have no way to observe tool
					// calls unless they also happen to be the Approver.
					if !yield(llm.StreamChunk{ToolCalls: chunk.ToolCalls}, nil) {
						return
					}
				}

				// A trailing usage-only chunk (no Delta/ReasoningDelta, per
				// providers that report token usage in a separate final
				// event) must still be forwarded, or callers can never see
				// real token counts.
				hasContent := chunk.Delta != "" || chunk.ReasoningDelta != ""
				hasUsage := chunk.Usage.TotalTokens > 0
				switch {
				case !hasToolCalls && (hasContent || hasUsage):
					responseBuilder.WriteString(chunk.Delta)
					if !yield(chunk, nil) {
						return
					}
				case hasToolCalls && hasUsage:
					// The content gate above exists to hide raw provider
					// deltas during a tool-call round, but it also swallowed
					// that round's token usage — so callers tracking context
					// occupancy (e.g. a UI's context bar deciding when to
					// compact) only ever saw the usage of the final,
					// tool-free round. Forward the usage on its own,
					// without the content.
					if !yield(llm.StreamChunk{Usage: chunk.Usage}, nil) {
						return
					}
				}
			}

			if !hasToolCalls {
				return
			}
		}

		yield(llm.StreamChunk{}, fmt.Errorf("max iterations (%d) reached", l.cfg.MaxIters))
	}
}

func (l *Loop) buildInitialMessages(input string, history []llm.Message) []llm.Message {
	systemPrompt := persona.AssembleSystemPrompt(l.cfg.Persona, l.cfg.AgentsMD)
	messages := make([]llm.Message, 0, len(history)+2)
	messages = append(messages, llm.Message{Role: llm.RoleSystem, Content: systemPrompt})
	messages = append(messages, history...)
	messages = append(messages, llm.Message{Role: llm.RoleUser, Content: input})
	return messages
}

func (l *Loop) getToolSchemas() []json.RawMessage {
	var schemas []json.RawMessage
	for _, tool := range l.cfg.Tools.List() {
		def := map[string]interface{}{
			"type": "function",
			"function": map[string]interface{}{
				"name":        tool.Name,
				"description": tool.Description,
				"parameters":  json.RawMessage(tool.InputSchema),
			},
		}
		data, _ := json.Marshal(def)
		schemas = append(schemas, data)
	}
	return schemas
}

func (l *Loop) executeTool(ctx context.Context, call llm.ToolCall) (tools.ToolResult, error) {
	tool, found := l.cfg.Tools.Get(call.Name)
	if !found {
		return tools.ToolResult{IsError: true}, fmt.Errorf("tool %q not found", call.Name)
	}

	if tool.Permission == tools.Ask && l.cfg.Approver != nil && !l.cfg.Approver(tool, call) {
		return tools.ToolResult{Content: "Tool execution denied by user"}, nil
	}

	if l.cfg.Sandbox != nil {
		if err := l.cfg.Sandbox.ValidateToolCall(tool, call); err != nil {
			return tools.ToolResult{IsError: true}, fmt.Errorf("sandbox violation: %w", err)
		}
	}

	if l.cfg.OnIteration != nil {
		l.cfg.OnIteration(Iteration{
			Number:    1,
			ToolCalls: []llm.ToolCall{call},
			ToolsUsed: 1,
		})
	}

	result, err := tool.Handler(ctx, call.Arguments)
	if err != nil {
		return tools.ToolResult{IsError: true, Content: fmt.Sprintf("Tool error: %v", err)}, nil
	}

	return result, nil
}