feat(agents): per-model output cap for agent job completions

Agent jobs (/api/agent/jobs, local and on distributed agent workers) run
cogito against a LocalAI client that sends no max_tokens of its own, so
every completion gets the client's fixed default of 16384. max_iterations
bounds the number of turns and loop_detection repeated tool calls, but
neither stops one completion that keeps generating, and 16384 is the same
for a small tool-picking model and a reasoning model that thinks at length.

Add agent.max_generation_tokens to the model config. When set, both job
paths apply it to the client with SetMaxTokens; requests that carry their
own cap still win, so cogito's length retry can raise it. Unset keeps
today's behaviour.

The specs run each job path against an MCP fixture and a recording LLM
endpoint and check the max_tokens the completions carry; they fail when
either call site drops the cap.

Assisted-by: Claude Code:claude-opus-5-5
Signed-off-by: Stefan Walcz <stefan.walcz@walcz.de>
(cherry picked from commit 5600a1cabe)

core/cli/agent_worker.go

@@ -445,6 +445,7 @@ func handleMCPCIJob(shutdownCtx context.Context, data []byte, apiURL, apiToken s
 
 	// Create LLM client pointing back to the frontend API
 	llm := clients.NewLocalAILLM(task.Model, apiToken, apiURL)
+	modelCfg.ApplyAgentGenerationCap(llm)
 
 	// Build cogito options
 	ctx, cancel := context.WithTimeout(shutdownCtx, jobTimeout)

core/cli/agent_worker_generation_cap_test.go

@@ -0,0 +1,98 @@
+package cli
+
+import (
+	"context"
+	"encoding/json"
+	"fmt"
+	"net/http"
+	"net/http/httptest"
+	"sync"
+	"time"
+
+	"github.com/modelcontextprotocol/go-sdk/mcp"
+	"gopkg.in/yaml.v3"
+
+	"github.com/mudler/LocalAI/core/config"
+	mcpTools "github.com/mudler/LocalAI/core/http/endpoints/mcp"
+	"github.com/mudler/LocalAI/core/services/jobs"
+	"github.com/mudler/LocalAI/core/services/testutil"
+
+	. "github.com/onsi/ginkgo/v2"
+	. "github.com/onsi/gomega"
+)
+
+var _ = Describe("handleMCPCIJob generation cap", func() {
+	var (
+		mcpFixture *httptest.Server
+		llmAPI     *httptest.Server
+		mu         sync.Mutex
+		sentCaps   []int
+	)
+
+	BeforeEach(func() {
+		srv := mcp.NewServer(&mcp.Implementation{Name: "fixture", Version: "v0.0.1"}, nil)
+		srv.AddTool(&mcp.Tool{
+			Name:        "noop",
+			Description: "fixture tool",
+			InputSchema: json.RawMessage(`{"type":"object","properties":{}}`),
+		}, func(context.Context, *mcp.CallToolRequest) (*mcp.CallToolResult, error) {
+			return &mcp.CallToolResult{Content: []mcp.Content{&mcp.TextContent{Text: "ok"}}}, nil
+		})
+		mcpFixture = httptest.NewServer(mcp.NewStreamableHTTPHandler(func(*http.Request) *mcp.Server { return srv }, nil))
+
+		// The LLM API records the output cap of every completion the job sends
+		// and then fails it: what this spec pins is the request, not the run.
+		sentCaps = nil
+		llmAPI = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			var body struct {
+				MaxTokens int `json:"max_tokens"`
+			}
+			if err := json.NewDecoder(r.Body).Decode(&body); err == nil {
+				mu.Lock()
+				sentCaps = append(sentCaps, body.MaxTokens)
+				mu.Unlock()
+			}
+			http.Error(w, "fixture", http.StatusBadRequest)
+		}))
+	})
+
+	AfterEach(func() {
+		mcpTools.CloseMCPSessions("capped")
+		mcpTools.CloseMCPSessions("uncapped")
+		llmAPI.Close()
+		mcpFixture.Close()
+	})
+
+	runJob := func(name, agentYAML string) []int {
+		var cfg config.ModelConfig
+		Expect(yaml.Unmarshal([]byte(fmt.Sprintf("name: %s\n%smcp:\n  remote: |\n    {\"mcpServers\": {\"fixture\": {\"url\": %q}}}\n", name, agentYAML, mcpFixture.URL)), &cfg)).To(Succeed())
+
+		payload, err := json.Marshal(jobs.JobEvent{
+			JobID:       "job-" + name,
+			TaskID:      "task-" + name,
+			Job:         &jobs.JobRecord{ID: "job-" + name},
+			Task:        &jobs.TaskRecord{ID: "task-" + name, Model: name, Prompt: "do it"},
+			ModelConfig: &cfg,
+		})
+		Expect(err).ToNot(HaveOccurred())
+
+		Expect(handleMCPCIJob(context.Background(), payload, llmAPI.URL, "", testutil.NewFakeBus(), 2*time.Second)).To(Succeed())
+
+		mu.Lock()
+		defer mu.Unlock()
+		return append([]int(nil), sentCaps...)
+	}
+
+	It("sends agent.max_generation_tokens as the output cap of the job's completions", func() {
+		caps := runJob("capped", "agent:\n  max_generation_tokens: 1234\n")
+		Expect(caps).ToNot(BeEmpty())
+		Expect(caps).To(HaveEach(1234))
+	})
+
+	It("leaves the client default in place when the model sets no cap", func() {
+		caps := runJob("uncapped", "")
+		Expect(caps).ToNot(BeEmpty())
+		Expect(caps).ToNot(ContainElement(1234))
+		Expect(caps).To(HaveEach(BeNumerically(">", 0)))
+	})
+})

