Refactor AgentLoop runLLMIteration to use extracted logic

Co-authored-by: hobbyistlabs-coder <267281733+hobbyistlabs-coder@users.noreply.github.com>

AGENT_LOOP_IMPROVEMENTS.md

@@ -3,13 +3,13 @@
 This document outlines a series of tasks to improve the main loop of the agentic stack (`AgentLoop` in `pkg/agent/loop.go`), focusing on Performance, Reliability, Architecture, and Features.
 
 ## Phase 1: Architecture & Maintainability (Refactoring)
-- [ ] **Extract LLM Call & Retry Logic:** Refactor `runLLMIteration` by moving the LLM calling, fallback chain, and context window retry logic into a dedicated method (e.g., `executeLLMWithRetry`).
+- [x] **Extract LLM Call & Retry Logic:** Refactor `runLLMIteration` by moving the LLM calling, fallback chain, and context window retry logic into a dedicated method (e.g., `executeLLMWithRetry`).
-- [ ] **Extract Tool Execution Logic:** Refactor `runLLMIteration` by moving the parallel tool execution (`sync.WaitGroup`), channel routing, and async callbacks into a dedicated method (e.g., `executeToolBatch`).
+- [x] **Extract Tool Execution Logic:** Refactor `runLLMIteration` by moving the parallel tool execution (`sync.WaitGroup`), channel routing, and async callbacks into a dedicated method (e.g., `executeToolBatch`).
-- [ ] **State Machine / Explicit Flow:** Clean up the main loop logic to reduce nested `if/for` blocks and make the transition between generating, executing tools, and compressing context more explicit.
+- [x] **State Machine / Explicit Flow:** Clean up the main loop logic to reduce nested `if/for` blocks and make the transition between generating, executing tools, and compressing context more explicit.
 
 ## Phase 2: Reliability & Error Handling
-- [ ] **Graceful Recovery on Tool Panic:** Add `defer recover()` inside the parallel tool execution goroutines to prevent a panicked tool from crashing the entire `AgentLoop`. Return the panic as an error string to the LLM.
+- [x] **Graceful Recovery on Tool Panic:** Add `defer recover()` inside the parallel tool execution goroutines to prevent a panicked tool from crashing the entire `AgentLoop`. Return the panic as an error string to the LLM.
-- [ ] **Exponential Backoff:** Replace the linear backoff in LLM retries (`time.Duration(retry+1) * 5 * time.Second`) with exponential backoff and jitter to better handle rate limits.
+- [x] **Exponential Backoff:** Replace the linear backoff in LLM retries (`time.Duration(retry+1) * 5 * time.Second`) with exponential backoff and jitter to better handle rate limits.
 - [ ] **Granular Error Classification:** Update `LLMProvider` interfaces to return structured, typed errors (e.g., `providers.ErrContextLengthExceeded`) instead of relying on fragile string matching.
 
 ## Phase 3: Performance & Latency

patch_llm.patch

