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. input's Role is overwritten to RoleUser regardless of
// what the caller sets, so callers only need to fill in Content/Parts.
func (l *Loop) Run(ctx context.Context, input llm.Message, 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
	nudges := 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 {
			// Same unparsed-tool-call recovery as RunStream: a tool call
			// written as plain text was never executed, so ending the turn
			// here would silently abandon the work mid-task.
			if nudges < maxUnparsedToolCallNudges && containsUnparsedToolCall(resp.Content+resp.Reasoning) {
				nudges++
				messages = append(messages,
					llm.Message{Role: llm.RoleAssistant, Content: resp.Content},
					llm.Message{Role: llm.RoleUser, Content: unparsedToolCallNudge},
				)
				continue
			}
			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. input's Role is overwritten to RoleUser regardless of
// what the caller sets, so callers only need to fill in Content/Parts.
func (l *Loop) RunStream(ctx context.Context, input llm.Message, 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
		nudges := 0
		budgetNudges := 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 budgetExceeded bool
			var responseBuilder strings.Builder
			// detectBuf collects this round's raw text (content AND
			// reasoning) only to spot tool calls the model wrote as plain
			// text — see the unparsed-tool-call recovery below the loop.
			var detectBuf strings.Builder
			for chunk, err := range stream {
				if err != nil {
					yield(llm.StreamChunk{}, err)
					return
				}
				if chunk.FinishReason == llm.FinishThinkingBudget {
					budgetExceeded = true
				}
				if detectBuf.Len() < unparsedDetectBudget {
					detectBuf.WriteString(chunk.ReasoningDelta)
					detectBuf.WriteString(chunk.Delta)
				}

				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 {
				// Recovery for a failure mode common with local models: the
				// model writes its tool call as plain text — typically
				// inside its reasoning block — so the server never parses
				// it into a real tool call. Ending the turn here (the old
				// behavior) silently abandons the work mid-task: the
				// transcript reads "now I'll update X:" and then... nothing,
				// because nothing was ever executed. Instead, tell the model
				// what happened and let it re-issue the call properly.
				if nudges < maxUnparsedToolCallNudges && containsUnparsedToolCall(detectBuf.String()) {
					nudges++
					messages = append(messages,
						llm.Message{Role: llm.RoleAssistant, Content: responseBuilder.String()},
						llm.Message{Role: llm.RoleUser, Content: unparsedToolCallNudge},
					)
					continue
				}
				// The provider cut this round because the model exceeded its
				// thinking budget without ever starting an answer or a tool
				// call (reasoning spiral). Ending the turn here would abandon
				// the task with nothing to show for it — instead tell the
				// model its reasoning was cut and demand direct action. Its
				// own nudge counter, so a spiral doesn't consume the
				// unparsed-tool-call retries (or vice versa).
				if budgetNudges < maxThinkingBudgetNudges && budgetExceeded {
					budgetNudges++
					content := responseBuilder.String()
					if content == "" {
						content = "(reasoning cut off: thinking budget exceeded)"
					}
					messages = append(messages,
						llm.Message{Role: llm.RoleAssistant, Content: content},
						llm.Message{Role: llm.RoleUser, Content: thinkingBudgetNudge},
					)
					continue
				}
				return
			}
		}

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

// maxUnparsedToolCallNudges bounds how many times per turn the loop re-prompts
// a model that keeps writing tool calls as plain text, so a model that never
// gets it right can't ping-pong forever.
const maxUnparsedToolCallNudges = 2

// unparsedDetectBudget caps how much of a round's raw text is buffered for
// unparsed-tool-call detection — markers appear well within this.
const unparsedDetectBudget = 64 * 1024

// unparsedToolCallNudge is the corrective message sent when a round produced
// tool-call markup as text but no parsed tool call.
const unparsedToolCallNudge = "Your tool call was written as plain text (inside your reasoning or answer), " +
	"so it was NOT executed - nothing has changed. Issue the tool call again now as a real tool call, " +
	"outside of any thinking block, without re-explaining your plan."

// maxThinkingBudgetNudges bounds how many times per turn the loop re-prompts a
// model whose reasoning was cut for exceeding the thinking budget. Separate
// from maxUnparsedToolCallNudges so one failure mode can't consume the other's
// retries. Each spiral still costs a full budget of reasoning tokens, so this
// is kept low.
const maxThinkingBudgetNudges = 2

// thinkingBudgetNudge is the corrective message sent when a round was cut by
// the provider's client-side thinking-budget enforcement.
const thinkingBudgetNudge = "Your reasoning exceeded the thinking budget and was cut off before you took any action. " +
	"Do not re-analyze from scratch: act now on your best current plan - issue the tool call or give " +
	"the final answer directly, with minimal further thinking."

// containsUnparsedToolCall reports whether s contains tool-call markup that
// should have been parsed by the provider but wasn't (Qwen-style
// <tool_call>/<function=...> markers are the ones seen in the wild).
func containsUnparsedToolCall(s string) bool {
	return strings.Contains(s, "<tool_call") || strings.Contains(s, "<function=")
}

func (l *Loop) buildInitialMessages(input llm.Message, 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...)
	input.Role = llm.RoleUser
	messages = append(messages, 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
}