core/config/agent_generation_cap_test.go

@@ -0,0 +1,81 @@
+package config_test
+
+import (
+	"context"
+	"encoding/json"
+	"net/http"
+	"net/http/httptest"
+	"sync"
+
+	"github.com/mudler/cogito/clients"
+	"github.com/sashabaranov/go-openai"
+	"gopkg.in/yaml.v3"
+
+	"github.com/mudler/LocalAI/core/config"
+
+	. "github.com/onsi/ginkgo/v2"
+	. "github.com/onsi/gomega"
+)
+
+var _ = Describe("agent.max_generation_tokens", func() {
+	var (
+		srv  *httptest.Server
+		mu   sync.Mutex
+		caps []int
+	)
+
+	BeforeEach(func() {
+		caps = nil
+		srv = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			var body struct {
+				MaxTokens int `json:"max_tokens"`
+			}
+			Expect(json.NewDecoder(r.Body).Decode(&body)).To(Succeed())
+			mu.Lock()
+			caps = append(caps, body.MaxTokens)
+			mu.Unlock()
+			w.Header().Set("Content-Type", "application/json")
+			_, _ = w.Write([]byte(`{"choices":[{"index":0,"finish_reason":"stop","message":{"role":"assistant","content":"ok"}}]}`))
+		}))
+		DeferCleanup(srv.Close)
+	})
+
+	// agentLLM builds the client the way the agent job paths do: the model
+	// config comes from YAML, the client points at the LocalAI API, and the
+	// config's cap is applied before the loop runs.
+	agentLLM := func(modelYAML string) *clients.LocalAIClient {
+		var cfg config.ModelConfig
+		Expect(yaml.Unmarshal([]byte(modelYAML), &cfg)).To(Succeed())
+		llm := clients.NewLocalAILLM(cfg.Name, "", srv.URL)
+		cfg.ApplyAgentGenerationCap(llm)
+		return llm
+	}
+
+	sentCap := func(llm *clients.LocalAIClient, req openai.ChatCompletionRequest) int {
+		req.Messages = []openai.ChatCompletionMessage{{Role: "user", Content: "hi"}}
+		_, _, err := llm.CreateChatCompletion(context.Background(), req)
+		Expect(err).ToNot(HaveOccurred())
+		mu.Lock()
+		defer mu.Unlock()
+		Expect(caps).ToNot(BeEmpty())
+		return caps[len(caps)-1]
+	}
+
+	It("sends the configured cap on every agent completion", func() {
+		llm := agentLLM("name: m\nagent:\n  max_generation_tokens: 2048\n")
+		Expect(sentCap(llm, openai.ChatCompletionRequest{})).To(Equal(2048))
+		Expect(sentCap(llm, openai.ChatCompletionRequest{})).To(Equal(2048))
+	})
+
+	It("keeps the client default when the model sets no cap", func() {
+		bare := sentCap(clients.NewLocalAILLM("m", "", srv.URL), openai.ChatCompletionRequest{})
+		Expect(bare).To(BeNumerically(">", 0))
+		llm := agentLLM("name: m\n")
+		Expect(sentCap(llm, openai.ChatCompletionRequest{})).To(Equal(bare))
+	})
+
+	It("lets a request that carries its own cap win, so cogito's length retry can raise it", func() {
+		llm := agentLLM("name: m\nagent:\n  max_generation_tokens: 2048\n")
+		Expect(sentCap(llm, openai.ChatCompletionRequest{MaxTokens: 4096})).To(Equal(4096))
+	})
+})

core/config/meta/registry.go

@@ -946,6 +946,15 @@ func DefaultRegistry() map[string]FieldMetaOverride {
 		},
 
 		// --- MCP ---
