Skip to content

Commit 4d627cc

Browse files
authored
Merge pull request #237 from krissetto/max-iterations-cagent-run
Support max_iterations config in agent yaml
2 parents 5ad3082 + 5190f4c commit 4d627cc

10 files changed

Lines changed: 283 additions & 3 deletions

File tree

cmd/root/run.go

Lines changed: 18 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -238,7 +238,7 @@ func doRunCommand(ctx context.Context, args []string, exec bool) error {
238238
return fmt.Errorf("failed to create runtime: %w", err)
239239
}
240240
rt = localRt
241-
sess = session.New()
241+
sess = session.New(session.WithMaxIterations(rt.CurrentAgent().MaxIterations()))
242242
sess.ToolsApproved = autoApprove
243243
slog.Debug("Using local runtime", "agent", agentName)
244244
}
@@ -432,6 +432,23 @@ func runWithoutTUI(ctx context.Context, agentFilename string, rt runtime.Runtime
432432
lastErr = fmt.Errorf("%s", e.Error)
433433
printError(lastErr)
434434
}
435+
case *runtime.MaxIterationsReachedEvent:
436+
if llmIsTyping {
437+
fmt.Println()
438+
llmIsTyping = false
439+
}
440+
441+
result := promptMaxIterationsContinue(e.MaxIterations)
442+
switch result {
443+
case ConfirmationApprove:
444+
rt.Resume(ctx, string(runtime.ResumeTypeApprove))
445+
case ConfirmationReject:
446+
rt.Resume(ctx, string(runtime.ResumeTypeReject))
447+
return nil
448+
case ConfirmationAbort:
449+
rt.Resume(ctx, string(runtime.ResumeTypeReject))
450+
return nil
451+
}
435452
}
436453
}
437454

pkg/agent/agent.go

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,7 @@ type Agent struct {
3939
parents []*Agent
4040
addDate bool
4141
addEnvironmentInfo bool
42+
maxIterations int
4243
toolWrapper toolWrapper
4344
memoryManager memorymanager.Manager
4445
}
@@ -75,6 +76,10 @@ func (a *Agent) AddEnvironmentInfo() bool {
7576
return a.addEnvironmentInfo
7677
}
7778

79+
func (a *Agent) MaxIterations() int {
80+
return a.maxIterations
81+
}
82+
7883
// Description returns the agent's description
7984
func (a *Agent) Description() string {
8085
return a.description

pkg/agent/opts.go

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -72,3 +72,9 @@ func WithMemoryManager(mm memorymanager.Manager) Opt {
7272
a.memoryManager = mm
7373
}
7474
}
75+
76+
func WithMaxIterations(maxIterations int) Opt {
77+
return func(a *Agent) {
78+
a.maxIterations = maxIterations
79+
}
80+
}

pkg/config/v2/types.go

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@ type AgentConfig struct {
2121
SubAgents []string `json:"sub_agents,omitempty" yaml:"sub_agents,omitempty"`
2222
AddDate bool `json:"add_date,omitempty" yaml:"add_date,omitempty"`
2323
AddEnvironmentInfo bool `json:"add_environment_info,omitempty" yaml:"add_environment_info,omitempty"`
24+
MaxIterations int `json:"max_iterations,omitempty" yaml:"max_iterations,omitempty"`
2425
}
2526

2627
// ModelConfig represents the configuration for a model

pkg/evaluation/evaluation.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -74,7 +74,7 @@ func runLoop(ctx context.Context, rt runtime.Runtime, eval *session.Session) ([]
7474
}
7575
}
7676

