Move maxIteration to session

krissetto · krissetto · commit ba150ae763bb · 2025-09-17T14:09:05.000+02:00
Signed-off-by: Christopher Petito &lt;chrisjpetito@gmail.com&gt;
diff --git a/README.md b/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
 ```
 
 ---
diff --git a/cmd/root/run.go b/cmd/root/run.go
@@ -349,7 +349,7 @@ func runWithoutTUI(ctx context.Context, agentFilename string, rt runtime.Runtime
 		lastAgent := rt.CurrentAgent().Name()
 		llmIsTyping := false
 		var lastConfirmedToolCallID string
-		for event := range rt.RunStream(loopCtx, sess, 0) {
+		for event := range rt.RunStream(loopCtx, sess) {
 			if event.GetAgentName() != "" && (firstLoop || lastAgent != event.GetAgentName()) {
 				if !firstLoop {
 					if llmIsTyping {
diff --git a/examples/golibrary/stream/main.go b/examples/golibrary/stream/main.go
@@ -44,7 +44,7 @@ func main() {
 
 	sess := session.New(session.WithUserMessage("", "How are you doing?"))
 
-	events := rt.RunStream(ctx, sess, 0)
+	events := rt.RunStream(ctx, sess)
 	for event := range events {
 		switch e := event.(type) {
 		case *runtime.AgentChoiceEvent:
diff --git a/pkg/app/app.go b/pkg/app/app.go
@@ -58,7 +58,7 @@ func (a *App) Run(ctx context.Context, message string) {
 
 		// User message
 		a.session.AddMessage(session.UserMessage(a.agentFilename, message))
-		for event := range a.runtime.RunStream(ctx, a.session, 0) {
+		for event := range a.runtime.RunStream(ctx, a.session) {
 			a.events <- event
 		}
 	}()
diff --git a/pkg/creator/agent.go b/pkg/creator/agent.go
@@ -232,8 +232,11 @@ func StreamCreateAgent(ctx context.Context, baseDir, prompt string, runConfig co
 		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, maxIterations), rt, nil
+	return rt.RunStream(ctx, sess), rt, nil
 }
diff --git a/pkg/runtime/remote_runtime.go b/pkg/runtime/remote_runtime.go
@@ -72,9 +72,8 @@ func (r *RemoteRuntime) CurrentAgent() *agent.Agent {
 	return agent.New(r.currentAgent, fmt.Sprintf("Remote agent: %s", r.currentAgent))
 }
 
-// RunStream starts the agent's interaction loop with an optional iteration limit
-// If maxIterations is 0, there is no limit
-func (r *RemoteRuntime) RunStream(ctx context.Context, sess *session.Session, maxIterations int) <-chan Event {
+// RunStream starts the agent's interaction loop and returns a channel of events
+func (r *RemoteRuntime) RunStream(ctx context.Context, sess *session.Session) <-chan Event {
 	slog.Debug("Starting remote runtime stream", "agent", r.currentAgent, "session_id", r.sessionID)
 	events := make(chan Event, 128)
 
@@ -126,7 +125,7 @@ func (r *RemoteRuntime) RunStream(ctx context.Context, sess *session.Session, ma
 
 // Run starts the agent's interaction loop and returns the final messages
 func (r *RemoteRuntime) Run(ctx context.Context, sess *session.Session) ([]session.Message, error) {
-	eventsChan := r.RunStream(ctx, sess, 0)
+	eventsChan := r.RunStream(ctx, sess)
 
 	for event := range eventsChan {
 		if errEvent, ok := event.(*ErrorEvent); ok {
diff --git a/pkg/runtime/runtime.go b/pkg/runtime/runtime.go
@@ -40,7 +40,7 @@ type Runtime interface {
 	// CurrentAgent returns the currently active agent
 	CurrentAgent() *agent.Agent
 	// RunStream starts the agent's interaction loop and returns a channel of events
-	RunStream(ctx context.Context, sess *session.Session, maxIterations int) <-chan Event
+	RunStream(ctx context.Context, sess *session.Session) <-chan Event
 	// Run starts the agent's interaction loop and returns the final messages
 	Run(ctx context.Context, sess *session.Session) ([]session.Message, error)
 	// Resume allows resuming execution after user confirmation
@@ -168,9 +168,8 @@ func (r *runtime) finalizeEventChannel(ctx context.Context, sess *session.Sessio
 	}
 }
 
-// RunStream starts the agent's interaction loop with an optional iteration limit
-// If maxIterations is 0, there is no limit
-func (r *runtime) RunStream(ctx context.Context, sess *session.Session, maxIterations int) <-chan Event {
+// 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)
 
@@ -208,18 +207,20 @@ func (r *runtime) RunStream(ctx context.Context, sess *session.Session, maxItera
 		}
 
 		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 maxIterations > 0 && iteration >= maxIterations {
-				slog.Debug("Maximum iterations reached", "agent", a.Name(), "iterations", iteration, "max", maxIterations)
-				events <- MaxIterationsReached(maxIterations)
+			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())
-						maxIterations = iteration + 10
+						runtimeMaxIterations = iteration + 10
 					} else {
 						slog.Debug("User chose to exit after max iterations", "agent", a.Name())
 						return
@@ -425,7 +426,7 @@ func (r *runtime) ResumeCodeReceived(_ context.Context, code string) error {
 
 // Run starts the agent's interaction loop
 func (r *runtime) Run(ctx context.Context, sess *session.Session) ([]session.Message, error) {
-	eventsChan := r.RunStream(ctx, sess, 0)
+	eventsChan := r.RunStream(ctx, sess)
 
 	for event := range eventsChan {
 		if errEvent, ok := event.(*ErrorEvent); ok {
@@ -839,12 +840,16 @@ 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
 
-	for event := range r.RunStream(ctx, s, 0) {
+	for event := range r.RunStream(ctx, s) {
 		evts <- event
 		if errEvent, ok := event.(*ErrorEvent); ok {
 			span.RecordError(fmt.Errorf("%s", errEvent.Error))
diff --git a/pkg/server/server.go b/pkg/server/server.go
@@ -883,7 +883,7 @@ func (s *Server) runAgent(c echo.Context) error {
 	c.Response().Header().Set("Connection", "keep-alive")
 	c.Response().WriteHeader(http.StatusOK)
 
-	streamChan := rt.RunStream(c.Request().Context(), sess, 0)
+	streamChan := rt.RunStream(c.Request().Context(), sess)
 	for event := range streamChan {
 		data, err := json.Marshal(event)
 		if err != nil {
diff --git a/pkg/session/migrations.go b/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
 	}
 }
diff --git a/pkg/session/session.go b/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()
diff --git a/pkg/session/store.go b/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
 	}

Original file line number	Diff line number	Diff line change
`@@ -58,7 +58,7 @@ func (a *App) Run(ctx context.Context, message string) {`
`58`	`58`
`59`	`59`	`// User message`
`60`	`60`	`a.session.AddMessage(session.UserMessage(a.agentFilename, message))`
`61`		`- for event := range a.runtime.RunStream(ctx, a.session, 0) {`
	`61`	`+ for event := range a.runtime.RunStream(ctx, a.session) {`
`62`	`62`	`a.events <- event`
`63`	`63`	`}`
`64`	`64`	`}()`
Original file line number	Diff line number	Diff line change
`@@ -232,8 +232,11 @@ func StreamCreateAgent(ctx context.Context, baseDir, prompt string, runConfig co`
`232`	`232`	`return nil, nil, fmt.Errorf("failed to create runtime: %w", err)`
`233`	`233`	`}`
`234`	`234`
`235`		`- sess := session.New(session.WithUserMessage("", prompt))`
	`235`	`+ sess := session.New(`
	`236`	`+ session.WithUserMessage("", prompt),`
	`237`	`+ session.WithMaxIterations(maxIterations),`
	`238`	`+ )`
`236`	`239`	`sess.ToolsApproved = true`
`237`	`240`
`238`		`- return rt.RunStream(ctx, sess, maxIterations), rt, nil`
	`241`	`+ return rt.RunStream(ctx, sess), rt, nil`
`239`	`242`	`}`
Original file line number	Diff line number	Diff line change
`@@ -193,6 +193,13 @@ func getAllMigrations() []Migration {`
`193`	`193`	UpSQL: `ALTER TABLE sessions ADD COLUMN send_user_message BOOLEAN DEFAULT 1`,
`194`	`194`	DownSQL: `ALTER TABLE sessions DROP COLUMN send_user_message`,
`195`	`195`	`},`
	`196`	`+ {`
	`197`	`+ ID: 7,`
	`198`	`+ Name: "007_add_max_iterations_column",`
	`199`	`+ Description: "Add max_iterations column to sessions table",`
	`200`	+ UpSQL: `ALTER TABLE sessions ADD COLUMN max_iterations INTEGER DEFAULT 0`,
	`201`	+ DownSQL: `ALTER TABLE sessions DROP COLUMN max_iterations`,
	`202`	`+ },`
`196`	`203`	`// Add more migrations here as needed`
`197`	`204`	`}`
`198`	`205`	`}`