+		"agent.max_generation_tokens": {
+			Section:     "agent",
+			Label:       "Max Generation Tokens",
+			Description: "Output-token cap for each completion of an agent job (/api/agent/jobs). Bounds a single completion that keeps generating, which max_iterations and loop_detection do not. Raise it for reasoning models that think at length. 0 or empty keeps the client default (16384).",
+			Component:   "number",
+			Min:         f64(0),
+			Order:       120,
+			Advanced:    true,
+		},
 		"mcp.remote": {
 			Section:     "mcp",
 			Label:       "Remote MCP Servers",

core/config/model_config.go

@@ -920,6 +920,22 @@ type AgentConfig struct {
 	LoopDetection         int  `yaml:"loop_detection,omitempty" json:"loop_detection,omitempty"`
 	MaxAdjustmentAttempts int  `yaml:"max_adjustment_attempts,omitempty" json:"max_adjustment_attempts,omitempty"`
 	ForceReasoningTool    bool `yaml:"force_reasoning_tool,omitempty" json:"force_reasoning_tool,omitempty"`
+	// MaxGenerationTokens caps the tokens a single completion may generate in
+	// the agent loop. Neither max_iterations nor loop_detection stops one
+	// completion that keeps generating; without a per-model value every model
+	// gets the client's fixed default, which is too low for some reasoning
+	// models and far above what others need. 0 keeps the client default.
+	MaxGenerationTokens int `yaml:"max_generation_tokens,omitempty" json:"max_generation_tokens,omitempty"`
+}
+
+// ApplyAgentGenerationCap puts agent.max_generation_tokens on the LLM client
+// the agent loop uses. The client applies it to every request that carries no
+// cap of its own, so a larger cap cogito sets when it retries a completion
+// truncated by length still wins. Unset leaves the client default in place.
+func (c *ModelConfig) ApplyAgentGenerationCap(llm interface{ SetMaxTokens(int) }) {
+	if c.Agent.MaxGenerationTokens > 0 {
+		llm.SetMaxTokens(c.Agent.MaxGenerationTokens)
+	}
 }
 
 // HasMCPServers returns true if any MCP servers (remote or stdio) are configured.

core/services/agentpool/agent_jobs.go

@@ -900,6 +900,7 @@ func (s *AgentJobService) ExecuteJobInternal(job schema.Job, task schema.Task, c
 
 	// Create LLM client
 	defaultLLM := clients.NewLocalAILLM(modelConfig.Name, apiKey, "http://127.0.0.1:"+port)
+	modelConfig.ApplyAgentGenerationCap(defaultLLM)
 
 	// Initialize traces slice
 	job.Traces = []schema.JobTrace{}

core/services/agentpool/agent_jobs_generation_cap_test.go

@@ -0,0 +1,120 @@
+package agentpool_test
+
+import (
+	"context"
+	"encoding/json"
+	"fmt"
+	"net/http"
+	"net/http/httptest"
+	"net/url"
+	"os"
+	"path/filepath"
+	"sync"
+	"time"
+
+	"github.com/modelcontextprotocol/go-sdk/mcp"
+
+	"github.com/mudler/LocalAI/core/config"
+	mcpTools "github.com/mudler/LocalAI/core/http/endpoints/mcp"
+	"github.com/mudler/LocalAI/core/schema"
+	"github.com/mudler/LocalAI/core/services/agentpool"
+	"github.com/mudler/LocalAI/core/templates"
+	"github.com/mudler/LocalAI/pkg/model"
+	"github.com/mudler/LocalAI/pkg/system"
+
+	. "github.com/onsi/ginkgo/v2"
+	. "github.com/onsi/gomega"
+)
+
+var _ = Describe("AgentJobService generation cap", func() {
+	var (
+		tempDir    string
+		mcpFixture *httptest.Server
+		llmAPI     *httptest.Server
+		mu         sync.Mutex
+		sentCaps   []int
+	)
+
+	BeforeEach(func() {
+		var err error
+		tempDir, err = os.MkdirTemp("", "agent_jobs_cap_test")
+		Expect(err).NotTo(HaveOccurred())
+
+		srv := mcp.NewServer(&mcp.Implementation{Name: "fixture", Version: "v0.0.1"}, nil)
+		srv.AddTool(&mcp.Tool{
+			Name:        "noop",
+			Description: "fixture tool",
+			InputSchema: json.RawMessage(`{"type":"object","properties":{}}`),
+		}, func(context.Context, *mcp.CallToolRequest) (*mcp.CallToolResult, error) {
+			return &mcp.CallToolResult{Content: []mcp.Content{&mcp.TextContent{Text: "ok"}}}, nil
+		})
+		mcpFixture = httptest.NewServer(mcp.NewStreamableHTTPHandler(func(*http.Request) *mcp.Server { return srv }, nil))
+
+		// Stands in for the LocalAI API the job calls back into: it records the
+		// output cap of every completion and fails it, since this spec pins the
+		// request rather than the run.
+		sentCaps = nil
+		llmAPI = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			var body struct {
+				MaxTokens int `json:"max_tokens"`
+			}
+			if err := json.NewDecoder(r.Body).Decode(&body); err == nil {
+				mu.Lock()
+				sentCaps = append(sentCaps, body.MaxTokens)
+				mu.Unlock()
+			}
+			http.Error(w, "fixture", http.StatusBadRequest)
+		}))
+	})
+
+	AfterEach(func() {
+		mcpTools.CloseMCPSessions("capped")
+		mcpTools.CloseMCPSessions("uncapped")
+		llmAPI.Close()
+		mcpFixture.Close()
+		_ = os.RemoveAll(tempDir)
+	})
+
+	runJob := func(name, agentYAML string) []int {
+		body := fmt.Sprintf("name: %s\n%smcp:\n  remote: |\n    {\"mcpServers\": {\"fixture\": {\"url\": %q}}}\n", name, agentYAML, mcpFixture.URL)
+		Expect(os.WriteFile(filepath.Join(tempDir, name+".yaml"), []byte(body), 0644)).To(Succeed())
+
+		systemState := &system.SystemState{}
+		systemState.Model.ModelsPath = tempDir
+		appConfig := config.NewApplicationConfig(
+			config.WithDynamicConfigDir(tempDir),
+			config.WithContext(context.Background()),
+		)
+		appConfig.SystemState = systemState
+		u, err := url.Parse(llmAPI.URL)
+		Expect(err).NotTo(HaveOccurred())
+		appConfig.APIAddress = u.Host
+
+		configLoader := config.NewModelConfigLoader(tempDir)
+		Expect(configLoader.LoadModelConfigsFromPath(tempDir)).To(Succeed())
+		service := agentpool.NewAgentJobService(appConfig, model.NewModelLoader(systemState), configLoader, templates.NewEvaluator(tempDir))
+
+		task := schema.Task{ID: "task-" + name, Name: name, Model: name, Prompt: "do it", Enabled: true}
+		job := schema.Job{ID: "job-" + name, TaskID: task.ID}
+		ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
+		defer cancel()
+		_ = service.ExecuteJobInternal(job, task, ctx)
+
+		mu.Lock()
+		defer mu.Unlock()
+		return append([]int(nil), sentCaps...)
+	}
+
+	It("sends agent.max_generation_tokens as the output cap of the job's completions", func() {
+		caps := runJob("capped", "agent:\n  max_generation_tokens: 1234\n")
+		Expect(caps).ToNot(BeEmpty())
+		Expect(caps).To(HaveEach(1234))
+	})
+
+	It("leaves the client default in place when the model sets no cap", func() {
+		caps := runJob("uncapped", "")
+		Expect(caps).ToNot(BeEmpty())
+		Expect(caps).ToNot(ContainElement(1234))
+		Expect(caps).To(HaveEach(BeNumerically(">", 0)))
+	})
+})

