diff --git a/agent-schema.json b/agent-schema.json index 99140e597..b7d2e3aa4 100644 --- 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 index ab8ac2096..55a59cdc0 100644 --- 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 index 9de938481..969463eaa 100644 --- 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 index cdeef9264..f18c2e68b 100644 --- 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 index 3421e39f9..0578e5653 100644 --- 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 index 3f63bd95d..e66a598c6 100644 --- 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 index fc9315759..53eca84b5 100644 --- 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 index edb7278b2..581fff30c 100644 --- 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 index 3e1a75212..ecdee7be8 100644 --- 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 index 85b60432e..6c26d3b58 100644 --- 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 index 262785f23..b9740d4ac 100644 --- 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 index 4a9b16b99..8774d980e 100644 --- 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)),