Merge pull request #228 from krissetto/max-iterations-cagent-new

max iteration support for `cagent new`

Author: krissetto

README.md

@@ -219,7 +219,10 @@ export GOOGLE_API_KEY=your_api_key_here     # if anthropic and openai keys are n
 `--max-tokens` can be specified to override the context limit used.  
 When using DMR, the default is 16k to limit memory usage. With all other providers the default is 64k
 
-Example of provider, model and context size overriding:
+`--max-iterations` can be specified to override how many times the agent is allowed to loop when doing tool calling etc.
+When using DMR, the default is set to 20 (small local models have the highest chance of getting confused and looping endlessly). For all other providers, the default is 0 (unlimited).
+
+Example of provider, model, context size and max iterations overriding:
 
 ```sh
 # Use GPT-5 via OpenAI
@@ -230,6 +233,9 @@ cagent new --model dmr/ai/gemma3-qat:12B
 
 # Override the max_tokens used during generation, default is 64k, 16k when using the dmr provider
 cagent new --model openai/gpt-5-mini --max-tokens 32000
+
+# Override max_iterations to limit how much the model can loop autonomously
+cagent new --model dmr/ai/gemma3n:2B-F16 --max-iterations 15
 ```
 
 ---

cmd/root/new.go

@@ -14,8 +14,9 @@ import (
 )
 
 var (
-	modelParam     string
-	maxTokensParam int
+	modelParam         string
+	maxTokensParam     int
+	maxIterationsParam int
 )
 
 // Cmd creates a new command to create a new agent configuration
@@ -93,7 +94,7 @@ func NewNewCmd() *cobra.Command {
 				fmt.Println()
 			}
 
-			out, err := creator.StreamCreateAgent(ctx, ".", prompt, runConfig, modelProvider, model, maxTokensParam)
+			out, rt, err := creator.StreamCreateAgent(ctx, ".", prompt, runConfig, modelProvider, model, maxTokensParam, maxIterationsParam)
 			if err != nil {
 				return err
 			}
@@ -126,6 +127,22 @@ func NewNewCmd() *cobra.Command {
 						llmIsTyping = false
 					}
 					printError(fmt.Errorf("%s", e.Error))
+				case *runtime.MaxIterationsReachedEvent:
+					if llmIsTyping {
+						fmt.Println()
+						llmIsTyping = false
+					}
+
+					result := promptMaxIterationsContinue(e.MaxIterations)
+					switch result {
+					case ConfirmationApprove:
+						rt.Resume(ctx, string(runtime.ResumeTypeApprove))
+					case ConfirmationReject:
+						rt.Resume(ctx, string(runtime.ResumeTypeReject))
+						return nil
+					case ConfirmationAbort:
+						rt.Resume(ctx, string(runtime.ResumeTypeReject))
+					}
 				}
 			}
 			fmt.Print("\n\n")
@@ -135,6 +152,7 @@ func NewNewCmd() *cobra.Command {
 	addGatewayFlags(cmd)
 	cmd.PersistentFlags().StringVar(&modelParam, "model", "", "Model to use, optionally as provider/model where provider is one of: anthropic, openai, google, dmr. If omitted, provider is auto-selected based on available credentials or gateway")
 	cmd.PersistentFlags().IntVar(&maxTokensParam, "max-tokens", 0, "Override max_tokens for the selected model (0 = default)")
+	cmd.PersistentFlags().IntVar(&maxIterationsParam, "max-iterations", 0, "Maximum number of agentic loop iterations to prevent infinite loops (default: 20 for DMR, unlimited for other providers)")
 
 	return cmd
 }

cmd/root/run_text_utils.go

@@ -18,6 +18,7 @@ var (
 	yellow = color.New(color.FgYellow).SprintfFunc()
 	red    = color.New(color.FgRed).SprintfFunc()
 	gray   = color.New(color.FgHiBlack).SprintfFunc()
+	green  = color.New(color.FgGreen).SprintfFunc()
 )
 
 // text styles