docs/content/advanced/model-configuration.md

@@ -1015,6 +1015,7 @@ Agent/autonomous agent configuration:
 | `agent.enable_planning` | bool | Enable planning capabilities |
 | `agent.enable_mcp_prompts` | bool | Enable MCP prompts |
 | `agent.enable_plan_re_evaluator` | bool | Enable plan re-evaluation |
+| `agent.max_generation_tokens` | int | Output-token cap for each completion in the agent loop (agent jobs, `/api/agent/jobs`). Bounds a single completion that keeps generating, which `max_iterations` and `loop_detection` do not. A request that already carries a cap keeps it, so a length retry can still raise it. `0` (default) keeps the client default of 16384 |
 
 ## Reasoning Configuration
 

docs/content/features/mcp.md

@@ -191,6 +191,7 @@ different tools share one loaded model.
 #### Agent Configuration (`agent`)
 
 - **`max_iterations`**: Maximum number of MCP tool execution loop iterations (default: 10). Each iteration allows the model to call tools and receive results before generating the next response.
+- **`max_generation_tokens`**: Output-token cap for each completion of an MCP agent job (default: the client default, 16384). `max_iterations` bounds the number of turns, not the length of one; this bounds a single completion that keeps generating. Raise it for reasoning models that think at length, lower it for models that should only pick a tool and answer briefly. It applies to agent jobs (`/api/agent/jobs`, also when run on a distributed agent worker); an MCP tool loop on `/v1/chat/completions` uses the request's own `max_tokens`.
 
 ## Agent-scoped MCP servers