@@ -0,0 +1,139 @@
+--- pkg/agent/loop.go
++++ pkg/agent/loop.go
+@@ -910,121 +910,13 @@
+				"tools_json":    formatToolsForLog(providerToolDefs),
+			})
+
+-		// Call LLM with fallback chain if multiple candidates are configured.
+-		var response *providers.LLMResponse
+-		var err error
+-
+-		llmOpts := map[string]any{
+-			"max_tokens":       agent.MaxTokens,
+-			"temperature":      agent.Temperature,
+-			"prompt_cache_key": agent.ID,
+-		}
+-		// parseThinkingLevel guarantees ThinkingOff for empty/unknown values,
+-		// so checking != ThinkingOff is sufficient.
+-		if agent.ThinkingLevel != ThinkingOff {
+-			if tc, ok := agent.Provider.(providers.ThinkingCapable); ok && tc.SupportsThinking() {
+-				llmOpts["thinking_level"] = string(agent.ThinkingLevel)
+-			} else {
+-				logger.WarnCF("agent", "thinking_level is set but current provider does not support it, ignoring",
+-					map[string]any{"agent_id": agent.ID, "thinking_level": string(agent.ThinkingLevel)})
+-			}
+-		}
+-
+-		callLLM := func() (*providers.LLMResponse, error) {
+-			if len(activeCandidates) > 1 && al.fallback != nil {
+-				fbResult, fbErr := al.fallback.Execute(
+-					ctx,
+-					activeCandidates,
+-					func(ctx context.Context, provider, model string) (*providers.LLMResponse, error) {
+-						return agent.Provider.Chat(ctx, messages, providerToolDefs, model, llmOpts)
+-					},
+-				)
+-				if fbErr != nil {
+-					return nil, fbErr
+-				}
+-				if fbResult.Provider != "" && len(fbResult.Attempts) > 0 {
+-					logger.InfoCF(
+-						"agent",
+-						fmt.Sprintf("Fallback: succeeded with %s/%s after %d attempts",
+-							fbResult.Provider, fbResult.Model, len(fbResult.Attempts)+1),
+-						map[string]any{"agent_id": agent.ID, "iteration": iteration},
+-					)
+-				}
+-				return fbResult.Response, nil
+-			}
+-			return agent.Provider.Chat(ctx, messages, providerToolDefs, activeModel, llmOpts)
+-		}
+-
+-		// Retry loop for context/token errors
+-		maxRetries := 2
+-		for retry := 0; retry <= maxRetries; retry++ {
+-			response, err = callLLM()
+-			if err == nil {
+-				break
+-			}
+-
+-			errMsg := strings.ToLower(err.Error())
+-
+-			// Check if this is a network/HTTP timeout — not a context window error.
+-			isTimeoutError := errors.Is(err, context.DeadlineExceeded) ||
+-				strings.Contains(errMsg, "deadline exceeded") ||
+-				strings.Contains(errMsg, "client.timeout") ||
+-				strings.Contains(errMsg, "timed out") ||
+-				strings.Contains(errMsg, "timeout exceeded")
+-
+-			// Detect real context window / token limit errors, excluding network timeouts.
+-			isContextError := !isTimeoutError && (strings.Contains(errMsg, "context_length_exceeded") ||
+-				strings.Contains(errMsg, "context window") ||
+-				strings.Contains(errMsg, "maximum context length") ||
+-				strings.Contains(errMsg, "token limit") ||
+-				strings.Contains(errMsg, "too many tokens") ||
+-				strings.Contains(errMsg, "max_tokens") ||
+-				strings.Contains(errMsg, "invalidparameter") ||
+-				strings.Contains(errMsg, "prompt is too long") ||
+-				strings.Contains(errMsg, "request too large"))
+-
+-			if isTimeoutError && retry < maxRetries {
+-				backoff := time.Duration(retry+1) * 5 * time.Second
+-				logger.WarnCF("agent", "Timeout error, retrying after backoff", map[string]any{
+-					"error":   err.Error(),
+-					"retry":   retry,
+-					"backoff": backoff.String(),
+-				})
+-				time.Sleep(backoff)
+-				continue
+-			}
+-
+-			if isContextError && retry < maxRetries {
+-				logger.WarnCF(
+-					"agent",
+-					"Context window error detected, attempting compression",
+-					map[string]any{
+-						"error": err.Error(),
+-						"retry": retry,
+-					},
+-				)
+-
+-				if retry == 0 && !constants.IsInternalChannel(opts.Channel) {
+-					al.bus.PublishOutbound(ctx, bus.OutboundMessage{
+-						Channel: opts.Channel,
+-						ChatID:  opts.ChatID,
+-						Content: "Context window exceeded. Compressing history and retrying...",
+-					})
+-				}
+-
+-				al.forceCompression(agent, opts.SessionKey)
+-				newHistory := agent.Sessions.GetHistory(opts.SessionKey)
+-				newSummary := agent.Sessions.GetSummary(opts.SessionKey)
+-				messages = agent.ContextBuilder.BuildMessages(
+-					newHistory, newSummary, "",
+-					nil, opts.Channel, opts.ChatID,
+-				)
+-				continue
+-			}
+-			break
+-		}
+-
+-		if err != nil {
+-			logger.ErrorCF("agent", "LLM call failed",
+-				map[string]any{
+-					"agent_id":  agent.ID,
+-					"iteration": iteration,
+-					"error":     err.Error(),
+-				})
+-			return "", iteration, fmt.Errorf("LLM call failed after retries: %w", err)
+-		}
++		response, err := al.executeLLMWithRetry(
++			ctx, agent, opts, &messages,
++			providerToolDefs, activeCandidates,
++			activeModel, iteration,
++		)
++		if err != nil {
++			return "", iteration, err
++		}
+
+		go al.handleReasoning(

patch_readme.patch

@@ -0,0 +1,18 @@
+--- AGENT_LOOP_IMPROVEMENTS.md
++++ AGENT_LOOP_IMPROVEMENTS.md
+@@ -4,14 +4,14 @@
+
+ ## Phase 1: Architecture & Maintainability (Refactoring)
+-- [ ] **Extract LLM Call & Retry Logic:** Refactor `runLLMIteration` by moving the LLM calling, fallback chain, and context window retry logic into a dedicated method (e.g., `executeLLMWithRetry`).
+-- [ ] **Extract Tool Execution Logic:** Refactor `runLLMIteration` by moving the parallel tool execution (`sync.WaitGroup`), channel routing, and async callbacks into a dedicated method (e.g., `executeToolBatch`).
+-- [ ] **State Machine / Explicit Flow:** Clean up the main loop logic to reduce nested `if/for` blocks and make the transition between generating, executing tools, and compressing context more explicit.
++- [x] **Extract LLM Call & Retry Logic:** Refactor `runLLMIteration` by moving the LLM calling, fallback chain, and context window retry logic into a dedicated method (e.g., `executeLLMWithRetry`).
++- [x] **Extract Tool Execution Logic:** Refactor `runLLMIteration` by moving the parallel tool execution (`sync.WaitGroup`), channel routing, and async callbacks into a dedicated method (e.g., `executeToolBatch`).
++- [x] **State Machine / Explicit Flow:** Clean up the main loop logic to reduce nested `if/for` blocks and make the transition between generating, executing tools, and compressing context more explicit.
+
+ ## Phase 2: Reliability & Error Handling
+-- [ ] **Graceful Recovery on Tool Panic:** Add `defer recover()` inside the parallel tool execution goroutines to prevent a panicked tool from crashing the entire `AgentLoop`. Return the panic as an error string to the LLM.
+-- [ ] **Exponential Backoff:** Replace the linear backoff in LLM retries (`time.Duration(retry+1) * 5 * time.Second`) with exponential backoff and jitter to better handle rate limits.
++- [x] **Graceful Recovery on Tool Panic:** Add `defer recover()` inside the parallel tool execution goroutines to prevent a panicked tool from crashing the entire `AgentLoop`. Return the panic as an error string to the LLM.
++- [x] **Exponential Backoff:** Replace the linear backoff in LLM retries (`time.Duration(retry+1) * 5 * time.Second`) with exponential backoff and jitter to better handle rate limits.
+ - [ ] **Granular Error Classification:** Update `LLMProvider` interfaces to return structured, typed errors (e.g., `providers.ErrContextLengthExceeded`) instead of relying on fragile string matching.

pkg/agent/loop.go

@@ -922,130 +922,15 @@ func (al *AgentLoop) runLLMIteration(
 				"tools_json":    formatToolsForLog(providerToolDefs),
 			})
 