@@ -115,6 +116,28 @@ func printToolCallResponse(toolCall tools.ToolCall, response string) {
 	fmt.Printf("\n%s\n", gray("%s response%s", bold(toolCall.Function.Name), formatToolCallResponse(response)))
 }
 
+func promptMaxIterationsContinue(maxIterations int) ConfirmationResult {
+	fmt.Printf("\n%s\n", yellow("⚠️  Maximum iterations (%d) reached. The agent may be stuck in a loop.", maxIterations))
+	fmt.Printf("%s\n", gray("This can happen with smaller or less capable models."))
+	fmt.Printf("\n%s (y/n): ", blue("Do you want to continue for 10 more iterations?"))
+
+	reader := bufio.NewReader(os.Stdin)
+	response, err := reader.ReadString('\n')
+	if err != nil {
+		fmt.Printf("\n%s\n", red("Failed to read input, exiting..."))
+		return ConfirmationAbort
+	}
+
+	response = strings.TrimSpace(strings.ToLower(response))
+	if response == "y" || response == "yes" {
+		fmt.Printf("%s\n\n", green("✓ Continuing..."))
+		return ConfirmationApprove
+	} else {
+		fmt.Printf("%s\n\n", gray("Exiting..."))
+		return ConfirmationReject
+	}
+}
+
 func formatToolCallArguments(arguments string) string {
 	if arguments == "" {
 		return "()"

pkg/creator/agent.go

@@ -121,7 +121,14 @@ func CreateAgent(ctx context.Context, baseDir, prompt string, runConfig config.R
 	return messages[len(messages)-1].Message.Content, fsToolset.path, nil
 }
 
-func StreamCreateAgent(ctx context.Context, baseDir, prompt string, runConfig config.RuntimeConfig, providerName, modelNameOverride string, maxTokensOverride int) (<-chan runtime.Event, error) {
+func StreamCreateAgent(ctx context.Context, baseDir, prompt string, runConfig config.RuntimeConfig, providerName, modelNameOverride string, maxTokensOverride, maxIterations int) (<-chan runtime.Event, runtime.Runtime, error) {
+	// Apply default max iterations if not specified (0 means use defaults)
+	if maxIterations == 0 {
+		// Only when using DMR we set a default limit. Local models are more prone to loops
+		if providerName == "dmr" {
+			maxIterations = 20
+		}
+	}
 	defaultModels := map[string]string{
 		"openai":    "gpt-5-mini",
 		"anthropic": "claude-sonnet-4-0",
@@ -184,7 +191,7 @@ func StreamCreateAgent(ctx context.Context, baseDir, prompt string, runConfig co
 		options.WithGateway(runConfig.ModelsGateway),
 	)
 	if err != nil {
-		return nil, fmt.Errorf("failed to create LLM client: %w", err)
+		return nil, nil, fmt.Errorf("failed to create LLM client: %w", err)
 	}
 
 	fmt.Println("Generating agent configuration....")
@@ -222,11 +229,14 @@ func StreamCreateAgent(ctx context.Context, baseDir, prompt string, runConfig co
 			)))
 	rt, err := runtime.New(newTeam)
 	if err != nil {
-		return nil, fmt.Errorf("failed to create runtime: %w", err)
+		return nil, nil, fmt.Errorf("failed to create runtime: %w", err)
 	}
 
-	sess := session.New(session.WithUserMessage("", prompt))
+	sess := session.New(
+		session.WithUserMessage("", prompt),
+		session.WithMaxIterations(maxIterations),
+	)
 	sess.ToolsApproved = true
 
-	return rt.RunStream(ctx, sess), nil
+	return rt.RunStream(ctx, sess), rt, nil
 }

pkg/runtime/event.go

@@ -289,3 +289,22 @@ func (e *AuthorizationRequiredEvent) isEvent() {}
 func (e *AuthorizationRequiredEvent) GetAgentName() string {
 	return ""
 }
+
+type MaxIterationsReachedEvent struct {
+	Type          string `json:"type"`
+	MaxIterations int    `json:"max_iterations"`
+	AgentContext
+}
+
+func MaxIterationsReached(maxIterations int) Event {
+	return &MaxIterationsReachedEvent{
+		Type:          "max_iterations_reached",
+		MaxIterations: maxIterations,
+	}
+}
+
+func (e *MaxIterationsReachedEvent) isEvent() {}
+
+func (e *MaxIterationsReachedEvent) GetAgentName() string {
+	return e.AgentName
+}

pkg/runtime/runtime.go

@@ -168,7 +168,7 @@ func (r *runtime) finalizeEventChannel(ctx context.Context, sess *session.Sessio
 	}
 }
 
-// Run starts the agent's interaction loop
+// RunStream starts the agent's interaction loop and returns a channel of events
 func (r *runtime) RunStream(ctx context.Context, sess *session.Session) <-chan Event {
 	slog.Debug("Starting runtime stream", "agent", r.currentAgent, "session_id", sess.ID)
 	events := make(chan Event, 128)
@@ -206,7 +206,31 @@ func (r *runtime) RunStream(ctx context.Context, sess *session.Session) <-chan E
 			slog.Debug("Failed to get model definition", "error", err)
 		}
 
+		iteration := 0
+		// Use a runtime copy of maxIterations so we don't modify the session's persistent config
+		runtimeMaxIterations := sess.MaxIterations
 		for {
+			// Check iteration limit
+			if runtimeMaxIterations > 0 && iteration >= runtimeMaxIterations {
+				slog.Debug("Maximum iterations reached", "agent", a.Name(), "iterations", iteration, "max", runtimeMaxIterations)
+				events <- MaxIterationsReached(runtimeMaxIterations)
+
+				// Wait for user decision
+				select {
+				case resumeType := <-r.resumeChan:
+					if resumeType == ResumeTypeApprove {
+						slog.Debug("User chose to continue after max iterations", "agent", a.Name())
+						runtimeMaxIterations = iteration + 10
+					} else {
+						slog.Debug("User chose to exit after max iterations", "agent", a.Name())
+						return
+					}
+				case <-ctx.Done():
+					slog.Debug("Context cancelled while waiting for max iterations decision", "agent", a.Name())
+					return
+				}
+			}
+			iteration++
 			// Exit immediately if the stream context has been cancelled (e.g., Ctrl+C)
 			if err := ctx.Err(); err != nil {
 				slog.Debug("Runtime stream context cancelled, stopping loop", "agent", a.Name(), "session_id", sess.ID)
@@ -816,7 +840,11 @@ func (r *runtime) handleTaskTransfer(ctx context.Context, sess *session.Session,
 	}
 
 	slog.Debug("Creating new session with parent session", "parent_session_id", sess.ID, "tools_approved", sess.ToolsApproved)
-	s := session.New(session.WithSystemMessage(memberAgentTask), session.WithUserMessage("", "Follow the default instructions"))
+	s := session.New(
+		session.WithSystemMessage(memberAgentTask),
+		session.WithUserMessage("", "Follow the default instructions"),
+		session.WithMaxIterations(sess.MaxIterations),
+	)
 	s.SendUserMessage = false
 	s.Title = "Transferred task"
 	s.ToolsApproved = sess.ToolsApproved

pkg/session/migrations.go

@@ -193,6 +193,13 @@ func getAllMigrations() []Migration {
 			UpSQL:       `ALTER TABLE sessions ADD COLUMN send_user_message BOOLEAN DEFAULT 1`,
 			DownSQL:     `ALTER TABLE sessions DROP COLUMN send_user_message`,
 		},
+		{
+			ID:          7,
+			Name:        "007_add_max_iterations_column",
+			Description: "Add max_iterations column to sessions table",
+			UpSQL:       `ALTER TABLE sessions ADD COLUMN max_iterations INTEGER DEFAULT 0`,
+			DownSQL:     `ALTER TABLE sessions DROP COLUMN max_iterations`,
+		},
 		// Add more migrations here as needed
 	}
 }

pkg/session/session.go

@@ -58,6 +58,10 @@ type Session struct {
 	// SendUserMessage is a flag to indicate if the user message should be sent
 	SendUserMessage bool
 
+	// MaxIterations is the maximum number of agentic loop iterations to prevent infinite loops
+	// If 0, there is no limit
+	MaxIterations int `json:"max_iterations"`
+
 	InputTokens  int     `json:"input_tokens"`
 	OutputTokens int     `json:"output_tokens"`
 	Cost         float64 `json:"cost"`
@@ -165,6 +169,12 @@ func WithSystemMessage(content string) Opt {
 	}
 }
 
+func WithMaxIterations(maxIterations int) Opt {
+	return func(s *Session) {
+		s.MaxIterations = maxIterations
+	}
+}
+
 // New creates a new agent session
 func New(opts ...Opt) *Session {
 	sessionID := uuid.New().String()

pkg/session/store.go

@@ -79,8 +79,8 @@ func (s *SQLiteSessionStore) AddSession(ctx context.Context, session *Session) e
 	}
 
 	_, err = s.db.ExecContext(ctx,
-		"INSERT INTO sessions (id, messages, tools_approved, input_tokens, output_tokens, title, send_user_message, created_at) VALUES (?, ?, ?, ?, ?, ?, ?, ?)",
-		session.ID, string(itemsJSON), session.ToolsApproved, session.InputTokens, session.OutputTokens, session.Title, session.SendUserMessage, session.CreatedAt.Format(time.RFC3339))
+		"INSERT INTO sessions (id, messages, tools_approved, input_tokens, output_tokens, title, send_user_message, max_iterations, created_at) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)",
+		session.ID, string(itemsJSON), session.ToolsApproved, session.InputTokens, session.OutputTokens, session.Title, session.SendUserMessage, session.MaxIterations, session.CreatedAt.Format(time.RFC3339))
 	return err
 }
 