77-
sess := session.New()
77+
sess := session.New(session.WithMaxIterations(rt.CurrentAgent().MaxIterations()))
7878
for i := range userMessages {
7979
sess.AddMessage(&userMessages[i])
8080
_, err := rt.Run(ctx, sess)

pkg/runtime/client.go

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -377,6 +377,9 @@ func (c *Client) runAgentWithAgentName(ctx context.Context, sessionID, agent, ag
377377
cost, _ := usage["cost"].(float64)
378378

379379
eventChan <- TokenUsage(int(inputTokens), int(outputTokens), int(contextLength), int(contextLimit), cost)
380+
case "max_iterations_reached":
381+
maxIterations, _ := event["max_iterations"].(float64)
382+
eventChan <- MaxIterationsReached(int(maxIterations))
380383
case "session_title":
381384
eventChan <- SessionTitle(event["session_id"].(string), event["title"].(string))
382385
case "session_summary":

pkg/runtime/runtime.go

Lines changed: 15 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -204,6 +204,14 @@ func (r *runtime) RunStream(ctx context.Context, sess *session.Session) <-chan E
204204
runtimeMaxIterations = iteration + 10
205205
} else {
206206
slog.Debug("User chose to exit after max iterations", "agent", a.Name())
207+
// Synthesize a final assistant message so callers (e.g., parent agents)
208+
// receive a non-empty response and providers are not given empty tool outputs.
209+
assistantMessage := chat.Message{
210+
Role: chat.MessageRoleAssistant,
211+
Content: fmt.Sprintf("I have reached the maximum number of iterations (%d). Stopping as requested by user.", runtimeMaxIterations),
212+
CreatedAt: time.Now().Format(time.RFC3339),
213+
}
214+
sess.AddMessage(session.NewAgentMessage(a, &assistantMessage))
207215
return
208216
}
209217
case <-ctx.Done():
@@ -821,10 +829,16 @@ func (r *runtime) handleTaskTransfer(ctx context.Context, sess *session.Session,
821829
}
822830

823831
slog.Debug("Creating new session with parent session", "parent_session_id", sess.ID, "tools_approved", sess.ToolsApproved)
832+
833+
subAgentMaxIter := 0
834+
if child := r.team.Agent(params.Agent); child != nil {
835+
subAgentMaxIter = child.MaxIterations()
836+
}
837+
824838
s := session.New(
825839
session.WithSystemMessage(memberAgentTask),
826840
session.WithUserMessage("", "Follow the default instructions"),
827-
session.WithMaxIterations(sess.MaxIterations),
841+
session.WithMaxIterations(subAgentMaxIter),
828842
)
829843
s.SendUserMessage = false
830844
s.Title = "Transferred task"

pkg/teamloader/teamloader.go

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -146,6 +146,7 @@ func Load(ctx context.Context, path string, runtimeConfig config.RuntimeConfig)
146146
agent.WithDescription(agentConfig.Description),
147147
agent.WithAddDate(agentConfig.AddDate),
148148
agent.WithAddEnvironmentInfo(agentConfig.AddEnvironmentInfo),
149+
agent.WithMaxIterations(agentConfig.MaxIterations),
149150
}
150151
for _, model := range models {
151152
opts = append(opts, agent.WithModel(model))

pkg/tui/dialog/max_iterations.go

Lines changed: 224 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,224 @@
1+
package dialog
2+
3+
import (
4+
"fmt"
5+
"strings"
6+
7+
"github.com/charmbracelet/bubbles/v2/key"
8+
tea "github.com/charmbracelet/bubbletea/v2"
9+
"github.com/charmbracelet/lipgloss/v2"
10+
11+
"github.com/docker/cagent/pkg/app"
12+
"github.com/docker/cagent/pkg/tui/core"
13+
)
14+
15+
// maxIterationsDialog implements DialogModel for max iterations confirmation
16+
type maxIterationsDialog struct {
17+
width, height int
18+
maxIterations int
19+
app *app.App
20+
keyMap maxIterationsKeyMap
21+
}
22+
23+
// SetSize implements Dialog.
24+
func (d *maxIterationsDialog) SetSize(width, height int) tea.Cmd {
25+
d.width = width
26+
d.height = height
27+
return nil
28+
}
29+
30+
// maxIterationsKeyMap defines key bindings for max iterations confirmation dialog
31+
type maxIterationsKeyMap struct {
32+
Yes key.Binding
33+
No key.Binding
34+
}
35+
36+
// defaultMaxIterationsKeyMap returns default key bindings
37+
func defaultMaxIterationsKeyMap() maxIterationsKeyMap {
38+
return maxIterationsKeyMap{
39+
Yes: key.NewBinding(
40+
key.WithKeys("y", "Y"),
41+
key.WithHelp("Y", "continue"),
42+
),
43+
No: key.NewBinding(
44+
key.WithKeys("n", "N"),
45+
key.WithHelp("N", "stop"),
46+
),
47+
}
48+
}
49+
50+
// NewMaxIterationsDialog creates a new max iterations confirmation dialog
51+
func NewMaxIterationsDialog(maxIterations int, appInstance *app.App) Dialog {
52+
return &maxIterationsDialog{
53+
maxIterations: maxIterations,
54+
app: appInstance,
55+
keyMap: defaultMaxIterationsKeyMap(),
56+
}
57+
}
58+
59+
// Init initializes the max iterations confirmation dialog
60+
func (d *maxIterationsDialog) Init() tea.Cmd {
61+
return nil
62+
}
63+
64+
// Update handles messages for the max iterations confirmation dialog
65+
func (d *maxIterationsDialog) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
66+
switch msg := msg.(type) {
67+
case tea.WindowSizeMsg:
68+
d.width = msg.Width
69+
d.height = msg.Height
70+
return d, nil
71+
72+
case tea.KeyPressMsg:
73+
switch {
74+
case key.Matches(msg, d.keyMap.Yes):
75+
if d.app != nil {
76+
d.app.Resume("approve")
77+
}
78+
return d, core.CmdHandler(CloseDialogMsg{})
79+
case key.Matches(msg, d.keyMap.No):
80+
if d.app != nil {
81+
d.app.Resume("reject")
82+
}
83+
return d, core.CmdHandler(CloseDialogMsg{})
84+
}
85+
86+
if msg.String() == "ctrl+c" {
87+
return d, tea.Quit
88+
}
89+
}
90+
91+
return d, nil
92+
}
93+
94+
// Position returns the dialog position (centered)
95+
func (d *maxIterationsDialog) Position() (row, col int) {
96+
// Render the dialog content to measure its actual dimensions
97+
dialogContent := d.View()
98+
99+
// Get the actual rendered dimensions
100+
dialogWidth := lipgloss.Width(dialogContent)
101+
dialogHeight := lipgloss.Height(dialogContent)
102+
103+
// Calculate centered position
104+
col = max(0, (d.width-dialogWidth)/2)
105+
row = max(0, (d.height-dialogHeight)/2)
106+
107+
// Ensure dialog fits on screen
108+
if col+dialogWidth > d.width {
109+
col = max(0, d.width-dialogWidth)
110+
}
111+
if row+dialogHeight > d.height {
112+
row = max(0, d.height-dialogHeight)
113+
}
114+
115+
return row, col
116+
}
117+
118+
// View renders the max iterations confirmation dialog
119+
func (d *maxIterationsDialog) View() string {
120+
// clamped width: ~60% of screen, bounded by [36, 84] and screen margin
121+
dialogWidth := d.width * 60 / 100
122+
if dialogWidth < 36 {
123+
dialogWidth = max(20, min(d.width-4, 36))
124+
}
125+
if dialogWidth > 84 {
126+
dialogWidth = min(84, d.width-4)
127+
}
128+
129+
padX := 2
130+
padY := 1
131+
132+
// Border takes one character on each side when present
133+
frameHorizontal := (padX * 2) + 2
134+
contentWidth := max(10, dialogWidth-frameHorizontal)
135+
136+
dialogStyle := lipgloss.NewStyle().
137+
Border(lipgloss.RoundedBorder()).
138+
BorderForeground(lipgloss.Color("#f59e0b")).
139+
Foreground(lipgloss.Color("#d1d5db")).
140+
Padding(padY, padX).
141+
Width(dialogWidth).
142+
Align(lipgloss.Left)
143+
144+
titleStyle := lipgloss.NewStyle().
145+
Bold(true).
146+
Foreground(lipgloss.Color("#f59e0b")).
147+
Align(lipgloss.Center).
148+
Width(contentWidth)
149+
150+
messageStyle := lipgloss.NewStyle().
151+
Foreground(lipgloss.Color("#d1d5db")).
152+
Width(contentWidth)
153+
154+
questionStyle := lipgloss.NewStyle().
155+
Bold(true).
156+
Foreground(lipgloss.Color("#d1d5db")).
157+
Align(lipgloss.Center).
158+
Width(contentWidth)
159+
160+
optionsStyle := lipgloss.NewStyle().
161+
Foreground(lipgloss.Color("#9ca3af")).
162+
Align(lipgloss.Center).
163+
Width(contentWidth)
164+
165+
title := titleStyle.Render("Maximum Iterations Reached")
166+
167+
separatorWidth := max(1, contentWidth)
168+
separator := lipgloss.NewStyle().
169+
Foreground(lipgloss.Color("#4b5563")).
170+
Align(lipgloss.Center).
171+
Width(contentWidth).
172+
Render(strings.Repeat("─", separatorWidth))
173+
174+
// Info section
175+
infoText := fmt.Sprintf("Max Iterations: %d", d.maxIterations)
176+
infoWrapped := wrapDisplayText(infoText, contentWidth)
177+
infoSection := lipgloss.NewStyle().
178+
Foreground(lipgloss.Color("#d1d5db")).
179+
Render(infoWrapped)
180+
181+
// Message section
182+
message := messageStyle.Render(wrapDisplayText("The agent may be stuck in a loop. This can happen with smaller or less capable models.", contentWidth))
183+
184+
// Question section
185+
question := questionStyle.Render(wrapDisplayText("Do you want to continue for 10 more iterations?", contentWidth))
186+
187+
// Options section
188+
options := optionsStyle.Render(wrapDisplayText("[Y]es [N]o", contentWidth))
189+
190+
// Combine all parts with proper spacing
191+
parts := []string{title, separator, infoSection, "", message, "", question, "", options}
192+
193+
content := lipgloss.JoinVertical(lipgloss.Left, parts...)
194+
return dialogStyle.Render(content)
195+
}
196+
197+
// width-aware wrapping based on display cell width
198+
func wrapDisplayText(text string, maxWidth int) string {
199+
if maxWidth <= 0 {
200+
return text
201+
}
202+
words := strings.Fields(text)
203+
if len(words) == 0 {
204+
return text
205+
}
206+
var lines []string
207+
var current string
208+
for _, w := range words {
209+
if lipgloss.Width(current) == 0 {
210+
current = w
211+
continue
212+
}
213+
if lipgloss.Width(current+" "+w) <= maxWidth {
214+
current += " " + w
215+
} else {
216+
lines = append(lines, current)
217+
current = w
218+
}
219+
}
220+
if current != "" {
221+
lines = append(lines, current)
222+
}
223+
return strings.Join(lines, "\n")
224+
}

pkg/tui/page/chat/chat.go

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -257,6 +257,15 @@ func (p *chatPage) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
257257
spinnerCmd := p.setWorking(true)
258258
cmd := p.messages.AddToolResult(msg, types.ToolStatusCompleted)
259259
return p, tea.Batch(cmd, p.messages.ScrollToBottom(), spinnerCmd)
260+
case *runtime.MaxIterationsReachedEvent:
261+
spinnerCmd := p.setWorking(false) // Stop working indicator during confirmation
262+
263+
// Open max iterations confirmation dialog
264+
dialogCmd := core.CmdHandler(dialog.OpenDialogMsg{
265+
Model: dialog.NewMaxIterationsDialog(msg.MaxIterations, p.app),
266+
})
267+
268+
return p, tea.Batch(spinnerCmd, dialogCmd)
260269
}
261270

262271
sidebarModel, sidebarCmd := p.sidebar.Update(msg)

0 commit comments

Comments
 (0)