Merge pull request #2192 from gtardif/configurable_max_tool_call_token

dgageot · web-flow · commit 9537c5be111b · 2026-03-20T15:13:50.000+01:00
feat: make maxOldToolCallTokens configurable in agent YAML
diff --git a/agent-schema.json b/agent-schema.json
@@ -187,6 +187,11 @@
           "description": "Maximum consecutive identical tool calls before the agent is terminated. Prevents degenerate loops. 0 uses the default of 5.",
           "minimum": 0
         },
+        "max_old_tool_call_tokens": {
+          "type": "integer",
+          "description": "Maximum number of tokens to keep from old tool call arguments and results. Older tool calls beyond this budget will have their content replaced with a placeholder. Tokens are approximated as len/4. Set to -1 to disable truncation (unlimited tool content). Default: 40000.",
+          "minimum": -1
+        },
         "num_history_items": {
           "type": "integer",
           "description": "Number of history items to keep",
diff --git a/cmd/root/run.go b/cmd/root/run.go
@@ -17,6 +17,7 @@ import (
 	"github.com/spf13/cobra"
 	"go.opentelemetry.io/otel"
 
+	"github.com/docker/docker-agent/pkg/agent"
 	"github.com/docker/docker-agent/pkg/app"
 	"github.com/docker/docker-agent/pkg/cli"
 	"github.com/docker/docker-agent/pkg/config"
@@ -339,7 +340,7 @@ func (f *runExecFlags) createRemoteRuntimeAndSession(ctx context.Context, origin
 func (f *runExecFlags) createLocalRuntimeAndSession(ctx context.Context, loadResult *teamloader.LoadResult) (runtime.Runtime, *session.Session, error) {
 	t := loadResult.Team
 
-	agent, err := t.Agent(f.agentName)
+	agt, err := t.Agent(f.agentName)
 	if err != nil {
 		return nil, nil, err
 	}
@@ -404,7 +405,7 @@ func (f *runExecFlags) createLocalRuntimeAndSession(ctx context.Context, loadRes
 		slog.Debug("Loaded existing session", "session_id", resolvedID, "session_ref", f.sessionID, "agent", f.agentName)
 	} else {
 		wd, _ := os.Getwd()
-		sess = session.New(f.buildSessionOpts(agent.MaxIterations(), wd)...)
+		sess = session.New(f.buildSessionOpts(agt, wd)...)
 		// Session is stored lazily on first UpdateSession call (when content is added)
 		// This avoids creating empty sessions in the database
 		slog.Debug("Using local runtime", "agent", f.agentName)
@@ -494,9 +495,11 @@ func (f *runExecFlags) buildAppOpts(args []string) ([]app.Opt, error) {
 // buildSessionOpts returns the canonical set of session options derived from
 // CLI flags and agent configuration. Both the initial session and spawned
 // sessions use this method so their options never drift apart.
-func (f *runExecFlags) buildSessionOpts(maxIterations int, workingDir string) []session.Opt {
+func (f *runExecFlags) buildSessionOpts(agt *agent.Agent, workingDir string) []session.Opt {
 	return []session.Opt{
-		session.WithMaxIterations(maxIterations),
+		session.WithMaxIterations(agt.MaxIterations()),
+		session.WithMaxConsecutiveToolCalls(agt.MaxConsecutiveToolCalls()),
+		session.WithMaxOldToolCallTokens(agt.MaxOldToolCallTokens()),
 		session.WithToolsApproved(f.autoApprove),
 		session.WithHideToolResults(f.hideToolResults),
 		session.WithWorkingDir(workingDir),
@@ -517,7 +520,7 @@ func (f *runExecFlags) createSessionSpawner(agentSource config.Source, sessStore
 		}
 
 		team := loadResult.Team
-		agent, err := team.Agent(f.agentName)
+		agt, err := team.Agent(f.agentName)
 		if err != nil {
 			return nil, nil, nil, err
 		}
@@ -543,7 +546,7 @@ func (f *runExecFlags) createSessionSpawner(agentSource config.Source, sessStore
 		}
 
 		// Create a new session
-		newSess := session.New(f.buildSessionOpts(agent.MaxIterations(), workingDir)...)
+		newSess := session.New(f.buildSessionOpts(agt, workingDir)...)
 
 		// Create cleanup function
 		cleanup := func() {
diff --git a/pkg/a2a/adapter.go b/pkg/a2a/adapter.go
@@ -47,6 +47,7 @@ func runDockerAgent(ctx agent.InvocationContext, t *team.Team, agentName string,
 			session.WithUserMessage(message),
 			session.WithMaxIterations(a.MaxIterations()),
 			session.WithMaxConsecutiveToolCalls(a.MaxConsecutiveToolCalls()),
+			session.WithMaxOldToolCallTokens(a.MaxOldToolCallTokens()),
 			session.WithToolsApproved(true),
 		)
 
diff --git a/pkg/acp/agent.go b/pkg/acp/agent.go
@@ -154,6 +154,7 @@ func (a *Agent) NewSession(ctx context.Context, params acp.NewSessionRequest) (a
 	sessOpts := []session.Opt{
 		session.WithMaxIterations(rootAgent.MaxIterations()),
 		session.WithMaxConsecutiveToolCalls(rootAgent.MaxConsecutiveToolCalls()),
+		session.WithMaxOldToolCallTokens(rootAgent.MaxOldToolCallTokens()),
 	}
 	if workingDir != "" {
 		sessOpts = append(sessOpts, session.WithWorkingDir(workingDir))
diff --git a/pkg/agent/agent.go b/pkg/agent/agent.go
@@ -34,6 +34,7 @@ type Agent struct {
 	addDescriptionParameter bool
 	maxIterations           int
 	maxConsecutiveToolCalls int
+	maxOldToolCallTokens    int
 	numHistoryItems         int
 	addPromptFiles          []string
 	tools                   []tools.Tool
@@ -81,6 +82,10 @@ func (a *Agent) MaxConsecutiveToolCalls() int {
 	return a.maxConsecutiveToolCalls
 }
 
+func (a *Agent) MaxOldToolCallTokens() int {
+	return a.maxOldToolCallTokens
+}
+
 func (a *Agent) NumHistoryItems() int {
 	return a.numHistoryItems
 }
diff --git a/pkg/agent/opts.go b/pkg/agent/opts.go
@@ -138,6 +138,15 @@ func WithMaxConsecutiveToolCalls(n int) Opt {
 	}
 }
 
+// WithMaxOldToolCallTokens sets the maximum token budget for old tool call content.
+// Set to -1 to disable truncation (unlimited tool content).
+// Set to 0 to use the default (40000).
+func WithMaxOldToolCallTokens(n int) Opt {
+	return func(a *Agent) {
+		a.maxOldToolCallTokens = n
+	}
+}
+
 func WithNumHistoryItems(numHistoryItems int) Opt {
 	return func(a *Agent) {
 		a.numHistoryItems = numHistoryItems
diff --git a/pkg/config/latest/types.go b/pkg/config/latest/types.go
@@ -252,6 +252,7 @@ type AgentConfig struct {
 	AddDescriptionParameter bool              `json:"add_description_parameter,omitempty"`
 	MaxIterations           int               `json:"max_iterations,omitempty"`
 	MaxConsecutiveToolCalls int               `json:"max_consecutive_tool_calls,omitempty"`
+	MaxOldToolCallTokens    int               `json:"max_old_tool_call_tokens,omitempty"`
 	NumHistoryItems         int               `json:"num_history_items,omitempty"`
 	AddPromptFiles          []string          `json:"add_prompt_files,omitempty" yaml:"add_prompt_files,omitempty"`
 	Commands                types.Commands    `json:"commands,omitempty"`
diff --git a/pkg/mcp/server.go b/pkg/mcp/server.go
@@ -161,6 +161,7 @@ func CreateToolHandler(t *team.Team, agentName string) func(context.Context, *mc
 			session.WithTitle("MCP tool call"),
 			session.WithMaxIterations(ag.MaxIterations()),
 			session.WithMaxConsecutiveToolCalls(ag.MaxConsecutiveToolCalls()),
+			session.WithMaxOldToolCallTokens(ag.MaxOldToolCallTokens()),
 			session.WithUserMessage(input.Message),
 			session.WithToolsApproved(true),
 		)
diff --git a/pkg/runtime/agent_delegation.go b/pkg/runtime/agent_delegation.go
@@ -106,6 +106,7 @@ func newSubSession(parent *session.Session, cfg SubSessionConfig, childAgent *ag
 		session.WithImplicitUserMessage(userMsg),
 		session.WithMaxIterations(childAgent.MaxIterations()),
 		session.WithMaxConsecutiveToolCalls(childAgent.MaxConsecutiveToolCalls()),
+		session.WithMaxOldToolCallTokens(childAgent.MaxOldToolCallTokens()),
 		session.WithTitle(cfg.Title),
 		session.WithToolsApproved(cfg.ToolsApproved),
 		session.WithSendUserMessage(false),
diff --git a/pkg/server/session_manager.go b/pkg/server/session_manager.go
@@ -77,6 +77,7 @@ func (sm *SessionManager) CreateSession(ctx context.Context, sessionTemplate *se
 	opts = append(opts,
 		session.WithMaxIterations(sessionTemplate.MaxIterations),
 		session.WithMaxConsecutiveToolCalls(sessionTemplate.MaxConsecutiveToolCalls),
+		session.WithMaxOldToolCallTokens(sessionTemplate.MaxOldToolCallTokens),
 		session.WithToolsApproved(sessionTemplate.ToolsApproved),
 	)
 
@@ -347,6 +348,7 @@ func (sm *SessionManager) runtimeForSession(ctx context.Context, sess *session.S
 	}
 	sess.MaxIterations = agent.MaxIterations()
 	sess.MaxConsecutiveToolCalls = agent.MaxConsecutiveToolCalls()
+	sess.MaxOldToolCallTokens = agent.MaxOldToolCallTokens()
 
 	opts := []runtime.Opt{
 		runtime.WithCurrentAgent(currentAgent),
diff --git a/pkg/session/session.go b/pkg/session/session.go
@@ -18,10 +18,10 @@ import (
 )
 
 const (
-	// MaxToolCallTokens is the maximum number of tokens to keep from tool call
+	// DefaultMaxOldToolCallTokens is the default maximum number of tokens to keep from tool call
 	// arguments and results. Older tool calls beyond this budget will have their
 	// content replaced with a placeholder. Tokens are approximated as len/4.
-	MaxToolCallTokens = 40000
+	DefaultMaxOldToolCallTokens = 40000
 
 	// toolContentPlaceholder is the text used to replace truncated tool content
 	toolContentPlaceholder = "[content truncated]"
@@ -94,6 +94,13 @@ type Session struct {
 	// repeatedly issues the same call without making progress. Default: 5.
 	MaxConsecutiveToolCalls int `json:"max_consecutive_tool_calls,omitempty"`
 
+	// MaxOldToolCallTokens is the maximum number of tokens to keep from old tool call
+	// arguments and results. Older tool calls beyond this budget will have their
+	// content replaced with a placeholder. Tokens are approximated as len/4.
+	// Set to -1 to disable truncation (unlimited tool content).
+	// Default: 40000 (when not configured or set to 0).
+	MaxOldToolCallTokens int `json:"max_old_tool_call_tokens,omitempty"`
+
 	// Starred indicates if this session has been starred by the user
 	Starred bool `json:"starred"`
 
@@ -442,6 +449,15 @@ func WithMaxConsecutiveToolCalls(n int) Opt {
 	}
 }
 
+// WithMaxOldToolCallTokens sets the maximum token budget for old tool call content.
+// Set to -1 to disable truncation (unlimited tool content).
+// Set to 0 to use the default (40000).
+func WithMaxOldToolCallTokens(n int) Opt {
+	return func(s *Session) {
+		s.MaxOldToolCallTokens = n
+	}
+}
+
 func WithWorkingDir(workingDir string) Opt {
 	return func(s *Session) {
 		s.WorkingDir = workingDir
@@ -781,7 +797,15 @@ func (s *Session) GetMessages(a *agent.Agent) []chat.Message {
 		messages = trimMessages(messages, maxItems)
 	}
 
-	messages = truncateOldToolContent(messages, MaxToolCallTokens)
+	// Use configured max tokens or fall back to default constant if zero or unset.
+	// -1 means unlimited (no truncation).
+	maxOldToolCallTokens := s.MaxOldToolCallTokens
+	if maxOldToolCallTokens == 0 {
+		maxOldToolCallTokens = DefaultMaxOldToolCallTokens
+	}
+	if maxOldToolCallTokens > 0 { // If maxOldToolCallTokens is -1, skip truncation (unlimited)
+		messages = truncateOldToolContent(messages, maxOldToolCallTokens)
+	}
 
 	systemCount := 0
 	conversationCount := 0
diff --git a/pkg/teamloader/teamloader.go b/pkg/teamloader/teamloader.go
@@ -171,6 +171,7 @@ func LoadWithConfig(ctx context.Context, agentSource config.Source, runConfig *c
 			agent.WithAddPromptFiles(promptFiles),
 			agent.WithMaxIterations(agentConfig.MaxIterations),
 			agent.WithMaxConsecutiveToolCalls(agentConfig.MaxConsecutiveToolCalls),
+			agent.WithMaxOldToolCallTokens(agentConfig.MaxOldToolCallTokens),
 			agent.WithNumHistoryItems(agentConfig.NumHistoryItems),
 			agent.WithCommands(expander.ExpandCommands(ctx, agentConfig.Commands)),
 			agent.WithHooks(config.MergeHooks(agentConfig.Hooks, cliHooks)),

Original file line number	Diff line number	Diff line change
`@@ -47,6 +47,7 @@ func runDockerAgent(ctx agent.InvocationContext, t *team.Team, agentName string,`
`47`	`47`	`session.WithUserMessage(message),`
`48`	`48`	`session.WithMaxIterations(a.MaxIterations()),`
`49`	`49`	`session.WithMaxConsecutiveToolCalls(a.MaxConsecutiveToolCalls()),`
	`50`	`+ session.WithMaxOldToolCallTokens(a.MaxOldToolCallTokens()),`
`50`	`51`	`session.WithToolsApproved(true),`
`51`	`52`	`)`
`52`	`53`
Original file line number	Diff line number	Diff line change
`@@ -154,6 +154,7 @@ func (a *Agent) NewSession(ctx context.Context, params acp.NewSessionRequest) (a`
`154`	`154`	`sessOpts := []session.Opt{`
`155`	`155`	`session.WithMaxIterations(rootAgent.MaxIterations()),`
`156`	`156`	`session.WithMaxConsecutiveToolCalls(rootAgent.MaxConsecutiveToolCalls()),`
	`157`	`+ session.WithMaxOldToolCallTokens(rootAgent.MaxOldToolCallTokens()),`
`157`	`158`	`}`
`158`	`159`	`if workingDir != "" {`
`159`	`160`	`sessOpts = append(sessOpts, session.WithWorkingDir(workingDir))`
Original file line number	Diff line number	Diff line change
`@@ -161,6 +161,7 @@ func CreateToolHandler(t team.Team, agentName string) func(context.Context, mc`
`161`	`161`	`session.WithTitle("MCP tool call"),`
`162`	`162`	`session.WithMaxIterations(ag.MaxIterations()),`
`163`	`163`	`session.WithMaxConsecutiveToolCalls(ag.MaxConsecutiveToolCalls()),`
	`164`	`+ session.WithMaxOldToolCallTokens(ag.MaxOldToolCallTokens()),`
`164`	`165`	`session.WithUserMessage(input.Message),`
`165`	`166`	`session.WithToolsApproved(true),`
`166`	`167`	`)`
Original file line number	Diff line number	Diff line change
`@@ -77,6 +77,7 @@ func (sm SessionManager) CreateSession(ctx context.Context, sessionTemplate se`
`77`	`77`	`opts = append(opts,`
`78`	`78`	`session.WithMaxIterations(sessionTemplate.MaxIterations),`
`79`	`79`	`session.WithMaxConsecutiveToolCalls(sessionTemplate.MaxConsecutiveToolCalls),`
	`80`	`+ session.WithMaxOldToolCallTokens(sessionTemplate.MaxOldToolCallTokens),`
`80`	`81`	`session.WithToolsApproved(sessionTemplate.ToolsApproved),`
`81`	`82`	`)`
`82`	`83`
`@@ -347,6 +348,7 @@ func (sm SessionManager) runtimeForSession(ctx context.Context, sess session.S`
`347`	`348`	`}`
`348`	`349`	`sess.MaxIterations = agent.MaxIterations()`
`349`	`350`	`sess.MaxConsecutiveToolCalls = agent.MaxConsecutiveToolCalls()`
	`351`	`+ sess.MaxOldToolCallTokens = agent.MaxOldToolCallTokens()`
`350`	`352`
`351`	`353`	`opts := []runtime.Opt{`
`352`	`354`	`runtime.WithCurrentAgent(currentAgent),`