@@ -91,12 +91,12 @@ func (s *SQLiteSessionStore) GetSession(ctx context.Context, id string) (*Sessio
 	}
 
 	row := s.db.QueryRowContext(ctx,
-		"SELECT id, messages, tools_approved, input_tokens, output_tokens, title, cost, send_user_message, created_at FROM sessions WHERE id = ?", id)
+		"SELECT id, messages, tools_approved, input_tokens, output_tokens, title, cost, send_user_message, max_iterations, created_at FROM sessions WHERE id = ?", id)
 
-	var messagesJSON, toolsApprovedStr, inputTokensStr, outputTokensStr, titleStr, costStr, sendUserMessageStr, createdAtStr string
+	var messagesJSON, toolsApprovedStr, inputTokensStr, outputTokensStr, titleStr, costStr, sendUserMessageStr, maxIterationsStr, createdAtStr string
 	var sessionID string
 
-	err := row.Scan(&sessionID, &messagesJSON, &toolsApprovedStr, &inputTokensStr, &outputTokensStr, &titleStr, &costStr, &sendUserMessageStr, &createdAtStr)
+	err := row.Scan(&sessionID, &messagesJSON, &toolsApprovedStr, &inputTokensStr, &outputTokensStr, &titleStr, &costStr, &sendUserMessageStr, &maxIterationsStr, &createdAtStr)
 	if err != nil {
 		if errors.Is(err, sql.ErrNoRows) {
 			return nil, ErrNotFound
@@ -146,6 +146,11 @@ func (s *SQLiteSessionStore) GetSession(ctx context.Context, id string) (*Sessio
 		return nil, err
 	}
 
+	maxIterations, err := strconv.Atoi(maxIterationsStr)
+	if err != nil {
+		return nil, err
+	}
+
 	createdAt, err := time.Parse(time.RFC3339, createdAtStr)
 	if err != nil {
 		return nil, err
@@ -160,25 +165,26 @@ func (s *SQLiteSessionStore) GetSession(ctx context.Context, id string) (*Sessio
 		OutputTokens:    outputTokens,
 		Cost:            cost,
 		SendUserMessage: sendUserMessage,
+		MaxIterations:   maxIterations,
 		CreatedAt:       createdAt,
 	}, nil
 }
 
 // GetSessions retrieves all sessions
 func (s *SQLiteSessionStore) GetSessions(ctx context.Context) ([]*Session, error) {
 	rows, err := s.db.QueryContext(ctx,
-		"SELECT id, messages, tools_approved, input_tokens, output_tokens, title, cost, send_user_message, created_at FROM sessions ORDER BY created_at DESC")
+		"SELECT id, messages, tools_approved, input_tokens, output_tokens, title, cost, send_user_message, max_iterations, created_at FROM sessions ORDER BY created_at DESC")
 	if err != nil {
 		return nil, err
 	}
 	defer rows.Close()
 
 	sessions := make([]*Session, 0)
 	for rows.Next() {
-		var messagesJSON, toolsApprovedStr, inputTokensStr, outputTokensStr, titleStr, costStr, sendUserMessageStr, createdAtStr string
+		var messagesJSON, toolsApprovedStr, inputTokensStr, outputTokensStr, titleStr, costStr, sendUserMessageStr, maxIterationsStr, createdAtStr string
 		var sessionID string
 
-		err := rows.Scan(&sessionID, &messagesJSON, &toolsApprovedStr, &inputTokensStr, &outputTokensStr, &titleStr, &costStr, &sendUserMessageStr, &createdAtStr)
+		err := rows.Scan(&sessionID, &messagesJSON, &toolsApprovedStr, &inputTokensStr, &outputTokensStr, &titleStr, &costStr, &sendUserMessageStr, &maxIterationsStr, &createdAtStr)
 		if err != nil {
 			return nil, err
 		}
@@ -225,6 +231,11 @@ func (s *SQLiteSessionStore) GetSessions(ctx context.Context) ([]*Session, error
 			return nil, err
 		}
 
+		maxIterations, err := strconv.Atoi(maxIterationsStr)
+		if err != nil {
+			return nil, err
+		}
+
 		createdAt, err := time.Parse(time.RFC3339, createdAtStr)
 		if err != nil {
 			return nil, err
@@ -239,6 +250,7 @@ func (s *SQLiteSessionStore) GetSessions(ctx context.Context) ([]*Session, error
 			OutputTokens:    outputTokens,
 			Cost:            cost,
 			SendUserMessage: sendUserMessage,
+			MaxIterations:   maxIterations,
 			CreatedAt:       createdAt,
 		}
 
@@ -283,8 +295,8 @@ func (s *SQLiteSessionStore) UpdateSession(ctx context.Context, session *Session
 	}
 
 	result, err := s.db.ExecContext(ctx,
-		"UPDATE sessions SET messages = ?, title = ?, tools_approved = ?, input_tokens = ?, output_tokens = ?, cost = ?, send_user_message = ? WHERE id = ?",
-		string(itemsJSON), session.Title, session.ToolsApproved, session.InputTokens, session.OutputTokens, session.Cost, session.SendUserMessage, session.ID)
+		"UPDATE sessions SET messages = ?, title = ?, tools_approved = ?, input_tokens = ?, output_tokens = ?, cost = ?, send_user_message = ?, max_iterations = ? WHERE id = ?",
+		string(itemsJSON), session.Title, session.ToolsApproved, session.InputTokens, session.OutputTokens, session.Cost, session.SendUserMessage, session.MaxIterations, session.ID)
 	if err != nil {
 		return err
 	}