ollie/agent/loop.go

710 lines
23 KiB
Go

package agent
import (
"context"
"encoding/json"
"errors"
"fmt"
"os"
"strings"
"sync"
"time"
"ollie/backend"
"ollie/toolsrv"
)
const maxTransientRetries = 3
// retryBaseDelay is the base delay for rate-limit retries (5<<attempt seconds).
// Tests can override this to speed up retry tests.
var retryBaseDelay = 5 * time.Second
// streamDropBaseDelay is the base delay for stream-drop retries (2<<attempt seconds).
var streamDropBaseDelay = 2 * time.Second
// Consecutive tool-error thresholds. At the soft limit the model is nudged
// to try a different approach; at the hard limit the loop aborts.
const (
consecutiveErrorSoftLimit = 5
consecutiveErrorHardLimit = 10
)
type toolExecutor func(ctx context.Context, name string, args json.RawMessage) (string, []backend.ContentBlock, error)
// cachedResult stores a tool result with file metadata for staleness detection.
type cachedResult struct {
Result string
ModTime int64 // unix nanoseconds of the file at cache time; 0 if no file
Size int64 // file size at cache time; -1 if no file
}
// extractFilePath attempts to pull a "path" field from JSON tool args.
func extractFilePath(args json.RawMessage) string {
var m map[string]json.RawMessage
if json.Unmarshal(args, &m) != nil {
return ""
}
raw, ok := m["path"]
if !ok {
return ""
}
var p string
if json.Unmarshal(raw, &p) != nil {
return ""
}
return p
}
// cacheValid checks if a cached result is still valid by stat-ing the file.
func cacheValid(entry cachedResult, args json.RawMessage) bool {
if entry.Size < 0 {
// No file was associated — always valid (non-file tool).
return true
}
path := extractFilePath(args)
if path == "" {
return true
}
info, err := os.Stat(path)
if err != nil {
// File gone or unreadable — cache is stale.
return false
}
return info.ModTime().UnixNano() == entry.ModTime && info.Size() == entry.Size
}
// fileStat returns (modtime nanoseconds, size) for a path, or (0, -1) if unavailable.
func fileStat(path string) (int64, int64) {
if path == "" {
return 0, -1
}
info, err := os.Stat(path)
if err != nil {
return 0, -1
}
return info.ModTime().UnixNano(), info.Size()
}
func toolOutputFormat(rt *Runtime, name string) string {
if ti, ok := rt.ToolMeta[name]; ok {
return ti.OutputFormat
}
return ""
}
func toolMemoryTier(rt *Runtime, name string) MemoryTier {
if ti, ok := rt.ToolMeta[name]; ok {
switch ti.Tier {
case "cold":
return MemoryCold
case "warm":
return MemoryWarm
}
}
return MemoryHot
}
// TurnCtx holds per-turn closures and state that vary between turns within
// the same agent. Everything else comes from *Runtime. Embeds context.Context
// so the turn's cancellation is carried with the turn state rather than as a
// separate parameter.
type TurnCtx struct {
context.Context
Output EventHandler
PopInject func() string // returns and clears pending inject, or ""
AutoCompact func(ctx context.Context) // called after each tool round; may compact in-place
Save func() // called after each state.update(); persists mid-turn progress
IncrToolCallCount func() int64
ResultCache *sync.Map
}
func run(rt *Runtime, ctx TurnCtx, h *History) error {
var step int
var consecutiveErrors int // rounds where every tool call returned an error
var lastErrorSig string // signature of previous round's error(s) for repeat detection
var repeatErrorCount int // consecutive rounds with identical error signature
// resultCache stores outputs of read-safe tool calls keyed by name+args.
// Only tools classified as parallel-safe (immutable reads) are cached.
// sync.Map is required because execOne may run in concurrent goroutines.
// Uses the cross-turn cache from TurnCtx if available.
resultCache := ctx.ResultCache
if resultCache == nil {
resultCache = &sync.Map{}
}
for {
emit(ctx, Event{Role: "state", Content: "thinking"})
// Proactive context gate: strip cold material before calling the backend.
if budget := contextBudget(ctx, rt.Backend); budget > 0 && h.estimateTokens() > budget {
h.stripCold(ctx, rt.Backend)
}
history := h.history()
if preamble := rt.PreambleString(); preamble != "" {
history = append([]backend.Message{{Role: "system", Content: preamble}}, history...)
}
// Stream the assistant's response, retrying on rate limits, transient
// backend errors (5xx, network), and mid-stream drops. One stable ID is
// shared by all chunks and the completed assistant message.
responseID := NewResponseID()
var content strings.Builder
var reasoning strings.Builder
var toolCalls []backend.ToolCall
var stopReason string
for attempt := range maxTransientRetries + 1 {
content.Reset()
reasoning.Reset()
toolCalls = nil
ch, err := rt.Backend.ChatStream(ctx, history, rt.Tools, rt.GenParams)
if err != nil {
if ctx.Err() != nil {
return ctx.Err()
}
wait, retryable := transientWait(err, attempt)
if !retryable || attempt >= maxTransientRetries {
return fmt.Errorf("step %d: %w", step, err)
}
var rlErr *backend.RateLimitError
if errors.As(err, &rlErr) {
emit(ctx, Event{Role: "limitretry"})
}
if err := retryCountdown(ctx, wait); err != nil {
return fmt.Errorf("step %d: %w", step, err)
}
continue
}
var done bool
var hadReasoning bool
for ev := range ch {
if ev.Reasoning != "" {
if !hadReasoning {
emit(ctx, Event{Role: "reasoning", Content: "<think>\n"})
hadReasoning = true
}
reasoning.WriteString(ev.Reasoning)
emit(ctx, Event{Role: "reasoning", Content: ev.Reasoning})
}
if ev.Content != "" {
if hadReasoning {
emit(ctx, Event{Role: "reasoning", Content: "\n</think>\n"})
hadReasoning = false
}
content.WriteString(ev.Content)
emit(ctx, Event{Role: "assistant", Content: ev.Content, ResponseID: responseID})
}
toolCalls = append(toolCalls, ev.ToolCalls...)
if ev.Done {
if hadReasoning {
emit(ctx, Event{Role: "reasoning", Content: "\n</think>\n"})
hadReasoning = false
}
stopReason = ev.StopReason
done = true
if ev.Usage.InputTokens > 0 || ev.Usage.OutputTokens > 0 {
emit(ctx, Event{
Role: "usage",
Content: fmt.Sprintf("%d %d 0 %g %d %d", ev.Usage.InputTokens, ev.Usage.OutputTokens, ev.Usage.CostUSD, ev.Usage.CachedInputTokens, ev.Usage.CacheCreationTokens),
})
} else {
// Backend didn't report usage; estimate from content.
inChars := 0
for _, m := range history {
inChars += len(m.Content)
for _, tc := range m.ToolCalls {
inChars += len(tc.Name) + len(tc.Arguments)
}
}
emit(ctx, Event{
Role: "usage",
Content: fmt.Sprintf("%d %d 1 0", inChars/4, content.Len()/4),
})
}
break
}
}
if done {
break
}
if ctx.Err() != nil {
// User pause: record partial state and return.
if hadReasoning {
emit(ctx, Event{Role: "reasoning", Content: "\n</think>\n"})
}
msg := backend.Message{ID: responseID, Role: "assistant", Content: content.String(), Reasoning: reasoning.String(), ToolCalls: toolCalls}
var results []toolResult
for _, tc := range toolCalls {
results = append(results, toolResult{
ToolCallID: tc.ID,
Name: tc.Name,
Content: `{"status":"cancelled","error":"stream interrupted"}`,
IsError: true,
})
}
h.update(msg, results)
ctx.Save()
return ctx.Err()
}
// Pure stream drop — retry if attempts remain.
if attempt >= maxTransientRetries {
return fmt.Errorf("step %d: stream dropped (no more retries)", step)
}
if hadReasoning {
emit(ctx, Event{Role: "reasoning", Content: "\n</think>\n"})
}
wait := streamDropBaseDelay << attempt
if err := retryCountdown(ctx, wait); err != nil {
return fmt.Errorf("step %d: %w", step, err)
}
}
switch stopReason {
case "stop", "tool_calls", "length", "error", "":
// normal
default:
return fmt.Errorf("step %d: %s", step, stopReason)
}
// Parse text-based tool calls: models that don't support the function
// calling API emit tool invocations as plain text (e.g.
// "file_read: args=[...]" (text-based tool call syntax)). Parse these into proper ToolCall structs.
if len(toolCalls) == 0 && len(rt.Tools) > 0 {
if parsed := parseTextToolCalls(content.String(), rt.Tools); len(parsed) > 0 {
toolCalls = parsed
}
}
// Execute tool calls, running consecutive parallel-read-safe tools concurrently.
msg := backend.Message{ID: responseID, Role: "assistant", Content: content.String(), Reasoning: reasoning.String(), ToolCalls: toolCalls}
results := make([]toolResult, 0, len(toolCalls))
interrupted := false
cancelledResult := func(call backend.ToolCall) toolResult {
return toolResult{
ToolCallID: call.ID,
Name: call.Name,
Content: `{"status":"cancelled","error":"interrupted"}`,
IsError: true,
}
}
// execOne runs a single tool call end-to-end and reports whether the
// context was cancelled during execution.
execOne := func(call backend.ToolCall) (toolResult, bool) {
if call.Name == "" {
return toolResult{ToolCallID: call.ID, Name: call.Name, Content: "error: empty tool name", IsError: true}, false
}
if ctx.Err() != nil {
cr := cancelledResult(call)
emit(ctx, Event{Role: "tool", Name: call.Name, Content: cr.Content, OutputFormat: toolOutputFormat(rt, call.Name)})
return cr, true
}
emit(ctx, Event{Role: "call", Name: call.Name, Content: string(call.Arguments)})
readSafe := rt.ToolMeta[call.Name].ReadOnly
if readSafe {
key := call.Name + "\x00" + string(call.Arguments)
if v, ok := resultCache.Load(key); ok {
entry := v.(cachedResult)
if cacheValid(entry, call.Arguments) {
emit(ctx, Event{Role: "tool", Name: call.Name, Content: entry.Result, OutputFormat: toolOutputFormat(rt, call.Name)})
return toolResult{ToolCallID: call.ID, Name: call.Name, Content: entry.Result}, false
}
// Stale — remove and re-execute.
resultCache.Delete(key)
}
}
var result string
var resultBlocks []backend.ContentBlock
var isErr bool
streamed := false
if rt.Exec != nil {
streamBytes := 0
streamCtx := toolsrv.WithOutputStream(ctx, func(data string) {
if streamBytes >= defaultToolResultMaxBytes {
return // already at ceiling, drop further chunks
}
if !streamed {
emit(ctx, Event{Role: "tool", Name: call.Name, OutputFormat: toolOutputFormat(rt, call.Name)})
streamed = true
}
streamBytes += len(data)
if streamBytes > defaultToolResultMaxBytes {
// Emit only the portion within the ceiling.
excess := streamBytes - defaultToolResultMaxBytes
data = data[:len(data)-excess]
}
emit(ctx, Event{Role: "tool", Name: call.Name, Content: data, OutputFormat: toolOutputFormat(rt, call.Name)})
})
out, blocks, err := rt.Exec(streamCtx, call.Name, call.Arguments)
if err != nil {
isErr = true
if ctx.Err() != nil {
result = "error: tool execution interrupted by user"
if ctx.PopInject != nil {
if injected := ctx.PopInject(); injected != "" {
result += "\n\n<system-user-interruption>\n" + injected + "\n</system-user-interruption>"
}
}
emit(ctx, Event{Role: "tool", Name: call.Name, Content: result, OutputFormat: toolOutputFormat(rt, call.Name)})
return toolResult{ToolCallID: call.ID, Name: call.Name, Content: result, IsError: true}, true
}
var rlErr *toolsrv.RateLimitedError
if errors.As(err, &rlErr) {
result = fmt.Sprintf("error: shell is rate-limited — blocked for %v. Do not call shell() again until the block expires. Use other tools or wait.", rlErr.Remaining)
} else {
result = fmt.Sprintf("error: %v", err)
}
} else {
result = out
resultBlocks = blocks
}
} else {
result = "error: no tool executor configured"
isErr = true
}
// Accumulate suffix text (user-interruptions, truncation hints)
// that must be emitted after streaming completes.
var suffix string
if ctx.PopInject != nil {
if injected := ctx.PopInject(); injected != "" {
suffix += "\n\n<system-user-interruption>\n" + injected + "\n</system-user-interruption>"
}
}
result += suffix
// Safety ceiling: unconditionally cap all tool results at 128KB.
if len(result) > defaultToolResultMaxBytes {
orig := len(result)
result = strings.ToValidUTF8(result[:defaultToolResultMaxBytes], "")
suffix = fmt.Sprintf("\n\n[HARD LIMIT: %s output truncated — %d of %d bytes shown. This is a safety ceiling, not a semantic boundary.]",
call.Name, defaultToolResultMaxBytes, orig)
result += suffix
}
if readSafe && !isErr {
path := extractFilePath(call.Arguments)
mtime, size := fileStat(path)
resultCache.Store(call.Name+"\x00"+string(call.Arguments), cachedResult{
Result: result,
ModTime: mtime,
Size: size,
})
}
if !streamed {
emit(ctx, Event{Role: "tool", Name: call.Name, Content: result, OutputFormat: toolOutputFormat(rt, call.Name)})
} else {
// Emit suffixes that were appended after streaming completed
// so they appear at the end of the chat output.
if suffix != "" {
emit(ctx, Event{Role: "tool", Name: call.Name, Content: suffix, OutputFormat: toolOutputFormat(rt, call.Name)})
}
// Signal the log handler to close the source fence and emit
// [[[end]]]. Without this, streamingRole stays set and the
// fence is only closed on the next role transition — or never,
// if the agent turn ends or is interrupted.
emit(ctx, Event{Role: "tool", Name: call.Name, OutputFormat: toolOutputFormat(rt, call.Name)})
}
tier := MemoryHot
if !isErr {
tier = toolMemoryTier(rt, call.Name)
}
return toolResult{ToolCallID: call.ID, Name: call.Name, Content: result, ContentBlocks: resultBlocks, IsError: isErr, Tier: tier}, false
}
isParallelSafe := func(name string) bool {
return name != "" && rt.ToolMeta[name].ReadOnly
}
for i := 0; i < len(toolCalls) && !interrupted; {
if ctx.Err() != nil {
for _, remaining := range toolCalls[i:] {
cr := cancelledResult(remaining)
emit(ctx, Event{Role: "tool", Name: remaining.Name, Content: cr.Content, OutputFormat: toolOutputFormat(rt, remaining.Name)})
results = append(results, cr)
}
interrupted = true
break
}
// Collect a run of consecutive parallel-safe calls.
j := i + 1
if isParallelSafe(toolCalls[i].Name) {
for j < len(toolCalls) && isParallelSafe(toolCalls[j].Name) {
j++
}
}
batch := toolCalls[i:j]
fillCancelled := func(from int) {
for _, remaining := range toolCalls[from:] {
cr := cancelledResult(remaining)
emit(ctx, Event{Role: "tool", Name: remaining.Name, Content: cr.Content, OutputFormat: toolOutputFormat(rt, remaining.Name)})
results = append(results, cr)
}
}
if len(batch) == 1 {
tr, wasInt := execOne(batch[0])
results = append(results, tr)
if wasInt {
fillCancelled(j)
interrupted = true
}
} else {
// Fan out the batch concurrently; preserve submission order in results.
// Deduplicate identical calls so only one executes per unique key.
type inflightResult struct {
tr toolResult
wasInt bool
}
inflight := make(map[string]int) // key -> index of first occurrence
batchResults := make([]toolResult, len(batch))
batchInt := make([]bool, len(batch))
var wg sync.WaitGroup
uniqueResults := make([]inflightResult, len(batch))
for k, call := range batch {
key := call.Name + "\x00" + string(call.Arguments)
if first, dup := inflight[key]; dup {
// Will copy result from first occurrence after wg.Wait.
batchResults[k] = toolResult{} // placeholder
_ = first // used below
continue
}
inflight[key] = k
wg.Add(1)
go func(k int, call backend.ToolCall) {
defer wg.Done()
uniqueResults[k].tr, uniqueResults[k].wasInt = execOne(call)
}(k, call)
}
wg.Wait()
// Fill results: unique calls get their own result, duplicates copy from first.
for k, call := range batch {
key := call.Name + "\x00" + string(call.Arguments)
first := inflight[key]
if k == first {
batchResults[k] = uniqueResults[k].tr
batchInt[k] = uniqueResults[k].wasInt
} else {
batchResults[k] = toolResult{
ToolCallID: call.ID,
Name: call.Name,
Content: uniqueResults[first].tr.Content,
IsError: uniqueResults[first].tr.IsError,
}
batchInt[k] = uniqueResults[first].wasInt
emit(ctx, Event{Role: "call", Name: call.Name, Content: string(call.Arguments)})
emit(ctx, Event{Role: "tool", Name: call.Name, Content: batchResults[k].Content, OutputFormat: toolOutputFormat(rt, call.Name)})
}
}
// Add all batch results — model needs a result for every tool call.
results = append(results, batchResults...)
for _, wasInt := range batchInt {
if wasInt {
fillCancelled(j)
interrupted = true
break
}
}
}
i = j
}
h.update(msg, results)
ctx.Save()
if interrupted {
return ctx.Err()
}
if len(toolCalls) == 0 {
break
}
// Track consecutive rounds where every tool call errored.
// Rate-limited errors don't count — they're a temporary block,
// not a real failure. Reset consecutiveErrors so the hard abort
// doesn't fire.
allErrors := len(results) > 0
allRateLimited := allErrors
for _, r := range results {
if !r.IsError {
allErrors = false
allRateLimited = false
break
}
if !strings.HasPrefix(r.Content, "error: shell is rate-limited") {
allRateLimited = false
}
}
if allRateLimited {
consecutiveErrors = 0
} else if allErrors {
consecutiveErrors++
} else {
consecutiveErrors = 0
}
if consecutiveErrors >= consecutiveErrorHardLimit {
emit(ctx, Event{Role: "error", Content: fmt.Sprintf("%d consecutive tool errors — aborting", consecutiveErrors)})
return fmt.Errorf("step %d: %d consecutive tool errors", step, consecutiveErrors)
}
if consecutiveErrors == consecutiveErrorSoftLimit {
// Nudge the model to try a different approach and keep the plan current.
h.update(backend.Message{
Role: "user",
Content: "<system-consecutive-errors>\nYour last several tool calls all failed. Try a different approach, or ask the user for help. Keep your plan current.\n</system-consecutive-errors>",
}, nil)
}
// Repeated identical error detection: if the same error occurs 3 times
// in a row, the model is stuck on a format/argument issue. Inject the
// error content as a corrective nudge immediately rather than waiting
// for the generic soft limit.
if allErrors && len(results) > 0 {
var errSig strings.Builder
for _, r := range results {
if r.IsError {
errSig.WriteString(r.Content)
errSig.WriteByte('\n')
}
}
sig := errSig.String()
if sig == lastErrorSig {
repeatErrorCount++
} else {
repeatErrorCount = 1
lastErrorSig = sig
}
if repeatErrorCount == 3 {
h.update(backend.Message{
Role: "user",
Content: "<system-repeated-error>\nYou have made the same malformed tool call 3 times in a row. Read the error message carefully and fix the arguments. The error was:\n" + strings.TrimSpace(results[0].Content) + "\n</system-repeated-error>",
}, nil)
}
} else {
repeatErrorCount = 0
lastErrorSig = ""
}
// Reset step counter when an action-oriented tool is called successfully.
if rt.ToolMeta != nil {
for _, r := range results {
if !r.IsError && rt.ToolMeta[r.Name].ResetsCounter {
step = 0
break
}
}
}
// Step budget warnings: nudge the agent at 50% and when exhausted.
if rt.MaxSteps > 0 {
halfBudget := rt.MaxSteps / 2
if step == halfBudget {
h.update(backend.Message{
Role: "user",
Content: fmt.Sprintf("<system-step-budget-warning>\nYou have used %d/%d research steps without taking action. Consider making progress — write code, edit files, or run commands. Action tools reset this counter.\n</system-step-budget-warning>", step, rt.MaxSteps),
}, nil)
} else if step >= rt.MaxSteps-1 {
emit(ctx, Event{Role: "maxsteps", Content: fmt.Sprintf("%d", step+1)})
h.update(backend.Message{
Role: "user",
Content: fmt.Sprintf("<system-step-budget-stop>\nStep budget exhausted (%d/%d steps used). Stop calling tools. Summarize what you have done and what remains, then stop.\n</system-step-budget-stop>", step+1, rt.MaxSteps),
}, nil)
break
}
}
if ctx.AutoCompact != nil {
ctx.AutoCompact(ctx)
}
step++
}
return nil
}
func emit(ctx TurnCtx, msg Event) {
if ctx.Output != nil {
ctx.Output(msg)
}
}
// contextBudget returns the token threshold (50% of context length) above which
// cold material should be proactively stripped. Returns 0 if unknown.
func contextBudget(ctx context.Context, b backend.Backend) int {
if b == nil {
return 0
}
ctxLen := b.ContextLength(ctx)
if ctxLen <= 0 {
return 0
}
return ctxLen / 2
}
// defaultToolResultMaxBytes caps tool result content sent back to the model.
const defaultToolResultMaxBytes = 131072
// MemoryTier classifies how long a tool result stays in the hot message list.
type MemoryTier int
const (
MemoryHot MemoryTier = iota // stays verbatim in messages (default)
MemoryWarm // summarized on next compaction pass
MemoryCold // immediately summarized at update time
)
// toolResult holds the output of a single tool call.
type toolResult struct {
ToolCallID string // may be empty (Ollama does not always set this)
Name string
Content string
ContentBlocks []backend.ContentBlock // image blocks extracted from tool output
IsError bool
Tier MemoryTier // default MemoryHot; tools may self-classify
}
// retryable. Rate limits use longer waits; transient/network errors use shorter.
func transientWait(err error, attempt int) (time.Duration, bool) {
var rlErr *backend.RateLimitError
if errors.As(err, &rlErr) {
wait := rlErr.RetryAfter
if wait == 0 {
wait = retryBaseDelay << attempt
}
return wait, true
}
var tErr *backend.TransientError
if errors.As(err, &tErr) {
return time.Duration(2<<attempt) * time.Second, true
}
return 0, false
}
func retryCountdown(ctx TurnCtx, wait time.Duration) error {
deadline := time.Now().Add(wait)
for {
remaining := time.Until(deadline)
if remaining <= 0 {
return nil
}
totalSecs := int(remaining.Seconds()) + 1
h := totalSecs / 3600
m := (totalSecs % 3600) / 60
s := totalSecs % 60
emit(ctx, Event{Role: "retry", Content: fmt.Sprintf("%02d:%02d:%02d", h, m, s)})
select {
case <-ctx.Done():
return ctx.Err()
case <-time.After(min(remaining, time.Second)):
}
}
}