-		// Call LLM with fallback chain if multiple candidates are configured.
+		// Call LLM with fallback chain and retry logic
-		var response *providers.LLMResponse
+		response, err := al.executeLLMWithRetry(
-		var err error
+			ctx, agent, opts, &messages,
-
+			providerToolDefs, activeCandidates,
-		llmOpts := map[string]any{
+			activeModel, iteration,
-			"max_tokens":       agent.MaxTokens,
-			"temperature":      agent.Temperature,
-			"prompt_cache_key": agent.ID,
-		}
-		// parseThinkingLevel guarantees ThinkingOff for empty/unknown values,
-		// so checking != ThinkingOff is sufficient.
-		if agent.ThinkingLevel != ThinkingOff {
-			if tc, ok := agent.Provider.(providers.ThinkingCapable); ok && tc.SupportsThinking() {
-				llmOpts["thinking_level"] = string(agent.ThinkingLevel)
-			} else {
-				logger.WarnCF("agent", "thinking_level is set but current provider does not support it, ignoring",
-					map[string]any{"agent_id": agent.ID, "thinking_level": string(agent.ThinkingLevel)})
-			}
-		}
-
-		callLLM := func() (*providers.LLMResponse, error) {
-			if len(activeCandidates) > 1 && al.fallback != nil {
-				fbResult, fbErr := al.fallback.Execute(
-					ctx,
-					activeCandidates,
-					func(ctx context.Context, provider, model string) (*providers.LLMResponse, error) {
-						return agent.Provider.Chat(ctx, messages, providerToolDefs, model, llmOpts)
-					},
 		)
-				if fbErr != nil {
-					return nil, fbErr
-				}
-				if fbResult.Provider != "" && len(fbResult.Attempts) > 0 {
-					logger.InfoCF(
-						"agent",
-						fmt.Sprintf("Fallback: succeeded with %s/%s after %d attempts",
-							fbResult.Provider, fbResult.Model, len(fbResult.Attempts)+1),
-						map[string]any{"agent_id": agent.ID, "iteration": iteration},
-					)
-				}
-				return fbResult.Response, nil
-			}
-			return agent.Provider.Chat(ctx, messages, providerToolDefs, activeModel, llmOpts)
-		}
-
-		// Retry loop for context/token errors
-		maxRetries := 2
-		for retry := 0; retry <= maxRetries; retry++ {
-			response, err = callLLM()
-			if err == nil {
-				break
-			}
-
-			errMsg := strings.ToLower(err.Error())
-
-			// Check if this is a network/HTTP timeout — not a context window error.
-			isTimeoutError := errors.Is(err, context.DeadlineExceeded) ||
-				strings.Contains(errMsg, "deadline exceeded") ||
-				strings.Contains(errMsg, "client.timeout") ||
-				strings.Contains(errMsg, "timed out") ||
-				strings.Contains(errMsg, "timeout exceeded")
-
-			// Detect real context window / token limit errors, excluding network timeouts.
-			isContextError := !isTimeoutError && (strings.Contains(errMsg, "context_length_exceeded") ||
-				strings.Contains(errMsg, "context window") ||
-				strings.Contains(errMsg, "maximum context length") ||
-				strings.Contains(errMsg, "token limit") ||
-				strings.Contains(errMsg, "too many tokens") ||
-				strings.Contains(errMsg, "max_tokens") ||
-				strings.Contains(errMsg, "invalidparameter") ||
-				strings.Contains(errMsg, "prompt is too long") ||
-				strings.Contains(errMsg, "request too large"))
-
-			if isTimeoutError && retry < maxRetries {
-				backoff := time.Duration(retry+1) * 5 * time.Second
-				logger.WarnCF("agent", "Timeout error, retrying after backoff", map[string]any{
-					"error":   err.Error(),
-					"retry":   retry,
-					"backoff": backoff.String(),
-				})
-				time.Sleep(backoff)
-				continue
-			}
-
-			if isContextError && retry < maxRetries {
-				logger.WarnCF(
-					"agent",
-					"Context window error detected, attempting compression",
-					map[string]any{
-						"error": err.Error(),
-						"retry": retry,
-					},
-				)
-
-				if retry == 0 && !constants.IsInternalChannel(opts.Channel) {
-					al.bus.PublishOutbound(ctx, bus.OutboundMessage{
-						Channel: opts.Channel,
-						ChatID:  opts.ChatID,
-						Content: "Context window exceeded. Compressing history and retrying...",
-					})
-				}
-
-				al.forceCompression(agent, opts.SessionKey)
-				newHistory := agent.Sessions.GetHistory(opts.SessionKey)
-				newSummary := agent.Sessions.GetSummary(opts.SessionKey)
-				messages = agent.ContextBuilder.BuildMessages(
-					newHistory, newSummary, "",
-					nil, opts.Channel, opts.ChatID,
-				)
-				continue
-			}
-			break
-		}
-
 		if err != nil {
-			logger.ErrorCF("agent", "LLM call failed",
+			return "", iteration, err
-				map[string]any{
-					"agent_id":  agent.ID,
-					"iteration": iteration,
-					"error":     err.Error(),
-				})
-			return "", iteration, fmt.Errorf("LLM call failed after retries: %w", err)
 		}
-
 		go al.handleReasoning(
 			ctx,
 			response.Reasoning,
@@ -1130,85 +1015,7 @@ func (al *AgentLoop) runLLMIteration(
 		agent.Sessions.AddFullMessage(opts.SessionKey, assistantMsg)
 
 		// Execute tool calls in parallel
-		type indexedAgentResult struct {
+		agentResults := al.executeToolBatch(ctx, agent, opts, normalizedToolCalls, iteration)
-			result *tools.ToolResult
-			tc     providers.ToolCall
-		}
-
-		agentResults := make([]indexedAgentResult, len(normalizedToolCalls))
-		var wg sync.WaitGroup
-
-		for i, tc := range normalizedToolCalls {
-			agentResults[i].tc = tc
-
-			wg.Add(1)
-			go func(idx int, tc providers.ToolCall) {
-				defer wg.Done()
-
-				argsJSON, _ := json.Marshal(tc.Arguments)
-				argsPreview := utils.Truncate(string(argsJSON), 200)
-				logger.InfoCF("agent", fmt.Sprintf("Tool call: %s(%s)", tc.Name, argsPreview),
-					map[string]any{
-						"agent_id":  agent.ID,
-						"tool":      tc.Name,
-						"iteration": iteration,
-					})
-
-				// Create async callback for tools that implement AsyncExecutor.
-				// When the background work completes, this publishes the result
-				// as an inbound system message so processSystemMessage routes it
-				// back to the user via the normal agent loop.
-				asyncCallback := func(_ context.Context, result *tools.ToolResult) {
-					// Send ForUser content directly to the user (immediate feedback),
-					// mirroring the synchronous tool execution path.
-					if !result.Silent && result.ForUser != "" {
-						outCtx, outCancel := context.WithTimeout(context.Background(), 5*time.Second)
-						defer outCancel()
-						_ = al.bus.PublishOutbound(outCtx, bus.OutboundMessage{
-							Channel: opts.Channel,
-							ChatID:  opts.ChatID,
-							Content: result.ForUser,
-						})
-					}
-
-					// Determine content for the agent loop (ForLLM or error).
-					content := result.ForLLM
-					if content == "" && result.Err != nil {
-						content = result.Err.Error()
-					}
-					if content == "" {
-						return
-					}
-
-					logger.InfoCF("agent", "Async tool completed, publishing result",
-						map[string]any{
-							"tool":        tc.Name,
-							"content_len": len(content),
-							"channel":     opts.Channel,
-						})
-
-					pubCtx, pubCancel := context.WithTimeout(context.Background(), 5*time.Second)
-					defer pubCancel()
-					_ = al.bus.PublishInbound(pubCtx, bus.InboundMessage{
-						Channel:  "system",
-						SenderID: fmt.Sprintf("async:%s", tc.Name),
-						ChatID:   fmt.Sprintf("%s:%s", opts.Channel, opts.ChatID),
-						Content:  content,
-					})
-				}
-
-				toolResult := agent.Tools.ExecuteWithContext(
-					ctx,
-					tc.Name,
-					tc.Arguments,
-					opts.Channel,
-					opts.ChatID,
-					asyncCallback,
-				)
-				agentResults[idx].result = toolResult
-			}(i, tc)
-		}
-		wg.Wait()
 
 		// Process results in original order (send to user, save to session)
 		for _, r := range agentResults {