From 2ea8721c37a3a04be05ec6f9e96415896ca3a60f Mon Sep 17 00:00:00 2001 From: Stefan Walcz Date: Fri, 2 Oct 2026 15:42:17 +0200 Subject: [PATCH 1/3] fix(tools): bound the forced reasoning steps by their own output cap The reasoning steps of WithForceReasoning (before the tool pick and before parameter generation) send no output cap of their own, so a model that never closes the reasoning argument runs to the client's global cap (16384 by default), and the length retry then doubles it. Observed with gemma4-26b-a4b behind LocalAI: one forced reasoning call streamed 12,969 tool-call deltas for 478-501 s until the cap, and the turn never answered. The reasoning steps now carry max_tokens from WithReasoningMaxTokens (default DefaultReasoningMaxTokens = 2048; 0 uses the default, a negative value keeps the previous uncapped behaviour). The cap goes on the request, where the bundled clients take it over their own default (see SetMaxTokens). A finish_reason=length on such a capped step is final: no length retry, the call returns an error wrapping errOutputCapReached; a complete tool call that is cut at the cap is still used. The parameter-reasoning step already falls back on any error; the tool-pick step now continues without the reasoning on this error only, other errors stay fatal. A reasoning that finishes normally is used as before. The option is forwarded to plan sub-executions. Signed-off-by: Stefan Walcz Assisted-by: Claude Opus 5.5 --- decision_output_cap_test.go | 84 ++++++++++++ options.go | 29 +++++ plan.go | 3 + reasoning_cap_test.go | 251 ++++++++++++++++++++++++++++++++++++ tools.go | 76 ++++++++++- 5 files changed, 436 insertions(+), 7 deletions(-) create mode 100644 decision_output_cap_test.go create mode 100644 reasoning_cap_test.go diff --git a/decision_output_cap_test.go b/decision_output_cap_test.go new file mode 100644 index 0000000..b2c7b0c --- /dev/null +++ b/decision_output_cap_test.go @@ -0,0 +1,84 @@ +package cogito + +import ( + "context" + "errors" + "testing" + + "github.com/sashabaranov/go-openai" +) + +// lengthCutLLM is a non-streaming fake whose reasoning call is always cut by +// the output cap, with the arguments left open, unless complete is set. +type lengthCutLLM struct { + complete bool + requests []openai.ChatCompletionRequest +} + +func (m *lengthCutLLM) Ask(ctx context.Context, f Fragment) (Fragment, error) { return f, nil } + +func (m *lengthCutLLM) CreateChatCompletion(ctx context.Context, req openai.ChatCompletionRequest) (LLMReply, LLMUsage, error) { + m.requests = append(m.requests, req) + limit := req.MaxTokens + if limit == 0 { + limit = 16384 + } + args := `{"reasoning":"a a a a` + if m.complete { + args = `{"reasoning":"short"}` + } + usage := LLMUsage{PromptTokens: 100, CompletionTokens: limit, TotalTokens: 100 + limit} + return LLMReply{ + ChatCompletionResponse: openai.ChatCompletionResponse{Choices: []openai.ChatCompletionChoice{{ + FinishReason: openai.FinishReasonLength, + Message: openai.ChatCompletionMessage{Role: "assistant", ToolCalls: []openai.ToolCall{{ + ID: "c1", Type: openai.ToolTypeFunction, + Function: openai.FunctionCall{Name: "reasoning", Arguments: args}, + }}}, + }}}, + MaxTokens: limit, + }, usage, nil +} + +var outputCapConv = []openai.ChatCompletionMessage{{Role: "user", Content: "find x"}} + +func TestDecisionCappedStopsAtItsOwnCap(t *testing.T) { + llm := &lengthCutLLM{} + _, err := decisionCapped(context.Background(), llm, outputCapConv, Tools{reasoningTool()}, "reasoning", 3, 512) + if !errors.Is(err, errOutputCapReached) { + t.Fatalf("expected errOutputCapReached, got %v", err) + } + if len(llm.requests) != 1 { + t.Fatalf("expected one attempt without a length retry, got %d", len(llm.requests)) + } + if got := llm.requests[0].MaxTokens; got != 512 { + t.Fatalf("expected the request to carry max tokens 512, got %d", got) + } +} + +func TestDecisionCappedKeepsACompleteCallCutAtTheCap(t *testing.T) { + llm := &lengthCutLLM{complete: true} + res, err := decisionCapped(context.Background(), llm, outputCapConv, Tools{reasoningTool()}, "reasoning", 3, 512) + if err != nil { + t.Fatalf("expected the complete call to be used, got %v", err) + } + if len(res.toolChoices) != 1 || res.toolChoices[0].Arguments["reasoning"] != "short" { + t.Fatalf("expected the reasoning call, got %+v", res.toolChoices) + } +} + +// Counter-check: without a step cap the decision behaves as before, it sends +// no cap of its own and raises the cap once on a cut tool call. +func TestDecisionWithoutCapKeepsTheLengthRetry(t *testing.T) { + llm := &lengthCutLLM{} + _, err := decision(context.Background(), llm, outputCapConv, Tools{reasoningTool()}, "reasoning", 3) + if err == nil || errors.Is(err, errOutputCapReached) { + t.Fatalf("expected the uncapped truncation error, got %v", err) + } + if len(llm.requests) != 2 { + t.Fatalf("expected the original attempt plus one length retry, got %d", len(llm.requests)) + } + if llm.requests[0].MaxTokens != 0 || llm.requests[1].MaxTokens != 32768 { + t.Fatalf("expected caps 0 then 32768, got %d then %d", llm.requests[0].MaxTokens, llm.requests[1].MaxTokens) + } +} diff --git a/options.go b/options.go index 1c1c50a..ffdd2a7 100644 --- a/options.go +++ b/options.go @@ -48,6 +48,7 @@ type Options struct { loopDetectionSteps int forceReasoning bool forceReasoningTool bool + reasoningMaxTokens int guidedTools bool parallelToolExecution bool toolImageForwarding bool @@ -409,6 +410,34 @@ func WithForceReasoningTool() func(o *Options) { } } +// DefaultReasoningMaxTokens is the output cap of the forced reasoning steps +// (tool and parameter reasoning) when WithReasoningMaxTokens is not set. +const DefaultReasoningMaxTokens = 2048 + +// WithReasoningMaxTokens sets the output-token cap of the forced reasoning +// steps (see WithForceReasoning). These steps are optional: when the cap cuts +// one, the loop continues without its reasoning instead of letting a model +// that never closes the reasoning argument run to the client's global cap. +// 0 (or unset) uses DefaultReasoningMaxTokens; a negative value removes the +// step's own cap, leaving the request to the client's default as before. +func WithReasoningMaxTokens(n int) func(o *Options) { + return func(o *Options) { + o.reasoningMaxTokens = n + } +} + +// reasoningOutputCap returns the cap the reasoning steps send, 0 for none. +func (o *Options) reasoningOutputCap() int { + switch { + case o.reasoningMaxTokens < 0: + return 0 + case o.reasoningMaxTokens == 0: + return DefaultReasoningMaxTokens + default: + return o.reasoningMaxTokens + } +} + // WithToolImageForwarding, when true, forwards image blocks from an MCP tool // result to the model as a follow-up user message (OpenAI tool-role messages // cannot carry image parts). Default false. diff --git a/plan.go b/plan.go index 8cbe7b4..ae16f6c 100644 --- a/plan.go +++ b/plan.go @@ -708,6 +708,9 @@ func convertOptionsToFunctions(o *Options) []Option { if o.forceReasoning { opts = append(opts, WithForceReasoning()) } + if o.reasoningMaxTokens != 0 { + opts = append(opts, WithReasoningMaxTokens(o.reasoningMaxTokens)) + } if o.maxRetries > 0 { opts = append(opts, WithMaxRetries(o.maxRetries)) } diff --git a/reasoning_cap_test.go b/reasoning_cap_test.go new file mode 100644 index 0000000..daa7773 --- /dev/null +++ b/reasoning_cap_test.go @@ -0,0 +1,251 @@ +package cogito_test + +import ( + "context" + "encoding/json" + "fmt" + "strings" + "sync" + + . "github.com/mudler/cogito" + "github.com/mudler/cogito/tests/mock" + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + "github.com/sashabaranov/go-openai" +) + +// clientDefaultCap stands in for the client's global output cap (the bundled +// clients send 16384 when the request carries none). +const clientDefaultCap = 16384 + +// reasoningStreamLLM is a streaming fake that answers by the forced tool of a +// request. Its reasoning tool call can be "endless": the argument never +// closes and the stream only ends when the request's output cap (or the +// client default) is spent, one token per delta, with finish_reason=length, +// which is what a model that loops inside the reasoning argument does. +type reasoningStreamLLM struct { + mu sync.Mutex + + endlessToolReasoning bool // the reasoning before the tool pick never ends + endlessParamReasoning bool // the reasoning before parameter generation never ends + shortReasoning string // reasoning text when it does end + + picks int + + reasoningRequests []openai.ChatCompletionRequest + reasoningDeltas []int + intentionRequests []openai.ChatCompletionRequest + paramRequests []openai.ChatCompletionRequest +} + +func forcedName(req openai.ChatCompletionRequest) string { + switch tc := req.ToolChoice.(type) { + case openai.ToolChoice: + return tc.Function.Name + case *openai.ToolChoice: + return tc.Function.Name + } + return "" +} + +func isToolPickReasoning(req openai.ChatCompletionRequest) bool { + for _, m := range req.Messages { + if strings.Contains(m.Content, "Analyze the current situation and available tools") { + return true + } + } + return false +} + +func (m *reasoningStreamLLM) Ask(ctx context.Context, f Fragment) (Fragment, error) { + return f.AddMessage(AssistantMessageRole, "final answer"), nil +} + +func (m *reasoningStreamLLM) CreateChatCompletion(ctx context.Context, req openai.ChatCompletionRequest) (LLMReply, LLMUsage, error) { + return LLMReply{}, LLMUsage{}, fmt.Errorf("unexpected non-streaming call (forced tool %q)", forcedName(req)) +} + +func (m *reasoningStreamLLM) CreateChatCompletionStream(ctx context.Context, req openai.ChatCompletionRequest) (<-chan StreamEvent, error) { + m.mu.Lock() + defer m.mu.Unlock() + + var events []StreamEvent + call := func(name, args string) { + events = append(events, + StreamEvent{Type: StreamEventToolCall, ToolCallIndex: 0, ToolCallID: "c1", ToolName: name, ToolArgs: args}, + StreamEvent{Type: StreamEventDone, FinishReason: "tool_calls", Usage: LLMUsage{CompletionTokens: 20}}) + } + + switch name := forcedName(req); name { + case "reasoning": + m.reasoningRequests = append(m.reasoningRequests, req) + endless := m.endlessParamReasoning + if isToolPickReasoning(req) { + endless = m.endlessToolReasoning + } + if !endless { + args, _ := json.Marshal(map[string]string{"reasoning": m.shortReasoning}) + m.reasoningDeltas = append(m.reasoningDeltas, 1) + call("reasoning", string(args)) + break + } + limit := req.MaxTokens + if limit == 0 { + limit = clientDefaultCap + } + events = append(events, StreamEvent{Type: StreamEventToolCall, ToolCallIndex: 0, ToolCallID: "c1", ToolName: "reasoning", ToolArgs: `{"reasoning":"`}) + for i := 1; i < limit; i++ { + events = append(events, StreamEvent{Type: StreamEventToolCall, ToolCallIndex: 0, ToolArgs: "a "}) + } + events = append(events, StreamEvent{Type: StreamEventDone, FinishReason: "length", MaxTokens: limit, + Usage: LLMUsage{PromptTokens: 100, CompletionTokens: limit, TotalTokens: 100 + limit}}) + m.reasoningDeltas = append(m.reasoningDeltas, limit) + case "pick_tool", "pick_tools", "": + if name == "" && len(req.Tools) == 0 { + // Plain answer turn. + events = append(events, + StreamEvent{Type: StreamEventContent, Content: "final answer"}, + StreamEvent{Type: StreamEventDone, FinishReason: "stop"}) + break + } + m.intentionRequests = append(m.intentionRequests, req) + m.picks++ + if m.picks == 1 { + call("pick_tool", `{"tool":"search","reasoning":"search first"}`) + } else { + call("pick_tool", `{"tool":"reply","reasoning":"done"}`) + } + case "search": + m.paramRequests = append(m.paramRequests, req) + call("search", `{}`) + default: + return nil, fmt.Errorf("unexpected forced tool %q", name) + } + + ch := make(chan StreamEvent, len(events)) + for _, ev := range events { + ch <- ev + } + close(ch) + return ch, nil +} + +func runWithReasoning(llm *reasoningStreamLLM, opts ...Option) (Fragment, error) { + f := NewEmptyFragment().AddMessage(UserMessageRole, "Wie buche ich eine Eingangsrechnung?") + search := mock.NewMockTool("search", "Search the knowledge base") + mock.SetRunResult(search, "search result") + base := []Option{ + WithTools(search), + WithForceReasoning(), + WithIterations(2), + WithMaxRetries(2), + WithStreamCallback(func(StreamEvent) {}), + } + return ExecuteTools(llm, f, append(base, opts...)...) +} + +func lastAssistant(f Fragment) string { + for i := len(f.Messages) - 1; i >= 0; i-- { + if f.Messages[i].Role == AssistantMessageRole.String() && f.Messages[i].Content != "" { + return f.Messages[i].Content + } + } + return "" +} + +var _ = Describe("Reasoning step output cap", func() { + It("caps an endless tool-pick reasoning and still reaches the tool and a final answer", func() { + llm := &reasoningStreamLLM{endlessToolReasoning: true} + + result, err := runWithReasoning(llm) + Expect(err).ToNot(HaveOccurred()) + + Expect(llm.reasoningRequests).ToNot(BeEmpty()) + for _, req := range llm.reasoningRequests { + Expect(req.MaxTokens).To(Equal(DefaultReasoningMaxTokens), "every reasoning request carries the step's own cap") + } + for _, n := range llm.reasoningDeltas { + Expect(n).To(BeNumerically("<=", DefaultReasoningMaxTokens), "the reasoning stream stops at the step cap, not at the client cap") + } + // A cut step is final: no length retry with a raised cap per pick. + Expect(llm.reasoningRequests).To(HaveLen(llm.picks)) + + // The loop continued: the tool was picked, parametrized and run. + Expect(llm.paramRequests).To(HaveLen(1)) + Expect(result.Status.ToolsCalled.Names()).To(ContainElement("search")) + Expect(lastAssistant(result)).ToNot(BeEmpty()) + }) + + It("caps an endless parameter reasoning and falls back to the original reasoning", func() { + llm := &reasoningStreamLLM{endlessParamReasoning: true, shortReasoning: "use search for the booking rule"} + + result, err := runWithReasoning(llm) + Expect(err).ToNot(HaveOccurred()) + + var capped int + for i, req := range llm.reasoningRequests { + Expect(req.MaxTokens).To(Equal(DefaultReasoningMaxTokens)) + if llm.reasoningDeltas[i] > 1 { + capped++ + Expect(llm.reasoningDeltas[i]).To(Equal(DefaultReasoningMaxTokens)) + } + } + Expect(capped).To(Equal(1), "exactly the one parameter reasoning ran into the cap, once") + + Expect(llm.paramRequests).To(HaveLen(1)) + Expect(llm.paramRequests[0].Messages[0].Content).To(ContainSubstring("use search for the booking rule"), + "the fallback keeps the tool-pick reasoning for parameter generation") + Expect(result.Status.ToolsCalled.Names()).To(ContainElement("search")) + }) + + It("uses a short reasoning that ends normally (counter-check)", func() { + llm := &reasoningStreamLLM{shortReasoning: "the user asks about booking, search the knowledge base"} + + result, err := runWithReasoning(llm) + Expect(err).ToNot(HaveOccurred()) + + // The reasoning reaches the tool pick as an assistant turn ... + Expect(llm.intentionRequests).ToNot(BeEmpty()) + var seen bool + for _, m := range llm.intentionRequests[0].Messages { + if strings.Contains(m.Content, "the user asks about booking, search the knowledge base") { + seen = true + } + } + Expect(seen).To(BeTrue(), "the reasoning is passed on to the tool pick") + // ... and the parameter generation gets the combined reasoning. + Expect(llm.paramRequests).To(HaveLen(1)) + Expect(llm.paramRequests[0].Messages[0].Content).To(ContainSubstring("Parameter Analysis")) + Expect(result.Status.ToolsCalled.Names()).To(ContainElement("search")) + }) + + It("sends a configured cap, the default for 0, and no step cap when negative", func() { + for _, tc := range []struct { + opt Option + want int + }{ + {WithReasoningMaxTokens(512), 512}, + {WithReasoningMaxTokens(0), DefaultReasoningMaxTokens}, + {nil, DefaultReasoningMaxTokens}, + {WithReasoningMaxTokens(-1), 0}, + } { + llm := &reasoningStreamLLM{endlessToolReasoning: true} + var opts []Option + if tc.opt != nil { + opts = append(opts, tc.opt) + } + _, err := runWithReasoning(llm, opts...) + Expect(llm.reasoningRequests).ToNot(BeEmpty()) + Expect(llm.reasoningRequests[0].MaxTokens).To(Equal(tc.want)) + if tc.want > 0 { + Expect(err).ToNot(HaveOccurred()) + continue + } + // Without a step cap the old behaviour stays: the endless + // reasoning runs to the client cap, the length retry doubles it, + // and the turn fails. This is the measured production failure. + Expect(err).To(HaveOccurred()) + Expect(llm.reasoningDeltas[0]).To(Equal(clientDefaultCap)) + } + }) +}) diff --git a/tools.go b/tools.go index fd1bb7d..32d5e80 100644 --- a/tools.go +++ b/tools.go @@ -282,16 +282,41 @@ func mergeConsecutiveAssistantMessages(messages []openai.ChatCompletionMessage) // Falls back to decision() when streaming is not possible. func decisionWithStreaming(ctx context.Context, llm LLM, conversation []openai.ChatCompletionMessage, tools Tools, forceTool string, maxRetries int, streamCB StreamCallback) (*decisionResult, error) { + return decisionWithStreamingCapped(ctx, llm, conversation, tools, forceTool, maxRetries, streamCB, 0) +} + +// errOutputCapReached reports that a decision with its own output cap (see +// decisionWithStreamingCapped) was cut by that cap. Callers treat the step as +// optional and continue without its result. +var errOutputCapReached = errors.New("output cap of an optional step reached") + +// outputCapReachedError wraps errOutputCapReached with the figures of the cut. +func outputCapReachedError(what string, outCap int, usage LLMUsage) error { + return fmt.Errorf("%s: %w (max tokens %d, completion tokens %d)", what, errOutputCapReached, outCap, usage.CompletionTokens) +} + +// decisionWithStreamingCapped is decisionWithStreaming with a request-level +// output cap. With outCap > 0 every attempt carries max_tokens=outCap, which +// the bundled clients honour over their own default, and a +// finish_reason=length is final: no length retry with a raised cap, the call +// returns an error wrapping errOutputCapReached instead. Meant for optional +// steps (the reasoning tool) that must not run to the client's global cap. +// With outCap <= 0 it behaves exactly like decisionWithStreaming. +func decisionWithStreamingCapped(ctx context.Context, llm LLM, conversation []openai.ChatCompletionMessage, + tools Tools, forceTool string, maxRetries int, streamCB StreamCallback, outCap int) (*decisionResult, error) { sllm, isStreaming := llm.(StreamingLLM) if !isStreaming || streamCB == nil { - return decision(ctx, llm, conversation, tools, forceTool, maxRetries) + return decisionCapped(ctx, llm, conversation, tools, forceTool, maxRetries, outCap) } req := openai.ChatCompletionRequest{ Messages: mergeConsecutiveAssistantMessages(normalizeSystemMessages(conversation)), Tools: tools.ToOpenAI(), } + if outCap > 0 { + req.MaxTokens = outCap + } if forceTool != "" { req.ToolChoice = openai.ToolChoice{ @@ -389,6 +414,16 @@ func decisionWithStreaming(ctx context.Context, llm LLM, conversation []openai.C xlog.Debug("[decisionWithStreaming] processed", "message", content, "reasoning", reasoning) + if outCap > 0 && finishReason == "length" { + // The step's own cap cut it. A complete tool call is still + // usable; anything else (no call, a call with cut arguments) ends + // the optional step here instead of retrying with a raised cap. + if choices, bad := parseToolCalls(toolCalls, finishReason); len(toolCalls) > 0 && bad == nil { + return &decisionResult{toolChoices: choices, message: content, reasoning: reasoning, usage: usage}, nil + } + return nil, outputCapReachedError("streaming decision", outCap, usage) + } + if len(toolCalls) == 0 { if content == "" { // The model produced no visible content and selected no tool. @@ -469,11 +504,21 @@ func backoffOrCancel(ctx context.Context, attempt int) error { // Similar to agent.go's decision function but adapted for cogito's architecture func decision(ctx context.Context, llm LLM, conversation []openai.ChatCompletionMessage, tools Tools, forceTool string, maxRetries int) (*decisionResult, error) { + return decisionCapped(ctx, llm, conversation, tools, forceTool, maxRetries, 0) +} + +// decisionCapped is decision with a request-level output cap; see +// decisionWithStreamingCapped for the semantics of outCap. +func decisionCapped(ctx context.Context, llm LLM, conversation []openai.ChatCompletionMessage, + tools Tools, forceTool string, maxRetries int, outCap int) (*decisionResult, error) { decision := openai.ChatCompletionRequest{ Messages: mergeConsecutiveAssistantMessages(normalizeSystemMessages(conversation)), Tools: tools.ToOpenAI(), } + if outCap > 0 { + decision.MaxTokens = outCap + } if forceTool != "" { decision.ToolChoice = openai.ToolChoice{ @@ -519,6 +564,14 @@ func decision(ctx context.Context, llm LLM, conversation []openai.ChatCompletion //reasoning := resp.Choices[0].Reasoning xlog.Debug("[decision] processed", "message", msg.Content, "reasoning", reasoning) + if outCap > 0 && resp.ChatCompletionResponse.Choices[0].FinishReason == openai.FinishReasonLength { + // See decisionWithStreamingCapped. + if choices, bad := parseToolCalls(msg.ToolCalls, string(openai.FinishReasonLength)); len(msg.ToolCalls) > 0 && bad == nil { + return &decisionResult{toolChoices: choices, message: msg.Content, reasoning: reasoning, usage: usage}, nil + } + return nil, outputCapReachedError("decision", outCap, usage) + } + if len(msg.ToolCalls) == 0 { // An empty reply cut by length is the same truncation the streaming // path handles; returning it as an empty text answer would end the @@ -610,13 +663,15 @@ func generateToolParameters(o *Options, llm LLM, tool ToolDefinitionInterface, c return nil, err } - // Use decision with reasoning tool to force structured output - paramReasoningResult, err := decisionWithStreaming(o.context, llm, + // Use decision with reasoning tool to force structured output. + // Bounded by the reasoning cap: on truncation the error below takes + // the same fallback as any other failure. + paramReasoningResult, err := decisionWithStreamingCapped(o.context, llm, append(conversation, openai.ChatCompletionMessage{ Role: "system", Content: paramPrompt, }), - Tools{reasoningTool()}, "reasoning", o.maxRetries, o.streamCallback) + Tools{reasoningTool()}, "reasoning", o.maxRetries, o.streamCallback, o.reasoningOutputCap()) if err != nil { xlog.Warn("Failed to get parameter reasoning, using original reasoning", "error", err) // Fall back to original single-step approach @@ -723,14 +778,21 @@ func pickTool(ctx context.Context, llm LLM, fragment Fragment, tools Tools, opts } } - reasoningResult, err := decisionWithStreaming(ctx, llm, + // The reasoning step is optional: bounded by its own output cap, and when + // that cap cuts it the tool pick continues without the reasoning instead + // of letting a model that never closes the argument run to the global cap. + reasoningResult, err := decisionWithStreamingCapped(ctx, llm, append(messages, openai.ChatCompletionMessage{ Role: "user", Content: reasoningPrompt, }), - Tools{reasoningTool()}, "reasoning", o.maxRetries, o.streamCallback) + Tools{reasoningTool()}, "reasoning", o.maxRetries, o.streamCallback, o.reasoningOutputCap()) if err != nil { - return nil, fmt.Errorf("failed to get reasoning: %w", err) + if !errors.Is(err, errOutputCapReached) { + return nil, fmt.Errorf("failed to get reasoning: %w", err) + } + xlog.Warn("[pickTool] reasoning step cut by its output cap, picking the tool without it", "error", err) + reasoningResult = &decisionResult{} } // Extract reasoning from the tool call response From 657aca6dfe5de3db6e0f330abf8cd31762beb392 Mon Sep 17 00:00:00 2001 From: Stefan Walcz Date: Fri, 2 Oct 2026 18:55:00 +0200 Subject: [PATCH 2/3] fix(tools): bound the sink state's parameter step by its own output cap With WithForceReasoning, a picked sink state that has parameters (an embedder's own sink such as LocalAGI's no_tool_to_call) gets its arguments from a forced parameter-generation request that carried no output cap, so it ran to the client's global cap (16384 by default). Observed with the daily-driver model behind LocalAI: the model answered inside the sink's "reasoning" argument and then repeated "17 mal 23 ist 391." until the cap (argumentsLen=46865). The backend reported finish_reason=tool_calls for the cut call, so the arguments correction re-ran the same request up to maxRetries times, and the turn hung for more than 420 s. With finish_reason=length the length retry would have doubled the cap instead. The sink's parameter step now carries max_tokens from WithSinkStateMaxTokens (default DefaultSinkStateMaxTokens = 2048; 0 uses the default, a negative value keeps the previous uncapped behaviour). Only the sink is capped: real tools keep the client's cap, since their arguments may legitimately be long. A capped step that spends its whole cap counts as cut even when the backend reports finish_reason=tool_calls: a complete call is still used, otherwise the step fails with errOutputCapReached and neither the length retry nor the arguments correction runs. This also applies to the reasoning steps capped by WithReasoningMaxTokens, which share the capped decision. On a cut, the parameter generation keeps its existing fallback (the arguments from the pick). The sink is not executed and the turn's answer comes from the closing completion after the loop, so nothing the user sees is lost and the turn ends. The option is forwarded to plan sub-executions. Signed-off-by: Stefan Walcz Assisted-by: Claude Opus 5.5 --- options.go | 31 ++++++ plan.go | 3 + sink_cap_test.go | 285 +++++++++++++++++++++++++++++++++++++++++++++++ tools.go | 22 +++- 4 files changed, 337 insertions(+), 4 deletions(-) create mode 100644 sink_cap_test.go diff --git a/options.go b/options.go index ffdd2a7..8129773 100644 --- a/options.go +++ b/options.go @@ -49,6 +49,7 @@ type Options struct { forceReasoning bool forceReasoningTool bool reasoningMaxTokens int + sinkStateMaxTokens int guidedTools bool parallelToolExecution bool toolImageForwarding bool @@ -438,6 +439,36 @@ func (o *Options) reasoningOutputCap() int { } } +// DefaultSinkStateMaxTokens is the output cap of the sink state's parameter +// generation when WithSinkStateMaxTokens is not set. +const DefaultSinkStateMaxTokens = 2048 + +// WithSinkStateMaxTokens sets the output-token cap of the step that generates +// the sink state's arguments (see WithSinkState, WithForceReasoning). The sink +// ends the tool loop and its arguments are not executed — the answer comes +// from the closing completion after the loop — so when the cap cuts them the +// loop ends without them instead of letting a model that repeats itself inside +// the argument run to the client's global cap. Real tools are not affected. +// 0 (or unset) uses DefaultSinkStateMaxTokens; a negative value removes the +// step's own cap, leaving the request to the client's default as before. +func WithSinkStateMaxTokens(n int) func(o *Options) { + return func(o *Options) { + o.sinkStateMaxTokens = n + } +} + +// sinkStateOutputCap returns the cap the sink's parameter step sends, 0 for none. +func (o *Options) sinkStateOutputCap() int { + switch { + case o.sinkStateMaxTokens < 0: + return 0 + case o.sinkStateMaxTokens == 0: + return DefaultSinkStateMaxTokens + default: + return o.sinkStateMaxTokens + } +} + // WithToolImageForwarding, when true, forwards image blocks from an MCP tool // result to the model as a follow-up user message (OpenAI tool-role messages // cannot carry image parts). Default false. diff --git a/plan.go b/plan.go index ae16f6c..f5b6ed0 100644 --- a/plan.go +++ b/plan.go @@ -711,6 +711,9 @@ func convertOptionsToFunctions(o *Options) []Option { if o.reasoningMaxTokens != 0 { opts = append(opts, WithReasoningMaxTokens(o.reasoningMaxTokens)) } + if o.sinkStateMaxTokens != 0 { + opts = append(opts, WithSinkStateMaxTokens(o.sinkStateMaxTokens)) + } if o.maxRetries > 0 { opts = append(opts, WithMaxRetries(o.maxRetries)) } diff --git a/sink_cap_test.go b/sink_cap_test.go new file mode 100644 index 0000000..a30a887 --- /dev/null +++ b/sink_cap_test.go @@ -0,0 +1,285 @@ +package cogito_test + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "strings" + "sync" + + . "github.com/mudler/cogito" + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + "github.com/sashabaranov/go-openai" +) + +// The loop an embedder's sink argument ran into in production: the model +// answered the question inside the sink's "reasoning" argument and then +// repeated the last sentence until the output cap. +const sinkLoopSentence = "17 mal 23 ist 391. " + +// sinkArgs mirrors the sink state an embedder registers (LocalAGI's +// no_tool_to_call): a tool with one free-text argument. +type sinkArgs struct { + Reasoning string `json:"reasoning"` +} + +type sinkRunner struct{} + +func (sinkRunner) Run(args sinkArgs) (string, any, error) { + return "No action needed: " + args.Reasoning, nil, nil +} + +// docArgs is a real tool whose argument is legitimately long (a document). +type docArgs struct { + Content string `json:"content"` +} + +type docRunner struct { + mu sync.Mutex + runs []string +} + +func (d *docRunner) Run(args docArgs) (string, any, error) { + d.mu.Lock() + defer d.mu.Unlock() + d.runs = append(d.runs, args.Content) + return "document written", nil, nil +} + +// sinkTurn is one scripted reply, rendered either as a stream or as a single +// completion, so both decision paths see the same model behaviour. +type sinkTurn struct { + content string + toolName string + toolArgs []string // argument deltas + finishReason string + maxTokens int + completion int +} + +// sinkLLM answers by the forced tool of a request. The parameter generation +// of the sink can be scripted to loop inside its argument until the output +// cap of the request (or the client default) is spent. +type sinkLLM struct { + mu sync.Mutex + + picks []string // tools the intention step picks, in order (last repeats) + + endlessSink bool // the sink argument never closes + maskedLength bool // the backend reports finish_reason=tool_calls instead of length + docDeltas int // length of the real tool's argument, in deltas + sinkText string // sink argument when it does end + + pickIndex int + sinkRequests []openai.ChatCompletionRequest + sinkDeltas []int + docRequests []openai.ChatCompletionRequest +} + +func (m *sinkLLM) Ask(ctx context.Context, f Fragment) (Fragment, error) { + return f.AddMessage(AssistantMessageRole, strings.TrimSpace(sinkLoopSentence)), nil +} + +func (m *sinkLLM) turn(req openai.ChatCompletionRequest) (sinkTurn, error) { + m.mu.Lock() + defer m.mu.Unlock() + + call := func(name, args string) sinkTurn { + return sinkTurn{toolName: name, toolArgs: []string{args}, finishReason: "tool_calls", completion: 20} + } + + switch name := forcedName(req); name { + case "reasoning": + return call("reasoning", `{"reasoning":"no tool is needed, answer directly"}`), nil + case "pick_tool": + pick := m.picks[min(m.pickIndex, len(m.picks)-1)] + m.pickIndex++ + return call("pick_tool", fmt.Sprintf(`{"tool":%q,"reasoning":"pick %s"}`, pick, pick)), nil + case "no_tool_to_call": + m.sinkRequests = append(m.sinkRequests, req) + if !m.endlessSink { + args, _ := json.Marshal(sinkArgs{Reasoning: m.sinkText}) + m.sinkDeltas = append(m.sinkDeltas, 1) + return call(name, string(args)), nil + } + limit := req.MaxTokens + if limit == 0 { + limit = clientDefaultCap + } + deltas := make([]string, 0, limit) + deltas = append(deltas, `{"reasoning":"`) + for i := 1; i < limit; i++ { + deltas = append(deltas, sinkLoopSentence) + } + finish := "length" + if m.maskedLength { + // What LocalAI reported in production: the call was cut by the + // cap, the finish reason still said tool_calls. + finish = "tool_calls" + } + m.sinkDeltas = append(m.sinkDeltas, limit) + return sinkTurn{toolName: name, toolArgs: deltas, finishReason: finish, maxTokens: limit, completion: limit}, nil + case "write_document": + m.docRequests = append(m.docRequests, req) + deltas := []string{`{"content":"`} + for i := 0; i < m.docDeltas; i++ { + deltas = append(deltas, "x ") + } + deltas = append(deltas, `"}`) + return sinkTurn{toolName: name, toolArgs: deltas, finishReason: "tool_calls", completion: len(deltas)}, nil + case "": + if len(req.Tools) == 0 { + return sinkTurn{content: strings.TrimSpace(sinkLoopSentence), finishReason: "stop"}, nil + } + } + return sinkTurn{}, fmt.Errorf("unexpected request (forced tool %q)", forcedName(req)) +} + +func (m *sinkLLM) CreateChatCompletion(ctx context.Context, req openai.ChatCompletionRequest) (LLMReply, LLMUsage, error) { + t, err := m.turn(req) + if err != nil { + return LLMReply{}, LLMUsage{}, err + } + msg := openai.ChatCompletionMessage{Role: openai.ChatMessageRoleAssistant, Content: t.content} + if t.toolName != "" { + msg.ToolCalls = []openai.ToolCall{{ID: "c1", Type: openai.ToolTypeFunction, + Function: openai.FunctionCall{Name: t.toolName, Arguments: strings.Join(t.toolArgs, "")}}} + } + usage := LLMUsage{PromptTokens: 100, CompletionTokens: t.completion, TotalTokens: 100 + t.completion} + return LLMReply{ + ChatCompletionResponse: openai.ChatCompletionResponse{Choices: []openai.ChatCompletionChoice{ + {Message: msg, FinishReason: openai.FinishReason(t.finishReason)}, + }}, + MaxTokens: t.maxTokens, + }, usage, nil +} + +func (m *sinkLLM) CreateChatCompletionStream(ctx context.Context, req openai.ChatCompletionRequest) (<-chan StreamEvent, error) { + t, err := m.turn(req) + if err != nil { + return nil, err + } + var events []StreamEvent + if t.content != "" { + events = append(events, StreamEvent{Type: StreamEventContent, Content: t.content}) + } + for i, d := range t.toolArgs { + ev := StreamEvent{Type: StreamEventToolCall, ToolCallIndex: 0, ToolArgs: d} + if i == 0 { + ev.ToolCallID, ev.ToolName = "c1", t.toolName + } + events = append(events, ev) + } + events = append(events, StreamEvent{Type: StreamEventDone, FinishReason: t.finishReason, MaxTokens: t.maxTokens, + Usage: LLMUsage{PromptTokens: 100, CompletionTokens: t.completion, TotalTokens: 100 + t.completion}}) + ch := make(chan StreamEvent, len(events)) + for _, ev := range events { + ch <- ev + } + close(ch) + return ch, nil +} + +func runWithSink(llm *sinkLLM, streaming bool, doc *docRunner, opts ...Option) (Fragment, error) { + f := NewEmptyFragment().AddMessage(UserMessageRole, "Was ist 17 mal 23?") + base := []Option{ + WithSinkState(NewToolDefinition(sinkRunner{}, sinkArgs{}, "no_tool_to_call", + "Called when no other tool is needed to respond to the user")), + WithForceReasoning(), + WithIterations(3), + WithMaxRetries(2), + } + if doc != nil { + base = append(base, WithTools(NewToolDefinition[docArgs](doc, docArgs{}, "write_document", "Write a document"))) + } + if streaming { + base = append(base, WithStreamCallback(func(StreamEvent) {})) + } + return ExecuteTools(llm, f, append(base, opts...)...) +} + +// The sink-only turn ends without a tool call recorded; cogito reports that +// as ErrNoToolSelected, which embedders (LocalAGI) treat as a normal answer. +func expectAnswered(result Fragment, err error) { + if err != nil { + Expect(errors.Is(err, ErrNoToolSelected)).To(BeTrue(), "unexpected error: %v", err) + } + Expect(lastAssistant(result)).To(Equal(strings.TrimSpace(sinkLoopSentence))) +} + +var _ = Describe("Sink state output cap", func() { + for _, streaming := range []bool{true, false} { + for _, masked := range []bool{false, true} { + streaming, masked := streaming, masked + It(fmt.Sprintf("caps a sink argument that never closes, once, and still answers (streaming=%v, masked length=%v)", streaming, masked), func() { + llm := &sinkLLM{picks: []string{"no_tool_to_call"}, endlessSink: true, maskedLength: masked} + + result, err := runWithSink(llm, streaming, nil) + expectAnswered(result, err) + + Expect(llm.sinkRequests).To(HaveLen(1), "a capped sink step is final: no length retry, no correction retry") + Expect(llm.sinkRequests[0].MaxTokens).To(Equal(DefaultSinkStateMaxTokens), "the sink step carries its own cap") + Expect(llm.sinkDeltas[0]).To(Equal(DefaultSinkStateMaxTokens), "the loop stops at the sink cap, not the client cap") + }) + } + } + + It("uses a sink argument that ends normally as before (counter-check)", func() { + llm := &sinkLLM{picks: []string{"no_tool_to_call"}, sinkText: "17 mal 23 ist 391."} + + result, err := runWithSink(llm, true, nil) + expectAnswered(result, err) + + Expect(llm.sinkRequests).To(HaveLen(1)) + Expect(llm.sinkRequests[0].MaxTokens).To(Equal(DefaultSinkStateMaxTokens)) + }) + + It("does not cap a real tool with long arguments (counter-check)", func() { + doc := &docRunner{} + // Longer than the sink cap, so a cap leaking onto real tools would cut it. + llm := &sinkLLM{picks: []string{"write_document", "no_tool_to_call"}, sinkText: "done", docDeltas: 3 * DefaultSinkStateMaxTokens} + + result, err := runWithSink(llm, true, doc) + Expect(err).ToNot(HaveOccurred()) + + Expect(llm.docRequests).To(HaveLen(1)) + Expect(llm.docRequests[0].MaxTokens).To(BeZero(), "a real tool's parameter generation carries no sink cap") + Expect(doc.runs).To(HaveLen(1)) + Expect(doc.runs[0]).To(HaveLen(len("x ") * 3 * DefaultSinkStateMaxTokens)) + Expect(result.Status.ToolsCalled.Names()).To(ContainElement("write_document")) + + Expect(llm.sinkRequests).To(HaveLen(1)) + Expect(llm.sinkRequests[0].MaxTokens).To(Equal(DefaultSinkStateMaxTokens)) + }) + + It("sends a configured cap, the default for 0, and no step cap when negative", func() { + for _, tc := range []struct { + opt Option + want int + }{ + {WithSinkStateMaxTokens(512), 512}, + {WithSinkStateMaxTokens(0), DefaultSinkStateMaxTokens}, + {nil, DefaultSinkStateMaxTokens}, + {WithSinkStateMaxTokens(-1), 0}, + } { + llm := &sinkLLM{picks: []string{"no_tool_to_call"}, endlessSink: true} + var opts []Option + if tc.opt != nil { + opts = append(opts, tc.opt) + } + result, err := runWithSink(llm, true, nil, opts...) + expectAnswered(result, err) + Expect(llm.sinkRequests[0].MaxTokens).To(Equal(tc.want)) + if tc.want > 0 { + Expect(llm.sinkRequests).To(HaveLen(1)) + continue + } + // Without a step cap the old behaviour stays: the loop runs to the + // client cap and the length retry doubles it. + Expect(llm.sinkDeltas[0]).To(Equal(clientDefaultCap)) + Expect(llm.sinkRequests).To(HaveLen(2)) + } + }) +}) diff --git a/tools.go b/tools.go index 32d5e80..d92821c 100644 --- a/tools.go +++ b/tools.go @@ -414,10 +414,14 @@ func decisionWithStreamingCapped(ctx context.Context, llm LLM, conversation []op xlog.Debug("[decisionWithStreaming] processed", "message", content, "reasoning", reasoning) - if outCap > 0 && finishReason == "length" { + if outCap > 0 && (finishReason == "length" || usage.CompletionTokens >= outCap) { // The step's own cap cut it. A complete tool call is still // usable; anything else (no call, a call with cut arguments) ends // the optional step here instead of retrying with a raised cap. + // Spending the whole cap counts as cut even when the backend + // reports finish_reason=tool_calls for a call it truncated + // (observed with LocalAI); otherwise the arguments correction + // would re-run the step maxRetries times. if choices, bad := parseToolCalls(toolCalls, finishReason); len(toolCalls) > 0 && bad == nil { return &decisionResult{toolChoices: choices, message: content, reasoning: reasoning, usage: usage}, nil } @@ -564,7 +568,7 @@ func decisionCapped(ctx context.Context, llm LLM, conversation []openai.ChatComp //reasoning := resp.Choices[0].Reasoning xlog.Debug("[decision] processed", "message", msg.Content, "reasoning", reasoning) - if outCap > 0 && resp.ChatCompletionResponse.Choices[0].FinishReason == openai.FinishReasonLength { + if outCap > 0 && (resp.ChatCompletionResponse.Choices[0].FinishReason == openai.FinishReasonLength || usage.CompletionTokens >= outCap) { // See decisionWithStreamingCapped. if choices, bad := parseToolCalls(msg.ToolCalls, string(openai.FinishReasonLength)); len(msg.ToolCalls) > 0 && bad == nil { return &decisionResult{toolChoices: choices, message: msg.Content, reasoning: reasoning, usage: usage}, nil @@ -706,8 +710,18 @@ func generateToolParameters(o *Options, llm LLM, tool ToolDefinitionInterface, c } } - // Use decision to force parameter generation - result, err := decisionWithStreaming(o.context, llm, conv, Tools{tool}, toolFunc.Name, o.maxRetries, o.streamCallback) + // Use decision to force parameter generation. Only the sink state's + // arguments are bounded by a cap of their own: the sink ends the loop and + // its arguments are not executed, so a cut leaves the caller's fallback + // (the arguments from the pick) without losing the answer, which the + // closing completion produces. Real tools keep the client's cap — their + // arguments may legitimately be long (a document to write). + paramCap := 0 + if o.sinkState && o.sinkStateTool != nil && o.sinkStateTool.Tool().Function != nil && + toolFunc.Name == o.sinkStateTool.Tool().Function.Name { + paramCap = o.sinkStateOutputCap() + } + result, err := decisionWithStreamingCapped(o.context, llm, conv, Tools{tool}, toolFunc.Name, o.maxRetries, o.streamCallback, paramCap) if err != nil { return nil, fmt.Errorf("failed to generate parameters for tool %s: %w", toolFunc.Name, err) } From bec8145f49239164f572fb190e9e5807d4c0b0ed Mon Sep 17 00:00:00 2001 From: Stefan Walcz Date: Fri, 2 Oct 2026 19:19:34 +0200 Subject: [PATCH 3/3] docs(readme): document the sink state's output cap Signed-off-by: Stefan Walcz Assisted-by: Claude Opus 5.5 --- README.md | 1 + 1 file changed, 1 insertion(+) diff --git a/README.md b/README.md index 98a10a0..9a8874a 100644 --- a/README.md +++ b/README.md @@ -629,6 +629,7 @@ result, err := cogito.ExecuteTools(llm, fragment, - When enabled, the sink state tool appears as an option in the tool selection enum - The sink state tool receives a `reasoning` parameter containing the LLM's reasoning about why no tool is needed - Custom sink state tools must accept a `reasoning` parameter in their arguments +- The sink state's arguments are generated under their own output cap, `DefaultSinkStateMaxTokens` (2048) by default. A model that never closes the `reasoning` argument is cut at the cap instead of running to the client's global limit, and a capped step is neither length-retried nor sent through argument correction; the loop then answers as usual. Set it with `cogito.WithSinkStateMaxTokens(n)` (`0` = default, a negative value restores the previous uncapped behaviour). Real tools are not affected by this cap. #### Sub-Agent Spawning