From 58961413ed31242d71cde099bae38cc3c54b61cb Mon Sep 17 00:00:00 2001
From: stefanwalcz <stefan.walcz@walcz.de>
Date: Sun, 2 Aug 2026 12:18:21 +0200
Subject: [PATCH] =?UTF-8?q?feat(localai-client):=20SetMaxTokens=20?=
 =?UTF-8?q?=E2=80=94=20per-completion=20generation=20cap?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

cogito never set max_tokens on any request, so an agent-loop completion with no
server-side cap could run to full context (runaway narration/emoji loops that
neither loop-detection nor iteration limits can stop). Add a maxTokens field +
SetMaxTokens setter (parity with SetTemperature/SetReasoningEffort) and inject
request.MaxTokens on both CreateChatCompletion and CreateChatCompletionStream
when > 0. Zero leaves it unset (backend/YAML default applies).

Gate test: TestLocalAIClientSetMaxTokens asserts max_tokens is sent when set and
omitted when unset.

Signed-off-by: stefanwalcz <stefan.walcz@walcz.de>
---
 clients/localai_client.go      | 16 +++++++++++++
 clients/localai_client_test.go | 41 ++++++++++++++++++++++++++++++++++
 2 files changed, 57 insertions(+)

diff --git a/clients/localai_client.go b/clients/localai_client.go
index c23094d..b21b2a6 100644
--- a/clients/localai_client.go
+++ b/clients/localai_client.go
@@ -30,6 +30,7 @@ type LocalAIClient struct {
 	metadata        map[string]string
 	reasoningEffort string
 	temperature     float32
+	maxTokens       int
 	client          *http.Client
 
 	nativePartsMu sync.Mutex
@@ -84,6 +85,15 @@ func (llm *LocalAIClient) SetTemperature(temperature float32) {
 	llm.temperature = temperature
 }
 
+// SetMaxTokens caps the number of tokens generated per request (llama.cpp
+// n_predict / OpenAI max_tokens). Zero leaves the field unset (the backend's
+// own default applies). This is the per-completion backstop against runaway
+// generation (narration/emoji loops to context end) in the agent loop, which
+// neither loop-detection nor iteration limits can stop.
+func (llm *LocalAIClient) SetMaxTokens(n int) {
+	llm.maxTokens = n
+}
+
 // SetMetadata sets per-request metadata forwarded to LocalAI under the
 // top-level "metadata" object (e.g. {"enable_thinking": "true"}). Pass nil
 // or an empty map to clear. LocalAI uses these flags to override per-model
@@ -296,6 +306,9 @@ func (llm *LocalAIClient) CreateChatCompletion(ctx context.Context, request open
 	if llm.temperature != 0 {
 		request.Temperature = llm.temperature
 	}
+	if llm.maxTokens > 0 {
+		request.MaxTokens = llm.maxTokens
+	}
 
 	body, err := llm.marshalRequest(request)
 	if err != nil {
@@ -426,6 +439,9 @@ func (llm *LocalAIClient) CreateChatCompletionStream(ctx context.Context, reques
 	if llm.temperature != 0 {
 		request.Temperature = llm.temperature
 	}
+	if llm.maxTokens > 0 {
+		request.MaxTokens = llm.maxTokens
+	}
 
 	body, err := llm.marshalRequest(request)
 	if err != nil {
diff --git a/clients/localai_client_test.go b/clients/localai_client_test.go
index dd449e2..2e22877 100644
--- a/clients/localai_client_test.go
+++ b/clients/localai_client_test.go
@@ -167,3 +167,44 @@ func TestLocalAIClientSetTemperature(t *testing.T) {
 		t.Fatalf("request temperature = %v, want 0.7", gotTemperature)
 	}
 }
+
+// TestLocalAIClientSetMaxTokens proves SetMaxTokens injects max_tokens into the
+// request body (the per-completion runaway cap), and that leaving it unset omits
+// the field. Covers the non-stream path; the injection is shared with the stream
+// path (same block in CreateChatCompletionStream).
+func TestLocalAIClientSetMaxTokens(t *testing.T) {
+	var gotBody map[string]any
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		b, _ := io.ReadAll(r.Body)
+		_ = json.Unmarshal(b, &gotBody)
+		w.Header().Set("Content-Type", "application/json")
+		_, _ = w.Write([]byte(`{"choices":[{"index":0,"message":{"role":"assistant","content":"ok"}}]}`))
+	}))
+	defer srv.Close()
+
+	// with cap set
+	llm := NewLocalAILLM("m", "k", srv.URL+"/v1")
+	llm.SetMaxTokens(4096)
+	_, _, err := llm.CreateChatCompletion(context.Background(), openai.ChatCompletionRequest{
+		Messages: []openai.ChatCompletionMessage{{Role: "user", Content: "hi"}},
+	})
+	if err != nil {
+		t.Fatalf("CreateChatCompletion: %v", err)
+	}
+	if mt, ok := gotBody["max_tokens"]; !ok || int(mt.(float64)) != 4096 {
+		t.Fatalf("max_tokens = %v (ok=%v), want 4096", gotBody["max_tokens"], ok)
+	}
+
+	// without cap -> field omitted (omitempty)
+	gotBody = nil
+	llm2 := NewLocalAILLM("m", "k", srv.URL+"/v1")
+	_, _, err = llm2.CreateChatCompletion(context.Background(), openai.ChatCompletionRequest{
+		Messages: []openai.ChatCompletionMessage{{Role: "user", Content: "hi"}},
+	})
+	if err != nil {
+		t.Fatalf("CreateChatCompletion (no cap): %v", err)
+	}
+	if _, ok := gotBody["max_tokens"]; ok {
+		t.Fatalf("max_tokens present without SetMaxTokens: %v", gotBody["max_tokens"])
+	}
+}