From f23ee82e22a50ec4926fc748b11a02d73ea941d0 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Fri, 25 Sep 2026 20:29:45 +0000 Subject: [PATCH 01/16] fix(chat): clamp max_tokens per request to the room the window has left The client sent its full output cap as max_tokens on every request. vLLM rejects prompt + max_tokens > window, so with a cap that is a large part of the window every long prompt, compaction summaries included, overflowed. The turn's LLM wrapper now sets max_tokens to min(cap, window - prompt estimate - 512), floored at 1024, when the request carries no explicit value. The session records the cap the factory and discovery resolved for the current model. Co-Authored-By: Claude Opus 5.5 --- chat/contextlive.go | 75 +++++++++++++++-- chat/modellimits.go | 34 ++++++++ chat/output_clamp_test.go | 172 ++++++++++++++++++++++++++++++++++++++ chat/session.go | 15 +++- 4 files changed, 290 insertions(+), 6 deletions(-) create mode 100644 chat/output_clamp_test.go diff --git a/chat/contextlive.go b/chat/contextlive.go index 749bc9d..cc4ba76 100644 --- a/chat/contextlive.go +++ b/chat/contextlive.go @@ -2,9 +2,11 @@ package chat import ( "context" + "encoding/json" "sync/atomic" "github.com/mudler/cogito" + "github.com/mudler/xlog" "github.com/sashabaranov/go-openai" ) @@ -76,12 +78,27 @@ func (l *liveUsage) promptTokens() int { // trackedLLM is the LLM the turn loop actually calls: the session's client, // with every request's reported usage recorded on the way back. +// +// It also sizes each request's output reservation (see clampOutputTokens). +// limits reports the session's output cap and context window; nil leaves every +// request as the caller built it. type trackedLLM struct { cogito.LLM - live *liveUsage + live *liveUsage + limits func() (cap, window int) +} + +// clamp applies clampOutputTokens with the session's current limits. +func (t *trackedLLM) clamp(request openai.ChatCompletionRequest) openai.ChatCompletionRequest { + if t.limits == nil { + return request + } + cap, window := t.limits() + return clampOutputTokens(request, cap, window) } func (t *trackedLLM) CreateChatCompletion(ctx context.Context, request openai.ChatCompletionRequest) (cogito.LLMReply, cogito.LLMUsage, error) { + request = t.clamp(request) var usage cogito.LLMUsage reply, err := retryRequest(ctx, func() (cogito.LLMReply, error) { r, u, e := t.LLM.CreateChatCompletion(ctx, request) @@ -120,6 +137,7 @@ type trackedStreamingLLM struct { // The relay goroutine is bounded by the source channel: it ends when the // provider closes it, which cogito's own consumer loop already depends on. func (t *trackedStreamingLLM) CreateChatCompletionStream(ctx context.Context, request openai.ChatCompletionRequest) (<-chan cogito.StreamEvent, error) { + request = t.clamp(request) src, err := retryRequest(ctx, func() (<-chan cogito.StreamEvent, error) { return t.stream.CreateChatCompletionStream(ctx, request) }) @@ -139,11 +157,58 @@ func (t *trackedStreamingLLM) CreateChatCompletionStream(ctx context.Context, re return out, nil } -// trackUsage wraps llm so the turn's requests report their size to live, +// trackUsage wraps llm so the turn's requests report their size to live and +// reserve only the output room the window has left (limits may be nil), // preserving streaming support when the client has it. -func trackUsage(llm cogito.LLM, live *liveUsage) cogito.LLM { +func trackUsage(llm cogito.LLM, live *liveUsage, limits func() (cap, window int)) cogito.LLM { + t := trackedLLM{LLM: llm, live: live, limits: limits} if s, ok := llm.(cogito.StreamingLLM); ok { - return &trackedStreamingLLM{trackedLLM: trackedLLM{LLM: llm, live: live}, stream: s} + return &trackedStreamingLLM{trackedLLM: t, stream: s} + } + return &t +} + +const ( + // minOutputTokens is the smallest output reservation a clamped request + // carries. A prompt that already fills the window still gets this much, + // and the overflow path deals with the rejection. + minOutputTokens = 1024 + // outputSafetyMargin absorbs the error of the bytes/4 prompt estimate. + outputSafetyMargin = 512 +) + +// clampOutputTokens sizes a request's output reservation to what the window +// has left after its prompt. +// +// The client sends its output cap as max_tokens on every request. A backend +// such as vLLM rejects a request when prompt + max_tokens exceeds the window, +// whatever the model actually generates. With a cap that is a large fraction +// of the window (regolo: 96000 of 210000), every prompt above window − cap +// failed, compaction summaries included. cogito applies the client's cap only +// when the request carries none, so setting it here decides the reservation. +// +// The result is min(cap, window − prompt − outputSafetyMargin), floored at +// minOutputTokens (and never above cap). The prompt estimate is bytes/4 over +// the messages and the JSON of the tool schemas, which are real prompt tokens. +// A request with an explicit value, or unknown limits (cap or window 0), is +// returned unchanged. +func clampOutputTokens(request openai.ChatCompletionRequest, cap, window int) openai.ChatCompletionRequest { + if cap <= 0 || window <= 0 || request.MaxTokens != 0 || request.MaxCompletionTokens != 0 { + return request + } + prompt := estimateTokens(request.Messages) + if len(request.Tools) > 0 { + if b, err := json.Marshal(request.Tools); err == nil { + prompt += len(b) / 4 + } + } + v := min(cap, window-prompt-outputSafetyMargin) + if v < minOutputTokens { + v = min(minOutputTokens, cap) + } + if v < cap { + xlog.Debug("output reservation clamped", "from", cap, "to", v, "prompt_estimate", prompt, "window", window) } - return &trackedLLM{LLM: llm, live: live} + request.MaxTokens = v + return request } diff --git a/chat/modellimits.go b/chat/modellimits.go index 9acffb3..76ea31a 100644 --- a/chat/modellimits.go +++ b/chat/modellimits.go @@ -3,8 +3,11 @@ package chat import ( "context" + "github.com/mudler/cogito" + "github.com/mudler/nib/auth" "github.com/mudler/nib/llmprovider" "github.com/mudler/nib/llmprovider/catalog" + "github.com/mudler/nib/types" ) // A model has two limits nib can only learn by asking the endpoint: the @@ -53,6 +56,11 @@ func (s *Session) ensureModelLimits(ctx context.Context) { cancel() if !res.Omit && res.MaxTokens > 0 { setter.SetMaxTokens(res.MaxTokens) + s.modelMu.Lock() + if s.llmModel == model { + s.outputCap = res.MaxTokens + } + s.modelMu.Unlock() } } @@ -67,3 +75,29 @@ func (s *Session) ensureModelLimits(ctx context.Context) { type maxTokensSetter interface { SetMaxTokens(int) } + +// factoryOutputCap is the output cap the LLM factory gave llm when it built it +// for provider: the same resolution (config, then catalog, no network), so the +// session knows what each request reserves without asking the client. A +// client that carries no cap (no SetMaxTokens) reports 0, which leaves its +// requests unclamped. +func factoryOutputCap(llm cogito.LLM, provider types.ModelProviderConfig, store *auth.Store) int { + if _, ok := llm.(maxTokensSetter); !ok { + return 0 + } + baseURL, apiKey, _ := llmprovider.ModelsEndpoint(provider, store) + res := catalog.ResolveMaxTokens(nil, provider, baseURL, apiKey) + if res.Omit || res.MaxTokens <= 0 { + return 0 + } + return res.MaxTokens +} + +// requestLimits reports the output cap and context window the turn's requests +// are clamped against (see clampOutputTokens), read under one lock so a model +// switch cannot pair one model's cap with another's window. +func (s *Session) requestLimits() (cap, window int) { + s.modelMu.RLock() + defer s.modelMu.RUnlock() + return s.outputCap, s.windowLocked() +} diff --git a/chat/output_clamp_test.go b/chat/output_clamp_test.go new file mode 100644 index 0000000..9c72660 --- /dev/null +++ b/chat/output_clamp_test.go @@ -0,0 +1,172 @@ +package chat + +import ( + "context" + "strings" + "testing" + + "github.com/mudler/cogito" + "github.com/mudler/nib/types" + "github.com/sashabaranov/go-openai" +) + +// promptOf builds a request whose estimated prompt is about tokens tokens +// (estimateTokens counts bytes/4). +func promptOf(tokens int) openai.ChatCompletionRequest { + return openai.ChatCompletionRequest{ + Messages: []openai.ChatCompletionMessage{{Role: "user", Content: strings.Repeat("a", tokens*4)}}, + } +} + +func TestClampOutputLeavesRoomForThePrompt(t *testing.T) { + got := clampOutputTokens(promptOf(70000), 96000, 96000) + want := 96000 - 70000 - outputSafetyMargin + if got.MaxTokens != want { + t.Fatalf("MaxTokens = %d, want %d", got.MaxTokens, want) + } +} + +func TestClampOutputNeverAboveCap(t *testing.T) { + got := clampOutputTokens(promptOf(100), 8000, 96000) + if got.MaxTokens != 8000 { + t.Fatalf("MaxTokens = %d, want the cap 8000", got.MaxTokens) + } +} + +func TestClampOutputFloorWhenPromptFillsWindow(t *testing.T) { + got := clampOutputTokens(promptOf(120000), 96000, 96000) + if got.MaxTokens != minOutputTokens { + t.Fatalf("MaxTokens = %d, want minOutputTokens %d", got.MaxTokens, minOutputTokens) + } +} + +func TestClampOutputLeavesExplicitValueAlone(t *testing.T) { + req := promptOf(70000) + req.MaxTokens = 50000 + if got := clampOutputTokens(req, 96000, 96000); got.MaxTokens != 50000 || got.MaxCompletionTokens != 0 { + t.Fatalf("explicit MaxTokens changed: %+v", got.MaxTokens) + } + req = promptOf(70000) + req.MaxCompletionTokens = 50000 + if got := clampOutputTokens(req, 96000, 96000); got.MaxTokens != 0 || got.MaxCompletionTokens != 50000 { + t.Fatalf("explicit MaxCompletionTokens changed: %d/%d", got.MaxTokens, got.MaxCompletionTokens) + } +} + +func TestClampOutputUnknownLimitsLeaveRequestAlone(t *testing.T) { + if got := clampOutputTokens(promptOf(70000), 0, 96000); got.MaxTokens != 0 { + t.Fatalf("cap 0: MaxTokens = %d, want 0", got.MaxTokens) + } + if got := clampOutputTokens(promptOf(70000), 96000, 0); got.MaxTokens != 0 { + t.Fatalf("window 0: MaxTokens = %d, want 0", got.MaxTokens) + } +} + +func TestClampOutputCountsToolSchemas(t *testing.T) { + req := promptOf(10000) + without := clampOutputTokens(req, 96000, 96000).MaxTokens + + req.Tools = []openai.Tool{{ + Type: openai.ToolTypeFunction, + Function: &openai.FunctionDefinition{ + Name: "big", + Description: strings.Repeat("d", 40000), // ~10k tokens of schema + }, + }} + with := clampOutputTokens(req, 96000, 96000).MaxTokens + if without-with < 9000 { + t.Fatalf("tool schemas not counted: without=%d with=%d", without, with) + } +} + +// captureLLMReq records the request it receives, on both paths. +type captureLLMReq struct { + got openai.ChatCompletionRequest + stream openai.ChatCompletionRequest +} + +func (c *captureLLMReq) Ask(ctx context.Context, f cogito.Fragment) (cogito.Fragment, error) { + return f, nil +} + +func (c *captureLLMReq) CreateChatCompletion(ctx context.Context, req openai.ChatCompletionRequest) (cogito.LLMReply, cogito.LLMUsage, error) { + c.got = req + return cogito.LLMReply{}, cogito.LLMUsage{}, nil +} + +func (c *captureLLMReq) CreateChatCompletionStream(ctx context.Context, req openai.ChatCompletionRequest) (<-chan cogito.StreamEvent, error) { + c.stream = req + ch := make(chan cogito.StreamEvent) + close(ch) + return ch, nil +} + +func TestClampOutputAppliedByTrackedLLM(t *testing.T) { + inner := &captureLLMReq{} + var live liveUsage + llm := trackUsage(inner, &live, func() (int, int) { return 96000, 96000 }) + + if _, _, err := llm.CreateChatCompletion(context.Background(), promptOf(70000)); err != nil { + t.Fatal(err) + } + want := 96000 - 70000 - outputSafetyMargin + if inner.got.MaxTokens != want { + t.Fatalf("blocking MaxTokens = %d, want %d", inner.got.MaxTokens, want) + } + + sl, ok := llm.(cogito.StreamingLLM) + if !ok { + t.Fatal("wrapper lost streaming support") + } + ch, err := sl.CreateChatCompletionStream(context.Background(), promptOf(70000)) + if err != nil { + t.Fatal(err) + } + for range ch { + } + if inner.stream.MaxTokens != want { + t.Fatalf("streaming MaxTokens = %d, want %d", inner.stream.MaxTokens, want) + } +} + +func TestClampOutputNilLimitsIsNoop(t *testing.T) { + inner := &captureLLMReq{} + var live liveUsage + llm := trackUsage(inner, &live, nil) + if _, _, err := llm.CreateChatCompletion(context.Background(), promptOf(70000)); err != nil { + t.Fatal(err) + } + if inner.got.MaxTokens != 0 { + t.Fatalf("MaxTokens = %d, want untouched 0", inner.got.MaxTokens) + } +} + +func TestClampOutputSessionLimits(t *testing.T) { + s := &Session{ + llmModel: "m", + outputCap: 96000, + compaction: types.CompactionConfig{MaxContextTokens: 96000}, + } + if c, w := s.requestLimits(); c != 96000 || w != 96000 { + t.Fatalf("requestLimits = %d,%d", c, w) + } + s.learnedWindow, s.learnedWindowModel = 210000, "m" + if _, w := s.requestLimits(); w != 210000 { + t.Fatalf("learned window not used: %d", w) + } +} + +// capSetterLLM is a client that carries an output cap, like LocalAIClient. +type capSetterLLM struct{ captureLLMReq } + +func (c *capSetterLLM) SetMaxTokens(int) {} + +func TestClampOutputFactoryCapRecorded(t *testing.T) { + p := types.ModelProviderConfig{Model: "some-model", MaxTokens: 96000, BaseURL: "http://localhost:1/v1"} + if got := factoryOutputCap(&capSetterLLM{}, p, nil); got != 96000 { + t.Fatalf("factoryOutputCap = %d, want 96000", got) + } + if got := factoryOutputCap(&captureLLMReq{}, p, nil); got != 0 { + t.Fatalf("adapter without a cap: factoryOutputCap = %d, want 0", got) + } +} diff --git a/chat/session.go b/chat/session.go index 88a0ea8..6749051 100644 --- a/chat/session.go +++ b/chat/session.go @@ -230,6 +230,13 @@ type Session struct { // modellimits.go. limitsFor string + // outputCap is the output-token cap the current model's client sends, + // as resolved for it (config, catalog, then discovery); guarded by + // modelMu. 0 means unknown, or a client that carries no cap. The turn's + // LLM wrapper clamps each request's reservation against it (see + // clampOutputTokens). + outputCap int + // prunedMu guards the tool-output pruning state below. The manipulator reads // it from inside cogito's loop, and nothing here should assume which // goroutine that is; Reload writes the policy from the turn goroutine. @@ -474,6 +481,8 @@ func NewSession(ctx context.Context, cfg types.Config, callbacks Callbacks, tran if err != nil { return nil, fmt.Errorf("create main LLM: %w", err) } + // Read before tracing wraps the client and hides its SetMaxTokens. + outCap := factoryOutputCap(llm, mainProvider, credStore) endpoints, configErrs := endpoint.New(cfg, credStore) classifier, err := provenance.ClassifierForConfig(cfg) if err != nil { @@ -542,6 +551,7 @@ func NewSession(ctx context.Context, cfg types.Config, callbacks Callbacks, tran agentLogs: newAgentLogStore(), llmModel: mainProvider.Model, mainProvider: mainProvider, + outputCap: outCap, configProvider: mainProvider, endpoints: endpoints, configErrs: configErrs, @@ -1649,7 +1659,7 @@ func (s *Session) SendMessage(text string, parts ...ContentPart) (string, error) // retrying client instead (WithAgentLLM below; see agentretry.go): // without it cogito would hand them this tracked one. agentLLM := retryForAgent(llm, &s.agentBackoff) - llm = trackUsage(llm, &s.live) + llm = trackUsage(llm, &s.live, s.requestLimits) // Build cogito options from config cogitoOpts := []cogito.Option{ @@ -2628,6 +2638,8 @@ func (s *Session) applyProvider(provider types.ModelProviderConfig, id string) e if err != nil { return err } + // Read before tracing wraps the client and hides its SetMaxTokens. + outCap := factoryOutputCap(llm, provider, s.credStore) if s.tracer != nil { llm = trace.NewRecordingLLM(llm, s.tracer, name, "") } @@ -2636,6 +2648,7 @@ func (s *Session) applyProvider(provider types.ModelProviderConfig, id string) e s.llm = llm s.llmModel = name s.mainProvider = provider + s.outputCap = outCap if id != "" { s.endpointID = id } From 58b5a1fafacab75ea5aad22255e14568f604c53b Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Fri, 25 Sep 2026 20:39:43 +0000 Subject: [PATCH 02/16] fix(chat): fit the compaction summary by moving the boundary, not truncating A summary that overflowed was retried once, at a size scaled by window/request. The output the request reserves does not scale with the prompt, so with a large reservation the retry overflowed too, and the session stayed stuck. The summary request now reserves min(output cap, summary_max_tokens, default 16384) through CreateChatCompletion, and its prompt is fitted into the window minus that reservation. On an overflow, summarizeFitting subtracts the stated overshoot from the prompt and moves the boundary earlier until the head fits, up to 6 attempts. Messages that leave the head stay in the history verbatim. End-of-turn, mid-turn and /compact all use this loop, and the artifact holds the head that was summarized. Co-Authored-By: Claude Opus 5.5 --- README.md | 10 ++ chat/compact.go | 224 +++++++++++++++++++-------- chat/compact_fit_test.go | 43 ++++-- chat/compact_internal_test.go | 25 ++- chat/midturn.go | 39 ++--- chat/overflow_retry_test.go | 31 +++- chat/summary_shift_test.go | 281 ++++++++++++++++++++++++++++++++++ chat/usage_test.go | 10 +- config/config.go | 4 + config/config_test.go | 2 +- config/reserve_config_test.go | 23 +++ selfdoc/README.md | 10 ++ types/config.go | 5 + 13 files changed, 586 insertions(+), 121 deletions(-) create mode 100644 chat/summary_shift_test.go diff --git a/README.md b/README.md index a7630d7..1505974 100644 --- a/README.md +++ b/README.md @@ -568,6 +568,16 @@ agents: # LOG_FORMAT=json switches to JSON lines. log_level: error +# Optional: conversation compaction. When the context fills up, nib +# summarizes the older turns and keeps the recent ones verbatim. +# summary_max_tokens caps the output the summary request reserves (default +# 16384; the model's own output cap wins when it is smaller). The summary +# prompt is fitted into the window minus that reservation. If the backend +# still rejects it as too large, nib summarizes fewer of the older turns and +# keeps the rest verbatim, instead of cutting the text it summarizes. +compaction: + summary_max_tokens: 16384 + # Optional: external MCP servers mcp_servers: filesystem: diff --git a/chat/compact.go b/chat/compact.go index e37872c..9bbf0d4 100644 --- a/chat/compact.go +++ b/chat/compact.go @@ -116,6 +116,15 @@ const ( // Threshold default already is). const defaultReserveTokens = 4096 +// defaultSummaryMaxTokens is the fallback for an unset CompactionConfig. +// SummaryMaxTokens, repeated from config.Load for the same reason. +const defaultSummaryMaxTokens = 16384 + +// maxSummaryAttempts bounds how many summary requests summarizeFitting sends +// for one compaction. Each overflow moves the boundary so the head shrinks by +// the backend's stated overshoot, so a fit normally takes one or two retries. +const maxSummaryAttempts = 6 + // rememberWindow records a context window a backend stated for a specific // model. Values outside the plausibility band are discarded rather than // stored, because an implausible figure is a parse artefact and storing one @@ -517,16 +526,20 @@ func fitSummaryInput(pieces []summaryPiece, maxTokens int) string { // summaryRetryTarget returns the prompt size to retry a summary at after the // backend rejected a prompt of sent (byte/4) tokens. // -// byte/4 is only an estimate, and a backend can count the same text as more -// tokens than that. The rejection states both figures in the backend's own -// count, so their ratio corrects the estimate. The 10% margin covers the -// estimate varying across the text. With no figures to go on, it halves. +// The rejection states the request size and the window in the backend's own +// count. Their difference is by how many tokens the request overshot, and that +// is a fixed number, not a ratio: the output the request reserves does not +// shrink with the prompt. Scaling the prompt by window/request left the +// reservation whole, so a large reservation overflowed again on every retry. +// The 10% margin covers the byte/4 estimate varying across the text. With no +// figures to go on, it halves. A result of zero or less means the reservation +// alone does not fit, and no prompt will. func summaryRetryTarget(err error, sent int) int { needs, allows, ok := overflowFigures(err) if !ok { return sent / 2 } - return int(float64(sent) * float64(allows) / float64(needs) * 0.9) + return int(float64(sent-(needs-allows)) * 0.9) } // overflowFigures reads the request size and the window from an overflow @@ -570,32 +583,23 @@ func (s *Session) compactHistory(ctx context.Context) (before, after int, err er before = estimateTokens(msgs) cfg := s.compactionConfig() - head, tail := splitForCompaction(msgs, cfg.KeepRecent) s.prunedMu.Lock() - head = stubbedView(head, s.prunedIDs) - s.prunedMu.Unlock() - pieces := renderMessages(head) - if len(pieces) == 0 { - return before, before, nil // nothing to compact + pruned := make(map[string]string, len(s.prunedIDs)) + for k, v := range s.prunedIDs { + pruned[k] = v } - // Save the full untruncated head as an artifact so nothing is lost - // when the lossy summary truncates or drops messages. The model can - // page through it via the read tool (artifact://N) or search it with - // search_artifacts. - var artifactURI string - if !cfg.DisableArtifactSpill && s.artifacts != nil { - var b strings.Builder - for _, p := range pieces { - b.WriteString(p.text) - } - if b.Len() > 0 { - artifactURI = s.artifacts.Save("compaction", b.String()) - } + s.prunedMu.Unlock() + view := func(m []openai.ChatCompletionMessage) []openai.ChatCompletionMessage { + return stubbedView(m, pruned) } - summary, err := s.summarize(ctx, pieces) + summary, head, tail, err := s.summarizeFitting(ctx, msgs, cfg.KeepRecent, view) if err != nil { return before, before, err } + if len(head) == 0 { + return before, before, nil // nothing to compact + } + artifactURI := s.spillCompactionHead(renderMessages(view(head))) // Build the new state up front; swap only after success (atomic). newFragMsgs := append([]openai.ChatCompletionMessage{summaryMessage(summary, artifactURI)}, tail...) @@ -627,58 +631,150 @@ func (s *Session) compactHistory(ctx context.Context) (before, after int, err er return before, after, nil } -// summarize asks the model for the compaction summary of pieces, fitting the -// prompt into the window. It counts the spend of every call it makes. -func (s *Session) summarize(ctx context.Context, pieces []summaryPiece) (string, error) { +// summaryPrefix is the instruction the rendered conversation follows in the +// summary prompt. +const summaryPrefix = compactInstruction + "\n\n--- CONVERSATION ---\n" + +// summaryOutputTokens is the output the summary request reserves: the smaller +// of the model's output cap and cfg.SummaryMaxTokens, and never more than half +// of a known window, so the prompt keeps room beside it. +// +// It is set on the request explicitly. Left to the client, the request +// reserves the model's whole output cap, and a backend that checks prompt + +// reservation against the window rejects every summary prompt larger than the +// window minus that cap, however short the summary would have been. +func (s *Session) summaryOutputTokens() int { cfg := s.compactionConfig() - // The summary prompt has to fit the same window as any other request, and - // the head it summarizes is by nature most of a window that just filled up. - const prefix = compactInstruction + "\n\n--- CONVERSATION ---\n" - limit := 0 - if window := s.contextWindow(); window > 0 { - limit = max(ContextBudget(cfg, window)-tokensOf(prefix), 1) - } - prompt := prefix + fitSummaryInput(pieces, limit) - llm, _ := s.currentLLM() - res, aerr := llm.Ask(ctx, cogito.NewFragment().AddMessage(cogito.UserMessageRole, prompt)) - // Exactly once: the backend's rejection states by how much the estimate was - // off, so a single corrected retry fits, and a second failure means the - // backend cannot take even the minimal prompt. - if aerr != nil && canRecoverFromOverflow(ctx, aerr) { - if res.Status != nil { - s.addUsage(res.Status.LastUsage) + out := cfg.SummaryMaxTokens + if out <= 0 { + out = defaultSummaryMaxTokens + } + cap, window := s.requestLimits() + if cap > 0 { + out = min(out, cap) + } + if window > 0 { + out = max(min(out, window/2), 1) + } + return out +} + +// summaryPromptLimit is the budget, in byte/4 tokens, for the rendered +// conversation in the first summary request: the window less the summary's +// output reservation (and never more than ContextBudget), less the +// instruction. 0 means no window is known, so no limit. +func (s *Session) summaryPromptLimit(output int) int { + window := s.contextWindow() + if window <= 0 { + return 0 + } + budget := min(ContextBudget(s.compactionConfig(), window), window-output) + return max(budget-tokensOf(summaryPrefix), 1) +} + +// summarizeFitting summarizes the head of msgs that splitForCompaction(msgs, +// keep) leaves, and returns the summary with the head it covers and the tail +// to keep verbatim. view maps messages to what the summary prompt should see +// (the pruned view of them). An empty head, with a nil error, means there was +// nothing to compact. +// +// When the backend rejects the summary as too large, the boundary moves +// earlier: keep grows until the head fits the retry target (see +// summaryRetryTarget), or until only the smallest head is left, which +// fitSummaryInput then truncates to the target. What leaves the head joins +// the tail, verbatim: nothing is lost, and the next compaction can summarize +// it. The loop stops on success, a non-overflow error, a cancelled context, a +// target of zero or less, or after maxSummaryAttempts, and returns the last +// error. It changes no session state other than the usage it counts. +func (s *Session) summarizeFitting(ctx context.Context, msgs []openai.ChatCompletionMessage, keep int, view func([]openai.ChatCompletionMessage) []openai.ChatCompletionMessage) (summary string, head, tail []openai.ChatCompletionMessage, err error) { + head, tail = splitForCompaction(msgs, keep) + if len(head) == 0 || len(renderMessages(view(head))) == 0 { + return "", nil, nil, nil + } + keep = len(tail) + prefix := tokensOf(summaryPrefix) + limit := s.summaryPromptLimit(s.summaryOutputTokens()) + for attempt := 0; ; attempt++ { + summary, sent, serr := s.summarize(ctx, renderMessages(view(head)), limit) + if serr == nil { + return summary, head, tail, nil + } + if attempt+1 >= maxSummaryAttempts || !canRecoverFromOverflow(ctx, serr) { + return "", nil, nil, serr + } + target := summaryRetryTarget(serr, sent) + if target <= 0 { + return "", nil, nil, serr } - limit = max(summaryRetryTarget(aerr, tokensOf(prompt))-tokensOf(prefix), 1) - prompt = prefix + fitSummaryInput(pieces, limit) - res, aerr = llm.Ask(ctx, cogito.NewFragment().AddMessage(cogito.UserMessageRole, prompt)) + limit = max(target-prefix, 1) + // Move the boundary earlier until the head fits the target. A head + // that renders nothing would leave nothing to summarize, so the + // smallest head kept is the smallest that still renders. + for estimateTokens(view(head))+prefix > target { + h, t := splitForCompaction(msgs, keep+1) + if len(h) == 0 || len(renderMessages(view(h))) == 0 { + break + } + keep++ + head, tail = h, t + } + } +} + +// spillCompactionHead saves the rendered head a summary stands for as an +// artifact, so nothing is lost to the lossy summary, and returns its URI ("" when +// spilling is off or there is no store). The model can page through it with +// the read tool (artifact://N) or search it with search_artifacts. +func (s *Session) spillCompactionHead(pieces []summaryPiece) string { + if s.compactionConfig().DisableArtifactSpill || s.artifacts == nil { + return "" } + var b strings.Builder + for _, p := range pieces { + b.WriteString(p.text) + } + if b.Len() == 0 { + return "" + } + return s.artifacts.Save("compaction", b.String()) +} + +// summarize sends one summary request for pieces, fitted into limit byte/4 +// tokens (0 means no limit) after the instruction, and returns the summary and +// the byte/4 size of the prompt it sent. It counts the spend of the call. +// +// The request goes out through CreateChatCompletion with MaxTokens set to +// summaryOutputTokens. Ask would leave the reservation to the client, which +// reserves the model's whole output cap. +func (s *Session) summarize(ctx context.Context, pieces []summaryPiece, limit int) (summary string, sentTokens int, err error) { + prompt := summaryPrefix + fitSummaryInput(pieces, limit) + sentTokens = tokensOf(prompt) + llm, _ := s.currentLLM() + reply, usage, aerr := llm.CreateChatCompletion(ctx, openai.ChatCompletionRequest{ + Messages: []openai.ChatCompletionMessage{{Role: cogito.UserMessageRole.String(), Content: prompt}}, + MaxTokens: s.summaryOutputTokens(), + }) // Compaction is not free, and it fires exactly when a session has already // grown expensive — so leaving it out would understate the runs that cost // the most. // - // LastUsage, not CumulativeUsage: Ask is a single call on a throwaway - // fragment that no ExecuteTools run ever stamped, so CumulativeUsage is - // zero here and reading it would count nothing at all. - // // Counted before BOTH exits below on purpose. What the backend served it // billed, whether the summary then came back empty or the call came back an // error, and a rejected summary must not also erase the spend — the same - // rule the interrupted-turn path in SendMessage follows. Placing this after - // either early return would make compaction the one path where paid-for - // tokens vanish, precisely on the sessions that spend the most. Clients - // that discard the fragment on error simply leave Status nil and add - // nothing, so the guard is what keeps this honest rather than optimistic. - if res.Status != nil { - s.addUsage(res.Status.LastUsage) - } + // rule the interrupted-turn path in SendMessage follows. A client that + // reports nothing on error simply adds zero. + s.addUsage(usage) if aerr != nil { - return "", fmt.Errorf("compaction summary failed: %w", aerr) + return "", sentTokens, fmt.Errorf("compaction summary failed: %w", aerr) + } + var content string + if ch := reply.ChatCompletionResponse.Choices; len(ch) > 0 { + content = ch[0].Message.Content } - last := res.LastMessage() - if last == nil || strings.TrimSpace(last.Content) == "" { - return "", fmt.Errorf("compaction produced an empty summary") + if strings.TrimSpace(content) == "" { + return "", sentTokens, fmt.Errorf("compaction produced an empty summary") } - return last.Content, nil + return content, sentTokens, nil } // summaryMessage wraps a compaction summary as the message that stands in for diff --git a/chat/compact_fit_test.go b/chat/compact_fit_test.go index 62980bf..b8695c9 100644 --- a/chat/compact_fit_test.go +++ b/chat/compact_fit_test.go @@ -4,6 +4,7 @@ import ( "context" "errors" "fmt" + "reflect" "strings" "testing" @@ -20,16 +21,18 @@ type promptCapturingLLM struct { } func (f *promptCapturingLLM) Ask(ctx context.Context, fr cogito.Fragment) (cogito.Fragment, error) { - prompt := fr.Messages[len(fr.Messages)-1].Content - f.prompts = append(f.prompts, prompt) - if f.limit > 0 && tokensOf(prompt) > f.limit { - return cogito.Fragment{}, fmt.Errorf("rpc error: code = Internal desc = request (%d tokens) exceeds the available context size (%d tokens), try increasing it", tokensOf(prompt), f.limit) - } - return fr.AddMessage(cogito.AssistantMessageRole, "SUMMARY"), nil + return cogito.Fragment{}, errors.New("the summary must be sent with CreateChatCompletion, not Ask") } func (f *promptCapturingLLM) CreateChatCompletion(ctx context.Context, req openai.ChatCompletionRequest) (cogito.LLMReply, cogito.LLMUsage, error) { - return cogito.LLMReply{}, cogito.LLMUsage{}, nil + prompt := req.Messages[len(req.Messages)-1].Content + f.prompts = append(f.prompts, prompt) + if f.limit > 0 && tokensOf(prompt) > f.limit { + return cogito.LLMReply{}, cogito.LLMUsage{}, fmt.Errorf("rpc error: code = Internal desc = request (%d tokens) exceeds the available context size (%d tokens), try increasing it", tokensOf(prompt), f.limit) + } + return cogito.LLMReply{ChatCompletionResponse: openai.ChatCompletionResponse{ + Choices: []openai.ChatCompletionChoice{{Message: openai.ChatCompletionMessage{Role: "assistant", Content: "SUMMARY"}}}, + }}, cogito.LLMUsage{}, nil } // toolTurn is one assistant tool call plus its result. @@ -103,7 +106,8 @@ func TestCompactHistoryFitsTheSummaryPromptInTheBudget(t *testing.T) { func TestCompactHistoryRetriesSmallerWhenTheSummaryOverflows(t *testing.T) { // The backend tokenizes denser than byte/4, so a prompt the estimate says - // fits can still overflow. The backend's own figures say by how much. + // fits can still overflow. The backend's own figures say by how much, and + // the retry moves the boundary so the head it summarizes fits. frag := []openai.ChatCompletionMessage{{Role: "user", Content: "goal"}} for i := range 6 { frag = append(frag, toolTurn(fmt.Sprintf("c%d", i), fmt.Sprintf("f%d.go", i), strings.Repeat("x", 8000))...) @@ -112,6 +116,7 @@ func TestCompactHistoryRetriesSmallerWhenTheSummaryOverflows(t *testing.T) { openai.ChatCompletionMessage{Role: "user", Content: "u2"}, openai.ChatCompletionMessage{Role: "assistant", Content: "a2"}, ) + orig := cloneMessages(frag) llm := &promptCapturingLLM{limit: 1500} s := newCompactTestSession(llm, 2, frag, frag) s.compaction.MaxContextTokens = 8000 @@ -119,21 +124,26 @@ func TestCompactHistoryRetriesSmallerWhenTheSummaryOverflows(t *testing.T) { if _, _, err := s.CompactHistory(); err != nil { t.Fatalf("CompactHistory: %v", err) } - if len(llm.prompts) != 2 { - t.Fatalf("summary calls = %d, want 2 (one overflow, one smaller retry)", len(llm.prompts)) + if n := len(llm.prompts); n < 2 || n > maxSummaryAttempts { + t.Fatalf("summary calls = %d, want an overflow and then smaller retries", n) } - if !strings.Contains(s.fragment.Messages[0].Content, "SUMMARY") { - t.Fatal("the retried summary was not installed") + if got := tokensOf(llm.prompts[len(llm.prompts)-1]); got > 1500 { + t.Fatalf("the accepted prompt is ~%d tokens, over the backend's 1500", got) } + // The messages the head gave up are kept verbatim instead. + checkVerbatimTail(t, orig, s.fragment.Messages) } -func TestCompactHistoryRetriesTheSummaryOnlyOnce(t *testing.T) { +func TestCompactHistoryGivesUpWhenNoSummaryFits(t *testing.T) { + // Not even the instruction fits, and the rejection carries no usable + // figures: the loop halves until it runs out of attempts. frag := []openai.ChatCompletionMessage{ {Role: "user", Content: strings.Repeat("u", 4000)}, {Role: "assistant", Content: "a1"}, {Role: "user", Content: "u2"}, {Role: "assistant", Content: "a2"}, } + orig := cloneMessages(frag) llm := &promptCapturingLLM{limit: 10} s := newCompactTestSession(llm, 1, frag, frag) s.compaction.MaxContextTokens = 1 << 20 @@ -142,7 +152,10 @@ func TestCompactHistoryRetriesTheSummaryOnlyOnce(t *testing.T) { if !isContextOverflow(errors.Unwrap(err)) { t.Fatalf("want the summarizer's overflow, got %v", err) } - if len(llm.prompts) != 2 { - t.Fatalf("summary calls = %d, want exactly 2", len(llm.prompts)) + if len(llm.prompts) != maxSummaryAttempts { + t.Fatalf("summary calls = %d, want exactly %d", len(llm.prompts), maxSummaryAttempts) + } + if !reflect.DeepEqual(s.fragment.Messages, orig) { + t.Fatal("a failed compaction changed the fragment") } } diff --git a/chat/compact_internal_test.go b/chat/compact_internal_test.go index 96a9271..c10baa2 100644 --- a/chat/compact_internal_test.go +++ b/chat/compact_internal_test.go @@ -96,11 +96,14 @@ func TestEstimateAndHumanTokens(t *testing.T) { } } -// fakeSummaryLLM is a minimal cogito.LLM that returns a canned summary from Ask. +// fakeSummaryLLM is a minimal cogito.LLM that returns a canned summary from +// Ask and from CreateChatCompletion, which is what compaction sends. lastReq is +// the last CreateChatCompletion request. type fakeSummaryLLM struct { - reply string - err error - calls int + reply string + err error + calls int + lastReq openai.ChatCompletionRequest } func (f *fakeSummaryLLM) Ask(ctx context.Context, fr cogito.Fragment) (cogito.Fragment, error) { @@ -116,7 +119,19 @@ func (f *fakeSummaryLLM) Ask(ctx context.Context, fr cogito.Fragment) (cogito.Fr } func (f *fakeSummaryLLM) CreateChatCompletion(ctx context.Context, req openai.ChatCompletionRequest) (cogito.LLMReply, cogito.LLMUsage, error) { - return cogito.LLMReply{}, cogito.LLMUsage{}, nil + f.calls++ + f.lastReq = req + if f.err != nil { + return cogito.LLMReply{}, cogito.LLMUsage{}, f.err + } + return replyWith(f.reply), cogito.LLMUsage{TotalTokens: 1}, nil +} + +// replyWith is a completion whose single choice says content. +func replyWith(content string) cogito.LLMReply { + return cogito.LLMReply{ChatCompletionResponse: openai.ChatCompletionResponse{ + Choices: []openai.ChatCompletionChoice{{Message: openai.ChatCompletionMessage{Role: "assistant", Content: content}}}, + }} } func newCompactTestSession(llm cogito.LLM, keepRecent int, frag []openai.ChatCompletionMessage, disp []openai.ChatCompletionMessage) *Session { diff --git a/chat/midturn.go b/chat/midturn.go index 9976366..1b7dccd 100644 --- a/chat/midturn.go +++ b/chat/midturn.go @@ -2,7 +2,6 @@ package chat import ( "context" - "strings" "github.com/mudler/cogito" "github.com/mudler/xlog" @@ -126,49 +125,43 @@ func (c *turnCompactor) compact(msgs, out []openai.ChatCompletionMessage, base, // The kept tail has to fit too, or the summary buys nothing and the next // step compacts again. Shrink it toward the last tool step when it does // not; splitForCompaction still never separates a call from its results. - var head, tail []openai.ChatCompletionMessage - for keep := max(c.s.compactionConfig().KeepRecent, 1); keep >= 1; keep-- { - h, t := splitForCompaction(out, keep) + keep := 0 + for k := max(c.s.compactionConfig().KeepRecent, 1); k >= 1; k-- { + h, t := splitForCompaction(out, k) if h == nil { continue } - head, tail = h, t + keep = k if !c.s.shouldCompactNow(estimateTokens(t) + overhead) { break } } // A head that is only the previous summary has nothing new in it. - if len(head) <= repl { + if keep == 0 { return out } - pieces := renderMessages(head) - if len(pieces) == 0 { + if h, _ := splitForCompaction(out, keep); len(h) <= repl { return out } - // Save the full untruncated head as a compaction artifact, the same - // as the end-of-turn path, so nothing is lost to the lossy summary. - var artifactURI string - cfg := c.s.compactionConfig() - if !cfg.DisableArtifactSpill && c.s.artifacts != nil { - var b strings.Builder - for _, p := range pieces { - b.WriteString(p.text) - } - if b.Len() > 0 { - artifactURI = c.s.artifacts.Save("compaction", b.String()) - } - } - if c.s.callbacks.OnStatus != nil { c.s.callbacks.OnStatus("Compacting conversation…") } - summary, err := c.s.summarize(c.ctx, pieces) + // out is already the pruned request, so the summary sees it as it is. + summary, head, tail, err := c.s.summarizeFitting(c.ctx, out, keep, func(m []openai.ChatCompletionMessage) []openai.ChatCompletionMessage { return m }) if err != nil { xlog.Warn("mid-turn compaction failed", "error", err) c.failed = true return out } + // The overflow retries may have moved the boundary back until the head + // is only the previous summary, which has nothing new in it. + if len(head) <= repl { + return out + } + // Save the head actually summarized as a compaction artifact, the same + // as the end-of-turn path, so nothing is lost to the lossy summary. + artifactURI := c.s.spillCompactionHead(renderMessages(head)) // renderMessages leaves system messages out of the summary, so they are // kept as they are: dropping them would lose the system prompt. diff --git a/chat/overflow_retry_test.go b/chat/overflow_retry_test.go index 6bbf3be..8280a6c 100644 --- a/chat/overflow_retry_test.go +++ b/chat/overflow_retry_test.go @@ -14,7 +14,9 @@ import ( ) // overflowLLM fails the first N completion calls with a backend overflow error, -// then succeeds. Ask always succeeds, so compaction's summarising call works. +// then succeeds. Compaction's summary request (sent through +// CreateChatCompletion, or Ask) always succeeds and is not counted as a +// completion call, so compaction's summarising call works. type overflowLLM struct { mu sync.Mutex failures int // remaining calls that should overflow @@ -26,6 +28,9 @@ func (o *overflowLLM) CreateChatCompletion(ctx context.Context, req openai.ChatC if err := ctx.Err(); err != nil { return cogito.LLMReply{}, cogito.LLMUsage{}, err } + if n := len(req.Messages); n > 0 && strings.HasPrefix(req.Messages[n-1].Content, compactInstruction) { + return replyWith(o.summary()), cogito.LLMUsage{PromptTokens: 20, CompletionTokens: 2, TotalTokens: 22}, nil + } o.mu.Lock() o.calls++ fail := o.failures > 0 @@ -57,19 +62,22 @@ func (o *overflowLLM) CreateChatCompletion(ctx context.Context, req openai.ChatC // hiding the infinite loop the cap exists to prevent. Varying the text puts the // cap, and only the cap, on the hook. func (o *overflowLLM) Ask(ctx context.Context, f cogito.Fragment) (cogito.Fragment, error) { - o.mu.Lock() - o.asks++ - n := o.asks - o.mu.Unlock() - - summary := "summary " + strconv.Itoa(n) + " of earlier turns" + strings.Repeat(" with more detail", n) - out := f.AddMessage("assistant", summary) + out := f.AddMessage("assistant", o.summary()) if out.Status != nil { out.Status.LastUsage = cogito.LLMUsage{PromptTokens: 20, CompletionTokens: 2, TotalTokens: 22} } return out, nil } +// summary counts a summary request and returns its (distinct) text. +func (o *overflowLLM) summary() string { + o.mu.Lock() + o.asks++ + n := o.asks + o.mu.Unlock() + return "summary " + strconv.Itoa(n) + " of earlier turns" + strings.Repeat(" with more detail", n) +} + func newOverflowSession(t *testing.T, llm cogito.LLM) *Session { t.Helper() s := &Session{ @@ -468,6 +476,13 @@ func (f *summaryFailingLLM) Ask(ctx context.Context, frag cogito.Fragment) (cogi return frag, errors.New("request (368203 tokens) exceeds the available context size (262144 tokens)") } +func (f *summaryFailingLLM) CreateChatCompletion(ctx context.Context, req openai.ChatCompletionRequest) (cogito.LLMReply, cogito.LLMUsage, error) { + if n := len(req.Messages); n > 0 && strings.HasPrefix(req.Messages[n-1].Content, compactInstruction) { + return cogito.LLMReply{}, cogito.LLMUsage{}, errors.New("request (368203 tokens) exceeds the available context size (262144 tokens)") + } + return f.overflowLLM.CreateChatCompletion(ctx, req) +} + // When compaction's own call fails, the user gets the ORIGINAL overflow — the // one they can act on — and no retry is attempted. func TestOverflowReportsOriginalErrorWhenCompactionFails(t *testing.T) { diff --git a/chat/summary_shift_test.go b/chat/summary_shift_test.go new file mode 100644 index 0000000..109ea1c --- /dev/null +++ b/chat/summary_shift_test.go @@ -0,0 +1,281 @@ +package chat + +import ( + "context" + "errors" + "fmt" + "reflect" + "strings" + "testing" + + "github.com/mudler/cogito" + "github.com/mudler/nib/mcp" + openai "github.com/sashabaranov/go-openai" +) + +// vllmOverflow is vLLM's current rejection of a request whose prompt plus the +// output it reserves exceeds the window. +func vllmOverflow(window, output, prompt int) error { + return fmt.Errorf("This model's maximum context length is %d tokens. However, you requested %d output tokens and your prompt contains %d input tokens, for a total of %d tokens. Please reduce the length of the input prompt or the number of requested output tokens.", + window, output, prompt, prompt+output) +} + +// reservingLLM is a backend that checks the output reservation the way vLLM +// does: it rejects a request when prompt + max_tokens exceeds limit, whatever +// the model would actually generate. The prompt is counted byte/4. +type reservingLLM struct { + // limit is the backend's window. 0 accepts every size. + limit int + // overBy, when positive, rejects every request as overflowing by that + // many tokens. + overBy int + // err, when set, is returned for every request. + err error + reqs []openai.ChatCompletionRequest +} + +func (f *reservingLLM) Ask(ctx context.Context, fr cogito.Fragment) (cogito.Fragment, error) { + return cogito.Fragment{}, errors.New("the summary must be sent with CreateChatCompletion, not Ask") +} + +func (f *reservingLLM) CreateChatCompletion(ctx context.Context, req openai.ChatCompletionRequest) (cogito.LLMReply, cogito.LLMUsage, error) { + f.reqs = append(f.reqs, req) + usage := cogito.LLMUsage{PromptTokens: 1, TotalTokens: 1} + if f.err != nil { + return cogito.LLMReply{}, usage, f.err + } + p := tokensOf(req.Messages[len(req.Messages)-1].Content) + if f.overBy > 0 { + return cogito.LLMReply{}, usage, vllmOverflow(p+req.MaxTokens-f.overBy, req.MaxTokens, p) + } + if f.limit > 0 && p+req.MaxTokens > f.limit { + return cogito.LLMReply{}, usage, vllmOverflow(f.limit, req.MaxTokens, p) + } + return cogito.LLMReply{ChatCompletionResponse: openai.ChatCompletionResponse{ + Choices: []openai.ChatCompletionChoice{{Message: openai.ChatCompletionMessage{Role: "assistant", Content: "SUMMARY"}}}, + }}, usage, nil +} + +// longHistory is a goal, n tool steps of ~2000 tokens each, and a short +// closing exchange. +func longHistory(n int) []openai.ChatCompletionMessage { + msgs := []openai.ChatCompletionMessage{{Role: "user", Content: "goal: fix the parser"}} + for i := range n { + msgs = append(msgs, toolTurn(fmt.Sprintf("c%d", i), fmt.Sprintf("f%d.go", i), fmt.Sprintf("BODY%02d", i)+strings.Repeat("x", 8000))...) + } + return append(msgs, + openai.ChatCompletionMessage{Role: "user", Content: "u2"}, + openai.ChatCompletionMessage{Role: "assistant", Content: "a2"}, + ) +} + +func cloneMessages(msgs []openai.ChatCompletionMessage) []openai.ChatCompletionMessage { + return append([]openai.ChatCompletionMessage(nil), msgs...) +} + +// checkVerbatimTail asserts that got is [summary] followed by a verbatim +// suffix of orig that starts on a clean boundary, and returns how many of +// orig's messages the summary stands for. +func checkVerbatimTail(t *testing.T, orig, got []openai.ChatCompletionMessage) int { + t.Helper() + if len(got) < 2 || !strings.Contains(got[0].Content, "SUMMARY") { + t.Fatalf("want the summary followed by the kept messages, got %d messages", len(got)) + } + kept := got[1:] + start := len(orig) - len(kept) + if start < 1 { + t.Fatalf("kept %d of %d messages", len(kept), len(orig)) + } + if !reflect.DeepEqual(kept, orig[start:]) { + t.Fatal("the kept messages are not a verbatim, in-order suffix of the history") + } + if kept[0].Role == "tool" || len(orig[start-1].ToolCalls) > 0 { + t.Fatal("the boundary separates a tool call from its result") + } + return start +} + +func TestSummaryShiftConvergesWithAFixedOutputReservation(t *testing.T) { + // The session believes the window is huge, the backend serves 8192 and + // reserves the full 4000 summary tokens on every request. Scaling the + // prompt by allows/needs does not converge: the reservation does not + // scale with it. + frag := longHistory(20) + orig := cloneMessages(frag) + llm := &reservingLLM{limit: 8192} + s := newCompactTestSession(llm, 2, frag, frag) + s.compaction.MaxContextTokens = 1 << 20 + s.compaction.SummaryMaxTokens = 4000 + + if _, _, err := s.CompactHistory(); err != nil { + t.Fatalf("CompactHistory: %v", err) + } + if n := len(llm.reqs); n < 2 || n > maxSummaryAttempts { + t.Fatalf("summary calls = %d, want an overflow then a fitting retry", n) + } + for i, r := range llm.reqs { + if r.MaxTokens != 4000 { + t.Fatalf("request %d MaxTokens = %d, want SummaryMaxTokens 4000", i, r.MaxTokens) + } + } + start := checkVerbatimTail(t, orig, s.fragment.Messages) + if len(orig)-start <= 2 { + t.Fatal("the boundary did not move: the kept tail is still KeepRecent long") + } + if !strings.Contains(llm.reqs[len(llm.reqs)-1].Messages[0].Content, "goal: fix the parser") { + t.Fatal("the shifted head lost the user's goal") + } +} + +func TestSummaryRequestCarriesTheSummaryOutputCap(t *testing.T) { + for _, tc := range []struct { + name string + cfg, outputCap int + wantMaxTokens int + }{ + {"configured", 1234, 0, 1234}, + {"default", 0, 0, 16384}, + {"model cap is smaller", 1234, 700, 700}, + } { + t.Run(tc.name, func(t *testing.T) { + frag := longHistory(2) + llm := &reservingLLM{} + s := newCompactTestSession(llm, 2, frag, frag) + s.compaction.MaxContextTokens = 1 << 20 + s.compaction.SummaryMaxTokens = tc.cfg + s.outputCap = tc.outputCap + if _, _, err := s.CompactHistory(); err != nil { + t.Fatalf("CompactHistory: %v", err) + } + if len(llm.reqs) != 1 || llm.reqs[0].MaxTokens != tc.wantMaxTokens { + t.Fatalf("requests = %d, MaxTokens = %d, want one with %d", len(llm.reqs), llm.reqs[0].MaxTokens, tc.wantMaxTokens) + } + }) + } +} + +func TestSummaryGivesUpAfterMaxAttempts(t *testing.T) { + frag := longHistory(6) + orig := cloneMessages(frag) + llm := &reservingLLM{overBy: 50} + s := newCompactTestSession(llm, 2, frag, frag) + s.compaction.MaxContextTokens = 1 << 20 + s.compaction.SummaryMaxTokens = 1000 + s.artifacts = mcp.NewArtifactStore() + + _, _, err := s.CompactHistory() + if !isContextOverflow(err) { + t.Fatalf("want the last overflow, got %v", err) + } + if len(llm.reqs) != maxSummaryAttempts { + t.Fatalf("summary calls = %d, want %d", len(llm.reqs), maxSummaryAttempts) + } + if !reflect.DeepEqual(s.fragment.Messages, orig) { + t.Fatal("a failed compaction changed the fragment") + } + if s.artifacts.Count() != 0 { + t.Fatal("a failed compaction spilled an artifact") + } +} + +func TestSummaryStopsOnANonOverflowError(t *testing.T) { + frag := longHistory(6) + orig := cloneMessages(frag) + llm := &reservingLLM{err: errors.New("status code: 400, bad request")} + s := newCompactTestSession(llm, 2, frag, frag) + s.compaction.MaxContextTokens = 1 << 20 + + if _, _, err := s.CompactHistory(); err == nil { + t.Fatal("want the backend error") + } + if len(llm.reqs) != 1 { + t.Fatalf("summary calls = %d, want 1", len(llm.reqs)) + } + if !reflect.DeepEqual(s.fragment.Messages, orig) { + t.Fatal("a failed compaction changed the fragment") + } +} + +func TestSummaryStopsWhenTheReservationAloneExceedsTheWindow(t *testing.T) { + frag := longHistory(6) + llm := &reservingLLM{limit: 8192} + s := newCompactTestSession(llm, 2, frag, frag) + s.compaction.MaxContextTokens = 1 << 20 + s.compaction.SummaryMaxTokens = 10000 + + if _, _, err := s.CompactHistory(); !isContextOverflow(err) { + t.Fatalf("want the overflow, got %v", err) + } + if len(llm.reqs) != 1 { + t.Fatalf("summary calls = %d, want 1: no prompt fits beside a 10000-token reservation in 8192", len(llm.reqs)) + } +} + +func TestSummaryArtifactHoldsExactlyTheSummarizedHead(t *testing.T) { + frag := longHistory(20) + orig := cloneMessages(frag) + llm := &reservingLLM{limit: 8192} + s := newCompactTestSession(llm, 2, frag, frag) + s.compaction.MaxContextTokens = 1 << 20 + s.compaction.SummaryMaxTokens = 4000 + s.artifacts = mcp.NewArtifactStore() + + if _, _, err := s.CompactHistory(); err != nil { + t.Fatalf("CompactHistory: %v", err) + } + start := checkVerbatimTail(t, orig, s.fragment.Messages) + if s.artifacts.Count() != 1 { + t.Fatalf("artifacts = %d, want 1", s.artifacts.Count()) + } + var want strings.Builder + for _, p := range renderMessages(orig[:start]) { + want.WriteString(p.text) + } + if got := s.artifacts.Get(1).Content; got != want.String() { + t.Fatalf("artifact is %d bytes, want exactly the %d-byte summarized head", len(got), want.Len()) + } +} + +func TestMidTurnCompactionShiftsTheBoundaryWhenTheSummaryOverflows(t *testing.T) { + msgs := longHistory(20) + llm := &reservingLLM{limit: 8192} + s := newCompactTestSession(llm, 2, msgs, msgs) + s.compaction.MaxContextTokens = 1 << 20 + s.compaction.SummaryMaxTokens = 4000 + c := s.newTurnCompactor(context.Background()) + + out := c.compact(msgs, msgs, 0, 0, 0) + if len(llm.reqs) < 2 { + t.Fatalf("summary calls = %d, want the loop to retry", len(llm.reqs)) + } + start := checkVerbatimTail(t, msgs, out) + if len(msgs)-start <= 2 { + t.Fatal("the boundary did not move") + } + if c.covered != start { + t.Fatalf("covered = %d, want %d", c.covered, start) + } +} + +func TestMidTurnCompactionSkipsAHeadOfOnlyThePreviousSummary(t *testing.T) { + raw := longHistory(5) + // The request already carries a summary in place of the first 3 messages. + prev := openai.ChatCompletionMessage{Role: "user", Content: "PREV"} + out := append([]openai.ChatCompletionMessage{prev}, cloneMessages(raw[3:])...) + llm := &reservingLLM{limit: 3000} + s := newCompactTestSession(llm, 2, raw, raw) + s.compaction.MaxContextTokens = 1 << 20 + s.compaction.SummaryMaxTokens = 1000 + c := s.newTurnCompactor(context.Background()) + + got := c.compact(raw, out, 3, 1, 0) + if len(llm.reqs) < 2 { + t.Fatalf("summary calls = %d, want the loop to shrink the head to the previous summary", len(llm.reqs)) + } + if !reflect.DeepEqual(got, out) { + t.Fatal("a head of only the previous summary should leave the request unchanged") + } + if c.covered != 0 { + t.Fatalf("covered = %d, want 0", c.covered) + } +} diff --git a/chat/usage_test.go b/chat/usage_test.go index e37d656..b1c6777 100644 --- a/chat/usage_test.go +++ b/chat/usage_test.go @@ -359,12 +359,12 @@ func TestRunningAgentEventAddsNothing(t *testing.T) { } } -// summarizingLLM answers Ask with a summary and reports what that call cost, +// summarizingLLM answers Ask and CreateChatCompletion with a summary and reports what that call cost, // which is the figure compaction has been spending invisibly. type summarizingLLM struct{} func (summarizingLLM) CreateChatCompletion(ctx context.Context, req openai.ChatCompletionRequest) (cogito.LLMReply, cogito.LLMUsage, error) { - return cogito.LLMReply{}, cogito.LLMUsage{}, nil + return replyWith("a summary of earlier turns"), cogito.LLMUsage{PromptTokens: 60, CompletionTokens: 6, TotalTokens: 66}, nil } func (summarizingLLM) Ask(ctx context.Context, f cogito.Fragment) (cogito.Fragment, error) { @@ -409,12 +409,12 @@ func TestCompactionUsageCountsTowardTheSession(t *testing.T) { // failingSummaryLLM models a compaction call the backend served and billed and // that then failed — a truncated response, a mid-stream disconnect. It reports -// the usage on the returned fragment before returning the error, which is the -// only shape in which a client can hand back spend it already incurred. +// the usage alongside the error (on the fragment, for Ask), which is the only +// shape in which a client can hand back spend it already incurred. type failingSummaryLLM struct{} func (failingSummaryLLM) CreateChatCompletion(ctx context.Context, req openai.ChatCompletionRequest) (cogito.LLMReply, cogito.LLMUsage, error) { - return cogito.LLMReply{}, cogito.LLMUsage{}, nil + return cogito.LLMReply{}, cogito.LLMUsage{PromptTokens: 40, CompletionTokens: 4, TotalTokens: 44}, errors.New("backend hung up after answering") } func (failingSummaryLLM) Ask(ctx context.Context, f cogito.Fragment) (cogito.Fragment, error) { diff --git a/config/config.go b/config/config.go index 83f40d1..0e21925 100644 --- a/config/config.go +++ b/config/config.go @@ -423,6 +423,10 @@ func withDefaults(cfg types.Config) types.Config { if cfg.Compaction.ReserveTokens == 0 { cfg.Compaction.ReserveTokens = 4096 } + // Must match chat's defaultSummaryMaxTokens, for the same reason. + if cfg.Compaction.SummaryMaxTokens == 0 { + cfg.Compaction.SummaryMaxTokens = 16384 + } // Tool-output pruning is defaulted as a block, not field by field: a zero // HighWaterTokens is meaningful on its own (it disables size pruning while diff --git a/config/config_test.go b/config/config_test.go index ab0b7ea..641f4dc 100644 --- a/config/config_test.go +++ b/config/config_test.go @@ -32,7 +32,7 @@ func TestWithDefaultsKeepsUserValues(t *testing.T) { // comparison below would fail for a reason that is not an override bug. in := types.Config{Compaction: types.CompactionConfig{ MaxContextTokens: 200000, Threshold: 0.5, KeepRecent: 2, Disabled: true, - ReserveTokens: 512, + ReserveTokens: 512, SummaryMaxTokens: 2048, }} got := withDefaults(in) if got.Compaction != in.Compaction { diff --git a/config/reserve_config_test.go b/config/reserve_config_test.go index 38e0d40..93e2910 100644 --- a/config/reserve_config_test.go +++ b/config/reserve_config_test.go @@ -38,3 +38,26 @@ func TestReserveTokensBindsThroughTheRealLoadPath(t *testing.T) { t.Fatalf("ReserveTokens = %d, want 1234: 'reserve_tokens' did not bind", cfg.Compaction.ReserveTokens) } } + +func TestSummaryMaxTokensDefaults(t *testing.T) { + cfg := withDefaults(types.Config{}) + if cfg.Compaction.SummaryMaxTokens != 16384 { + t.Fatalf("SummaryMaxTokens = %d, want 16384", cfg.Compaction.SummaryMaxTokens) + } + cfg = withDefaults(types.Config{Compaction: types.CompactionConfig{SummaryMaxTokens: 2048}}) + if cfg.Compaction.SummaryMaxTokens != 2048 { + t.Fatalf("SummaryMaxTokens = %d, want 2048 preserved", cfg.Compaction.SummaryMaxTokens) + } +} + +func TestSummaryMaxTokensBindsThroughTheRealLoadPath(t *testing.T) { + dir := t.TempDir() + body := "compaction:\n summary_max_tokens: 3000\n" + if err := os.WriteFile(filepath.Join(dir, "config.yaml"), []byte(body), 0o644); err != nil { + t.Fatal(err) + } + cfg := LoadWith(LoadOptions{BaseDir: dir, SkipBareEnv: true}) + if cfg.Compaction.SummaryMaxTokens != 3000 { + t.Fatalf("SummaryMaxTokens = %d, want 3000: 'summary_max_tokens' did not bind", cfg.Compaction.SummaryMaxTokens) + } +} diff --git a/selfdoc/README.md b/selfdoc/README.md index a7630d7..1505974 100644 --- a/selfdoc/README.md +++ b/selfdoc/README.md @@ -568,6 +568,16 @@ agents: # LOG_FORMAT=json switches to JSON lines. log_level: error +# Optional: conversation compaction. When the context fills up, nib +# summarizes the older turns and keeps the recent ones verbatim. +# summary_max_tokens caps the output the summary request reserves (default +# 16384; the model's own output cap wins when it is smaller). The summary +# prompt is fitted into the window minus that reservation. If the backend +# still rejects it as too large, nib summarizes fewer of the older turns and +# keeps the rest verbatim, instead of cutting the text it summarizes. +compaction: + summary_max_tokens: 16384 + # Optional: external MCP servers mcp_servers: filesystem: diff --git a/types/config.go b/types/config.go index c4ea5b0..3cc774a 100644 --- a/types/config.go +++ b/types/config.go @@ -49,6 +49,11 @@ type CompactionConfig struct { // compaction with nothing to trigger on (see chat.ContextBudget). At the // 4096 default the cap applies only below an 8192-token window. ReserveTokens int `yaml:"reserve_tokens"` + // SummaryMaxTokens caps the output the compaction summary request + // reserves. The request asks for the smaller of this and the model's own + // output cap, and its prompt is fitted into the window minus that + // reservation. 0 → default 16384. + SummaryMaxTokens int `yaml:"summary_max_tokens"` // DisableArtifactSpill turns OFF saving the full conversation head as an // artifact during compaction. Zero value (false) = spill ON (default): // the head is saved and the summary includes an artifact://N reference. From c133426befb19996159e1faf24df684ce6855a72 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Fri, 25 Sep 2026 20:45:58 +0000 Subject: [PATCH 03/16] fix(chat): summarize an overflowing head in rolling chunks When the summary of the whole head overflowed, the boundary moved earlier and the cut messages stayed verbatim. Nothing was lost, but a far move left the history barely smaller, so the next request compacted again. Now the head is summarized in chunks from oldest to newest, each fitted to the retry target. Every chunk after the first carries the running summary and asks to merge it. A chunk that overflows anyway shrinks and passes the rest to the next chunk. Chunks end only on tool-call boundaries, and an over-large message is cut down, not skipped. After 8 chunks the uncovered messages stay verbatim, as before. Co-Authored-By: Claude Opus 5.5 --- chat/compact.go | 166 +++++++++++++++++---- chat/compact_fit_test.go | 8 +- chat/summary_rolling_test.go | 282 +++++++++++++++++++++++++++++++++++ chat/summary_shift_test.go | 28 ++-- 4 files changed, 439 insertions(+), 45 deletions(-) create mode 100644 chat/summary_rolling_test.go diff --git a/chat/compact.go b/chat/compact.go index 9bbf0d4..b0390e6 100644 --- a/chat/compact.go +++ b/chat/compact.go @@ -678,47 +678,147 @@ func (s *Session) summaryPromptLimit(output int) int { // (the pruned view of them). An empty head, with a nil error, means there was // nothing to compact. // -// When the backend rejects the summary as too large, the boundary moves -// earlier: keep grows until the head fits the retry target (see -// summaryRetryTarget), or until only the smallest head is left, which -// fitSummaryInput then truncates to the target. What leaves the head joins -// the tail, verbatim: nothing is lost, and the next compaction can summarize -// it. The loop stops on success, a non-overflow error, a cancelled context, a -// target of zero or less, or after maxSummaryAttempts, and returns the last -// error. It changes no session state other than the usage it counts. +// The whole head goes out in one request first. When the backend rejects it +// as too large, the head is summarized in chunks instead (see rollingCover), +// each fitted to the retry target summaryRetryTarget derives from the +// rejection. When the chunks cover the whole head, the result is the summary +// and the original tail. When they stop at maxRollingChunks, the head ends +// where they stopped and the rest joins the tail, verbatim: nothing is lost, +// and the next compaction can summarize it. The loop stops on a non-overflow +// error, a cancelled context, a target of zero or less, or after +// maxSummaryAttempts overflows in a row, and returns the last error. It +// changes no session state other than the usage it counts. func (s *Session) summarizeFitting(ctx context.Context, msgs []openai.ChatCompletionMessage, keep int, view func([]openai.ChatCompletionMessage) []openai.ChatCompletionMessage) (summary string, head, tail []openai.ChatCompletionMessage, err error) { head, tail = splitForCompaction(msgs, keep) if len(head) == 0 || len(renderMessages(view(head))) == 0 { return "", nil, nil, nil } - keep = len(tail) - prefix := tokensOf(summaryPrefix) limit := s.summaryPromptLimit(s.summaryOutputTokens()) - for attempt := 0; ; attempt++ { - summary, sent, serr := s.summarize(ctx, renderMessages(view(head)), limit) - if serr == nil { - return summary, head, tail, nil - } - if attempt+1 >= maxSummaryAttempts || !canRecoverFromOverflow(ctx, serr) { - return "", nil, nil, serr + summary, sent, serr := s.summarize(ctx, renderMessages(view(head)), limit) + if serr == nil { + return summary, head, tail, nil + } + if maxSummaryAttempts <= 1 || !canRecoverFromOverflow(ctx, serr) { + return "", nil, nil, serr + } + target := summaryRetryTarget(serr, sent) + if target <= 0 { + return "", nil, nil, serr + } + summary, covered, err := s.rollingCover(ctx, head, view, target, 1) + if err != nil { + return "", nil, nil, err + } + // head is msgs[:len(head)], so what the chunks did not cover and the + // original tail are one contiguous run of msgs. + return summary, msgs[:covered], msgs[covered:], nil +} + +// maxRollingChunks bounds how many chunks rollingCover summarizes for one +// compaction. What is left after that stays verbatim. +const maxRollingChunks = 8 + +// rollingInstruction is the instruction for every chunk after the first: the +// running summary comes first in the prompt, the chunk's messages after it. +const rollingInstruction = compactInstruction + "\n\n" + + "A summary of the earlier part of this conversation is given first. " + + "Merge it with the new messages into one summary in the same format; " + + "keep every file path, identifier and decision from both." + +// rollingPrefix is the prompt a later chunk's rendered messages follow. +func rollingPrefix(running string) string { + return rollingInstruction + "\n\n--- EARLIER SUMMARY ---\n" + running + "\n\n--- NEW MESSAGES ---\n" +} + +// rollingSummary summarizes head in consecutive chunks, each fitted to a +// prompt of target byte/4 tokens, and returns the one summary that stands +// for all of head. It fails when maxRollingChunks chunks do not cover head; +// summarizeFitting, which keeps the uncovered rest verbatim instead, uses +// rollingCover directly. +func (s *Session) rollingSummary(ctx context.Context, head []openai.ChatCompletionMessage, view func([]openai.ChatCompletionMessage) []openai.ChatCompletionMessage, target int) (summary string, err error) { + summary, covered, err := s.rollingCover(ctx, head, view, target, 0) + if err != nil { + return "", err + } + if covered < len(head) { + return "", fmt.Errorf("compaction summary covered %d of %d messages in %d chunks", covered, len(head), maxRollingChunks) + } + return summary, nil +} + +// rollingCover summarizes head from oldest to newest in chunks, and returns +// the running summary and how many of head's messages it stands for. +// +// Each chunk is the longest run that fits target: the first after the +// compaction instruction, each later one after rollingInstruction and the +// running summary, whose reply then replaces the running summary. A chunk +// ends only where splitForCompaction could, so it never separates a tool +// call from its results. A run too large for target on its own still forms +// a chunk, and fitSummaryInput cuts it down. When the backend rejects a chunk +// anyway, target drops by the stated overshoot (summaryRetryTarget) and the +// chunk is rebuilt from the same start, so the rest moves on to the next +// chunk. failures is how many overflows in a row the caller already had; +// maxSummaryAttempts of them in a row stop it. It stops after +// maxRollingChunks chunks, and on any error it returns nothing covered. +func (s *Session) rollingCover(ctx context.Context, head []openai.ChatCompletionMessage, view func([]openai.ChatCompletionMessage) []openai.ChatCompletionMessage, target, failures int) (summary string, covered int, err error) { + viewed := view(head) + if len(viewed) != len(head) { + return "", 0, fmt.Errorf("compaction view changed the message count") + } + sizes := make([]int, len(viewed)) + for i := range viewed { + for _, p := range renderMessages(viewed[i : i+1]) { + sizes[i] += len(p.text) } - target := summaryRetryTarget(serr, sent) - if target <= 0 { - return "", nil, nil, serr + } + boundary := func(i int) bool { + return i == len(head) || (head[i].Role != "tool" && len(head[i-1].ToolCalls) == 0) + } + + running := "" + start, chunks := 0, 0 + for start < len(head) && chunks < maxRollingChunks { + prefix := summaryPrefix + if chunks > 0 { + prefix = rollingPrefix(running) } - limit = max(target-prefix, 1) - // Move the boundary earlier until the head fits the target. A head - // that renders nothing would leave nothing to summarize, so the - // smallest head kept is the smallest that still renders. - for estimateTokens(view(head))+prefix > target { - h, t := splitForCompaction(msgs, keep+1) - if len(h) == 0 || len(renderMessages(view(h))) == 0 { + budget := target - tokensOf(prefix) + end, acc := 0, 0 + for e := start + 1; e <= len(head); e++ { + acc += sizes[e-1] + if !boundary(e) { + continue + } + if end == 0 || acc/4 <= budget { + end = e + } + if acc/4 > budget { break } - keep++ - head, tail = h, t } + pieces := renderMessages(viewed[start:end]) + if len(pieces) == 0 { + start = end // nothing in it for a summary to cover + continue + } + reply, sent, serr := s.summarizeWith(ctx, prefix, pieces, max(budget, 1)) + if serr == nil { + running, start, failures = reply, end, 0 + chunks++ + continue + } + failures++ + if failures >= maxSummaryAttempts || !canRecoverFromOverflow(ctx, serr) { + return "", 0, serr + } + if target = summaryRetryTarget(serr, sent); target <= 0 { + return "", 0, serr + } + } + if chunks == 0 { + return "", 0, fmt.Errorf("compaction produced an empty summary") } + return running, start, nil } // spillCompactionHead saves the rendered head a summary stands for as an @@ -747,7 +847,13 @@ func (s *Session) spillCompactionHead(pieces []summaryPiece) string { // summaryOutputTokens. Ask would leave the reservation to the client, which // reserves the model's whole output cap. func (s *Session) summarize(ctx context.Context, pieces []summaryPiece, limit int) (summary string, sentTokens int, err error) { - prompt := summaryPrefix + fitSummaryInput(pieces, limit) + return s.summarizeWith(ctx, summaryPrefix, pieces, limit) +} + +// summarizeWith is summarize with the instruction prefix the rendered pieces +// follow in the prompt. +func (s *Session) summarizeWith(ctx context.Context, prefix string, pieces []summaryPiece, limit int) (summary string, sentTokens int, err error) { + prompt := prefix + fitSummaryInput(pieces, limit) sentTokens = tokensOf(prompt) llm, _ := s.currentLLM() reply, usage, aerr := llm.CreateChatCompletion(ctx, openai.ChatCompletionRequest{ diff --git a/chat/compact_fit_test.go b/chat/compact_fit_test.go index b8695c9..212a025 100644 --- a/chat/compact_fit_test.go +++ b/chat/compact_fit_test.go @@ -107,7 +107,7 @@ func TestCompactHistoryFitsTheSummaryPromptInTheBudget(t *testing.T) { func TestCompactHistoryRetriesSmallerWhenTheSummaryOverflows(t *testing.T) { // The backend tokenizes denser than byte/4, so a prompt the estimate says // fits can still overflow. The backend's own figures say by how much, and - // the retry moves the boundary so the head it summarizes fits. + // the head is then summarized in chunks that fit. frag := []openai.ChatCompletionMessage{{Role: "user", Content: "goal"}} for i := range 6 { frag = append(frag, toolTurn(fmt.Sprintf("c%d", i), fmt.Sprintf("f%d.go", i), strings.Repeat("x", 8000))...) @@ -124,13 +124,13 @@ func TestCompactHistoryRetriesSmallerWhenTheSummaryOverflows(t *testing.T) { if _, _, err := s.CompactHistory(); err != nil { t.Fatalf("CompactHistory: %v", err) } - if n := len(llm.prompts); n < 2 || n > maxSummaryAttempts { - t.Fatalf("summary calls = %d, want an overflow and then smaller retries", n) + if n := len(llm.prompts); n < 2 || n > 1+maxRollingChunks { + t.Fatalf("summary calls = %d, want an overflow and then smaller chunks", n) } if got := tokensOf(llm.prompts[len(llm.prompts)-1]); got > 1500 { t.Fatalf("the accepted prompt is ~%d tokens, over the backend's 1500", got) } - // The messages the head gave up are kept verbatim instead. + // The messages no chunk covered are kept verbatim instead. checkVerbatimTail(t, orig, s.fragment.Messages) } diff --git a/chat/summary_rolling_test.go b/chat/summary_rolling_test.go new file mode 100644 index 0000000..20afb69 --- /dev/null +++ b/chat/summary_rolling_test.go @@ -0,0 +1,282 @@ +package chat + +import ( + "context" + "fmt" + "reflect" + "strings" + "testing" + + "github.com/mudler/cogito" + "github.com/mudler/nib/mcp" + "github.com/mudler/nib/types" + openai "github.com/sashabaranov/go-openai" +) + +// rollingLLM is a backend that checks the output reservation the way vLLM +// does (prompt + max_tokens against limit, prompt counted byte/4), and +// numbers its summaries SUMMARY-1, SUMMARY-2, ... so a test can see which +// summary a later prompt carries. +type rollingLLM struct { + limit int + // overflowCall, when positive, rejects that call (1-based) as + // overflowing by overBy tokens, whatever its size. + overflowCall, overBy int + reqs []openai.ChatCompletionRequest + served int +} + +func (f *rollingLLM) Ask(ctx context.Context, fr cogito.Fragment) (cogito.Fragment, error) { + return cogito.Fragment{}, fmt.Errorf("the summary must be sent with CreateChatCompletion, not Ask") +} + +func (f *rollingLLM) CreateChatCompletion(ctx context.Context, req openai.ChatCompletionRequest) (cogito.LLMReply, cogito.LLMUsage, error) { + f.reqs = append(f.reqs, req) + usage := cogito.LLMUsage{PromptTokens: 10, CompletionTokens: 1, TotalTokens: 11} + p := tokensOf(req.Messages[len(req.Messages)-1].Content) + if len(f.reqs) == f.overflowCall { + return cogito.LLMReply{}, usage, vllmOverflow(p+req.MaxTokens-f.overBy, req.MaxTokens, p) + } + if f.limit > 0 && p+req.MaxTokens > f.limit { + return cogito.LLMReply{}, usage, vllmOverflow(f.limit, req.MaxTokens, p) + } + f.served++ + return replyWith(fmt.Sprintf("SUMMARY-%d", f.served)), usage, nil +} + +func (f *rollingLLM) prompt(i int) string { return f.reqs[i].Messages[0].Content } + +// shortHistory is a goal, n tool steps of ~500 tokens each, and a short +// closing exchange. +func shortHistory(n int) []openai.ChatCompletionMessage { + msgs := []openai.ChatCompletionMessage{{Role: "user", Content: "goal: fix the parser"}} + for i := range n { + msgs = append(msgs, toolTurn(fmt.Sprintf("c%d", i), fmt.Sprintf("f%d.go", i), fmt.Sprintf("BODY%02d", i)+strings.Repeat("x", 2000))...) + } + return append(msgs, + openai.ChatCompletionMessage{Role: "user", Content: "u2"}, + openai.ChatCompletionMessage{Role: "assistant", Content: "a2"}, + ) +} + +// rollingSession is a session that believes the window is huge, while the +// backend serves 8192 and every summary reserves 4000 of it. +func rollingSession(llm cogito.LLM, frag []openai.ChatCompletionMessage) *Session { + s := newCompactTestSession(llm, 2, frag, frag) + s.compaction.MaxContextTokens = 1 << 20 + s.compaction.SummaryMaxTokens = 4000 + s.artifacts = mcp.NewArtifactStore() + return s +} + +// chunkPrompts are the prompts of the rolling chunks that the backend +// served, in order: the first, whole-head attempt is left out. +func servedChunkPrompts(llm *rollingLLM) []string { + var out []string + for i := 1; i < len(llm.reqs); i++ { + if i+1 == llm.overflowCall { + continue + } + out = append(out, llm.prompt(i)) + } + return out +} + +func TestRollingSummaryCoversTheWholeHeadInChunks(t *testing.T) { + frag := shortHistory(24) // ~13k tokens against a ~3.7k fitting size + orig := cloneMessages(frag) + llm := &rollingLLM{limit: 8192} + s := rollingSession(llm, frag) + + if _, _, err := s.CompactHistory(); err != nil { + t.Fatalf("CompactHistory: %v", err) + } + chunks := len(llm.reqs) - 1 + if chunks < 3 || chunks > maxRollingChunks { + t.Fatalf("chunks = %d, want the head covered in several chunks", chunks) + } + got := s.fragment.Messages + if len(got) != 3 || got[0].Content == "" || !strings.Contains(got[0].Content, fmt.Sprintf("SUMMARY-%d", chunks)) { + t.Fatalf("want [final summary] + the original tail, got %d messages", len(got)) + } + if !reflect.DeepEqual(got[1:], orig[len(orig)-2:]) { + t.Fatal("the tail is not the original KeepRecent tail") + } + real := ContextBudget(types.CompactionConfig{}, 8192) + if n := estimateTokens(got); n >= real { + t.Fatalf("compacted history is %d tokens, want below the real budget %d", n, real) + } + for i, r := range llm.reqs { + if r.MaxTokens != 4000 { + t.Fatalf("request %d MaxTokens = %d, want 4000", i, r.MaxTokens) + } + } + // Every message of the head reached a chunk exactly once. + all := strings.Join(servedChunkPrompts(llm), "") + for i := range 24 { + if c := strings.Count(all, fmt.Sprintf("BODY%02d", i)); c != 1 { + t.Fatalf("BODY%02d appears in %d chunks, want 1", i, c) + } + } + if !strings.Contains(llm.prompt(1), "goal: fix the parser") { + t.Fatal("the first chunk lost the user's goal") + } +} + +func TestRollingSummaryChunksCarryTheRunningSummaryFirst(t *testing.T) { + llm := &rollingLLM{limit: 8192} + s := rollingSession(llm, shortHistory(24)) + if _, _, err := s.CompactHistory(); err != nil { + t.Fatalf("CompactHistory: %v", err) + } + prompts := servedChunkPrompts(llm) + if len(prompts) < 2 { + t.Fatalf("chunks = %d, want several", len(prompts)) + } + if !strings.HasPrefix(prompts[0], summaryPrefix) { + t.Fatal("the first chunk must use the normal compaction instruction") + } + for k := 1; k < len(prompts); k++ { + p := prompts[k] + if !strings.HasPrefix(p, rollingInstruction) { + t.Fatalf("chunk %d does not use rollingInstruction", k+1) + } + prev := fmt.Sprintf("SUMMARY-%d", k) + at := strings.Index(p, prev) + body := strings.Index(p, "BODY") + if at < 0 || body < 0 || at > body { + t.Fatalf("chunk %d does not carry %s before its messages", k+1, prev) + } + } +} + +func TestRollingChunkBoundariesKeepToolCallsWithResults(t *testing.T) { + llm := &rollingLLM{limit: 8192} + s := rollingSession(llm, shortHistory(24)) + if _, _, err := s.CompactHistory(); err != nil { + t.Fatalf("CompactHistory: %v", err) + } + for k, p := range servedChunkPrompts(llm) { + calls := strings.Count(p, "[tool call read(") + results := strings.Count(p, "\ntool: BODY") + if calls != results { + t.Fatalf("chunk %d has %d tool calls and %d results", k+1, calls, results) + } + } +} + +func TestRollingChunkThatOverflowsIsShrunk(t *testing.T) { + frag := shortHistory(24) + orig := cloneMessages(frag) + // The first chunk overflows on the backend anyway, by 400 tokens. + llm := &rollingLLM{limit: 8192, overflowCall: 2, overBy: 400} + s := rollingSession(llm, frag) + + if _, _, err := s.CompactHistory(); err != nil { + t.Fatalf("CompactHistory: %v", err) + } + if len(llm.reqs) < 4 { + t.Fatalf("requests = %d, want the overflowing chunk retried", len(llm.reqs)) + } + if tokensOf(llm.prompt(2)) >= tokensOf(llm.prompt(1)) { + t.Fatalf("the retried chunk is %d tokens, want smaller than the %d that overflowed", tokensOf(llm.prompt(2)), tokensOf(llm.prompt(1))) + } + got := s.fragment.Messages + if len(got) != 3 || !reflect.DeepEqual(got[1:], orig[len(orig)-2:]) { + t.Fatalf("want [summary] + original tail, got %d messages", len(got)) + } + all := strings.Join(servedChunkPrompts(llm), "") + for i := range 24 { + if c := strings.Count(all, fmt.Sprintf("BODY%02d", i)); c != 1 { + t.Fatalf("BODY%02d appears in %d served chunks, want 1: the rest must move to the next chunk", i, c) + } + } +} + +func TestRollingSummaryStopsAfterMaxRollingChunks(t *testing.T) { + frag := longHistory(20) // one ~2000-token step per chunk + orig := cloneMessages(frag) + llm := &rollingLLM{limit: 8192} + s := rollingSession(llm, frag) + + if _, _, err := s.CompactHistory(); err != nil { + t.Fatalf("CompactHistory: %v", err) + } + if n := len(llm.reqs); n != 1+maxRollingChunks { + t.Fatalf("requests = %d, want the first attempt and %d chunks", n, maxRollingChunks) + } + start := checkVerbatimTail(t, orig, s.fragment.Messages) + if len(orig)-start <= 2 { + t.Fatal("want the messages no chunk covered kept verbatim") + } + // The final summary stands for exactly the covered messages. + if s.artifacts.Count() != 1 { + t.Fatalf("artifacts = %d, want 1", s.artifacts.Count()) + } + var want strings.Builder + for _, p := range renderMessages(orig[:start]) { + want.WriteString(p.text) + } + if got := s.artifacts.Get(1).Content; got != want.String() { + t.Fatalf("artifact is %d bytes, want exactly the %d covered bytes", len(got), want.Len()) + } + all := strings.Join(servedChunkPrompts(llm), "") + if !strings.Contains(all, orig[start-1].Content) { + t.Fatal("the last covered message is not in any chunk") + } + if strings.Contains(all, orig[start+1].Content) { + t.Fatal("a message kept verbatim was also summarized") + } +} + +func TestRollingSummaryCountsEveryChunkCall(t *testing.T) { + llm := &rollingLLM{limit: 8192, overflowCall: 3, overBy: 500} + s := rollingSession(llm, shortHistory(24)) + if _, _, err := s.CompactHistory(); err != nil { + t.Fatalf("CompactHistory: %v", err) + } + if got, want := s.Usage().TotalTokens, 11*len(llm.reqs); got != want { + t.Fatalf("TotalTokens = %d, want %d for %d calls", got, want, len(llm.reqs)) + } +} + +func TestRollingArtifactHoldsExactlyTheCoveredHead(t *testing.T) { + frag := shortHistory(24) + orig := cloneMessages(frag) + llm := &rollingLLM{limit: 8192} + s := rollingSession(llm, frag) + if _, _, err := s.CompactHistory(); err != nil { + t.Fatalf("CompactHistory: %v", err) + } + var want strings.Builder + for _, p := range renderMessages(orig[:len(orig)-2]) { + want.WriteString(p.text) + } + if s.artifacts.Count() != 1 || s.artifacts.Get(1).Content != want.String() { + t.Fatal("the artifact must hold exactly the whole covered head") + } +} + +func TestRollingSummaryCutsAMessageLargerThanTheTarget(t *testing.T) { + frag := shortHistory(4) + // One result alone is far larger than any prompt the backend accepts. + frag[4].Content = "HUGE" + strings.Repeat("y", 40000) + orig := cloneMessages(frag) + llm := &rollingLLM{limit: 8192} + s := rollingSession(llm, frag) + + if _, _, err := s.CompactHistory(); err != nil { + t.Fatalf("CompactHistory: %v", err) + } + got := s.fragment.Messages + if len(got) != 3 || !reflect.DeepEqual(got[1:], orig[len(orig)-2:]) { + t.Fatalf("want the whole head covered, got %d messages", len(got)) + } + all := strings.Join(servedChunkPrompts(llm), "") + if !strings.Contains(all, "tool: HUGE") || !strings.Contains(all, "omitted to fit the summary") { + t.Fatal("the over-large message must be cut down, not skipped") + } + if !strings.Contains(s.artifacts.Get(1).Content, frag[4].Content) { + t.Fatal("the artifact must hold the over-large message in full") + } +} diff --git a/chat/summary_shift_test.go b/chat/summary_shift_test.go index 109ea1c..6c0ddc3 100644 --- a/chat/summary_shift_test.go +++ b/chat/summary_shift_test.go @@ -110,8 +110,8 @@ func TestSummaryShiftConvergesWithAFixedOutputReservation(t *testing.T) { if _, _, err := s.CompactHistory(); err != nil { t.Fatalf("CompactHistory: %v", err) } - if n := len(llm.reqs); n < 2 || n > maxSummaryAttempts { - t.Fatalf("summary calls = %d, want an overflow then a fitting retry", n) + if n := len(llm.reqs); n < 2 || n > 1+maxRollingChunks { + t.Fatalf("summary calls = %d, want an overflow then fitting chunks", n) } for i, r := range llm.reqs { if r.MaxTokens != 4000 { @@ -122,8 +122,8 @@ func TestSummaryShiftConvergesWithAFixedOutputReservation(t *testing.T) { if len(orig)-start <= 2 { t.Fatal("the boundary did not move: the kept tail is still KeepRecent long") } - if !strings.Contains(llm.reqs[len(llm.reqs)-1].Messages[0].Content, "goal: fix the parser") { - t.Fatal("the shifted head lost the user's goal") + if !strings.Contains(llm.reqs[1].Messages[0].Content, "goal: fix the parser") { + t.Fatal("the first chunk lost the user's goal") } } @@ -257,7 +257,7 @@ func TestMidTurnCompactionShiftsTheBoundaryWhenTheSummaryOverflows(t *testing.T) } } -func TestMidTurnCompactionSkipsAHeadOfOnlyThePreviousSummary(t *testing.T) { +func TestMidTurnCompactionRollsThePreviousSummaryIntoTheNewOne(t *testing.T) { raw := longHistory(5) // The request already carries a summary in place of the first 3 messages. prev := openai.ChatCompletionMessage{Role: "user", Content: "PREV"} @@ -268,14 +268,20 @@ func TestMidTurnCompactionSkipsAHeadOfOnlyThePreviousSummary(t *testing.T) { s.compaction.SummaryMaxTokens = 1000 c := s.newTurnCompactor(context.Background()) + // No step fits beside the previous summary, so the chunks cover it and + // then one cut-down step at a time, and the new summary replaces the + // whole head: the previous summary is merged, not kept beside it. got := c.compact(raw, out, 3, 1, 0) - if len(llm.reqs) < 2 { - t.Fatalf("summary calls = %d, want the loop to shrink the head to the previous summary", len(llm.reqs)) + if len(llm.reqs) < 3 { + t.Fatalf("summary calls = %d, want an overflow and then several chunks", len(llm.reqs)) + } + if !strings.Contains(llm.reqs[1].Messages[0].Content, "PREV") { + t.Fatal("the first chunk does not start with the previous summary") } - if !reflect.DeepEqual(got, out) { - t.Fatal("a head of only the previous summary should leave the request unchanged") + if len(got) != 3 || !reflect.DeepEqual(got[1:], out[len(out)-2:]) { + t.Fatalf("want [summary] + the KeepRecent tail, got %d messages", len(got)) } - if c.covered != 0 { - t.Fatalf("covered = %d, want 0", c.covered) + if want := len(raw) - 2; c.covered != want { + t.Fatalf("covered = %d, want %d", c.covered, want) } } From 42ebe43a328c7a71fa275cfe7326382bf767a879 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Fri, 25 Sep 2026 20:50:57 +0000 Subject: [PATCH 04/16] feat(chat): validate a compaction result before it is sent Add validateCompaction, which rejects an empty result, a result that grew, a no-op, a result over the budget, and broken tool-call pairing. The sentinel errors are exported so overflow recovery can switch on them. Co-Authored-By: Claude Opus 5.5 --- chat/compact_validate.go | 69 +++++++++++++++++++ chat/compact_validate_test.go | 124 ++++++++++++++++++++++++++++++++++ 2 files changed, 193 insertions(+) create mode 100644 chat/compact_validate.go create mode 100644 chat/compact_validate_test.go diff --git a/chat/compact_validate.go b/chat/compact_validate.go new file mode 100644 index 0000000..67f83f5 --- /dev/null +++ b/chat/compact_validate.go @@ -0,0 +1,69 @@ +package chat + +import ( + "errors" + "fmt" + + openai "github.com/sashabaranov/go-openai" +) + +// Sentinel errors returned by validateCompaction. They are exported so the +// overflow-recovery caller can switch on them with errors.Is. +var ( + ErrCompactionNoOp = errors.New("compaction was a no-op") + ErrCompactionGrew = errors.New("compaction result is larger than the original") + ErrCompactionOverBudget = errors.New("compaction result exceeds the context budget") + ErrCompactionEmpty = errors.New("compaction result is empty") +) + +// validateCompaction gates a compaction result before it is sent. It returns +// the first failing check: empty result, growth, no-op, over budget, then +// broken tool-call pairing. Growth is checked before no-op, because a result +// that grew is never equal in size and the growth check must be reachable. +// budget is decided by the caller (the context budget, minus any +// non-compactable floor it knows about). +func validateCompaction(before, after []openai.ChatCompletionMessage, budget int) error { + if len(after) == 0 { + return ErrCompactionEmpty + } + was, now := estimateTokens(before), estimateTokens(after) + if now > was { + return fmt.Errorf("%w: %d > %d tokens", ErrCompactionGrew, now, was) + } + if now == was { + return fmt.Errorf("%w: still %d tokens", ErrCompactionNoOp, now) + } + if now > budget { + return fmt.Errorf("%w: %d > %d tokens", ErrCompactionOverBudget, now, budget) + } + return validateToolPairing(after) +} + +// validateToolPairing reports a tool result whose ToolCallID has no earlier +// assistant tool call, and an assistant tool call with no tool result after +// it. Either one makes the backend reject the request. +func validateToolPairing(msgs []openai.ChatCompletionMessage) error { + pending := map[string]bool{} // call ID -> seen a result yet + for i, m := range msgs { + for _, tc := range m.ToolCalls { + pending[tc.ID] = false + } + if m.Role == openai.ChatMessageRoleTool { + answered, ok := pending[m.ToolCallID] + if !ok { + return fmt.Errorf("tool result %q at message %d has no matching tool call", m.ToolCallID, i) + } + if !answered { + pending[m.ToolCallID] = true + } + } + } + for i, m := range msgs { + for _, tc := range m.ToolCalls { + if !pending[tc.ID] { + return fmt.Errorf("tool call %q at message %d has no tool result", tc.ID, i) + } + } + } + return nil +} diff --git a/chat/compact_validate_test.go b/chat/compact_validate_test.go new file mode 100644 index 0000000..c8ae330 --- /dev/null +++ b/chat/compact_validate_test.go @@ -0,0 +1,124 @@ +package chat + +import ( + "errors" + "strings" + "testing" + + openai "github.com/sashabaranov/go-openai" +) + +func vmsg(role string, n int) openai.ChatCompletionMessage { + return openai.ChatCompletionMessage{Role: role, Content: strings.Repeat("x", n*4)} +} + +func vcall(id string) openai.ChatCompletionMessage { + return openai.ChatCompletionMessage{ + Role: openai.ChatMessageRoleAssistant, + ToolCalls: []openai.ToolCall{{ + ID: id, Type: openai.ToolTypeFunction, + Function: openai.FunctionCall{Name: "read", Arguments: "{}"}, + }}, + } +} + +func vresult(id string) openai.ChatCompletionMessage { + return openai.ChatCompletionMessage{Role: openai.ChatMessageRoleTool, ToolCallID: id, Content: "ok"} +} + +func TestValidateCompactionAcceptsAValidResult(t *testing.T) { + before := []openai.ChatCompletionMessage{vmsg("user", 100), vcall("a"), vresult("a"), vmsg("assistant", 100)} + after := []openai.ChatCompletionMessage{vmsg("user", 10), vcall("a"), vresult("a")} + if err := validateCompaction(before, after, 1000); err != nil { + t.Fatalf("valid compaction rejected: %v", err) + } +} + +func TestValidateCompactionRejectsEmpty(t *testing.T) { + before := []openai.ChatCompletionMessage{vmsg("user", 100)} + if err := validateCompaction(before, nil, 1000); !errors.Is(err, ErrCompactionEmpty) { + t.Fatalf("err = %v, want ErrCompactionEmpty", err) + } +} + +func TestValidateCompactionRejectsGrowth(t *testing.T) { + before := []openai.ChatCompletionMessage{vmsg("user", 10)} + after := []openai.ChatCompletionMessage{vmsg("user", 20)} + if err := validateCompaction(before, after, 1000); !errors.Is(err, ErrCompactionGrew) { + t.Fatalf("err = %v, want ErrCompactionGrew", err) + } +} + +func TestValidateCompactionRejectsNoOp(t *testing.T) { + before := []openai.ChatCompletionMessage{vmsg("user", 50), vmsg("assistant", 50)} + after := []openai.ChatCompletionMessage{vmsg("user", 100)} + if err := validateCompaction(before, after, 1000); !errors.Is(err, ErrCompactionNoOp) { + t.Fatalf("err = %v, want ErrCompactionNoOp", err) + } +} + +func TestValidateCompactionRejectsNoOpFromSplit(t *testing.T) { + // A split whose head comes back unchanged reassembles into the original + // history, the same size: the known no-op cause. + before := []openai.ChatCompletionMessage{vmsg("user", 50), vmsg("assistant", 50)} + head, tail := splitForCompaction(before, 1) + after := append(append([]openai.ChatCompletionMessage{}, head...), tail...) + if err := validateCompaction(before, after, 1000); !errors.Is(err, ErrCompactionNoOp) { + t.Fatalf("err = %v, want ErrCompactionNoOp", err) + } +} + +func TestValidateCompactionRejectsOverBudget(t *testing.T) { + before := []openai.ChatCompletionMessage{vmsg("user", 500)} + after := []openai.ChatCompletionMessage{vmsg("user", 200)} + if err := validateCompaction(before, after, 100); !errors.Is(err, ErrCompactionOverBudget) { + t.Fatalf("err = %v, want ErrCompactionOverBudget", err) + } +} + +func TestValidateCompactionGrowthBeatsNoOpAndBudget(t *testing.T) { + before := []openai.ChatCompletionMessage{vmsg("user", 10)} + after := []openai.ChatCompletionMessage{vmsg("user", 500)} + if err := validateCompaction(before, after, 100); !errors.Is(err, ErrCompactionGrew) { + t.Fatalf("err = %v, want ErrCompactionGrew", err) + } +} + +func TestValidateCompactionRejectsOrphanToolResult(t *testing.T) { + before := []openai.ChatCompletionMessage{vmsg("user", 100), vcall("a"), vresult("a")} + after := []openai.ChatCompletionMessage{vmsg("user", 10), vresult("a")} + err := validateCompaction(before, after, 1000) + if err == nil || errors.Is(err, ErrCompactionEmpty) || errors.Is(err, ErrCompactionGrew) || + errors.Is(err, ErrCompactionNoOp) || errors.Is(err, ErrCompactionOverBudget) { + t.Fatalf("err = %v, want a tool-pairing error", err) + } +} + +func TestValidateCompactionRejectsOrphanToolCall(t *testing.T) { + before := []openai.ChatCompletionMessage{vmsg("user", 100), vcall("a"), vresult("a")} + after := []openai.ChatCompletionMessage{vmsg("user", 10), vcall("a")} + if err := validateCompaction(before, after, 1000); err == nil { + t.Fatal("orphan tool call accepted") + } +} + +func TestValidateCompactionToolPairing(t *testing.T) { + cases := []struct { + name string + msgs []openai.ChatCompletionMessage + ok bool + }{ + {"paired", []openai.ChatCompletionMessage{vcall("a"), vresult("a")}, true}, + {"no tools", []openai.ChatCompletionMessage{vmsg("user", 1)}, true}, + {"orphan result", []openai.ChatCompletionMessage{vresult("a")}, false}, + {"result before call", []openai.ChatCompletionMessage{vresult("a"), vcall("a")}, false}, + {"orphan call", []openai.ChatCompletionMessage{vcall("a"), vmsg("user", 1)}, false}, + {"mismatched id", []openai.ChatCompletionMessage{vcall("a"), vresult("b")}, false}, + } + for _, c := range cases { + err := validateToolPairing(c.msgs) + if (err == nil) != c.ok { + t.Errorf("%s: err = %v, want ok=%v", c.name, err, c.ok) + } + } +} From 3d407ff24c32f61dd6267677dba7119147a3c52f Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Fri, 25 Sep 2026 20:53:59 +0000 Subject: [PATCH 05/16] feat(chat): estimate the context the tool schemas and system prompt take SchemaBudget measures the fixed per-request overhead (tool schemas plus system prompt, bytes/4) against the context budget and window, and logs a warning when it passes 60% of the budget or 90% of the window. toolOptions records the tool definitions it registers so the session can measure them; cogito keeps its own list private. Co-Authored-By: Claude Opus 5.5 --- chat/schema_budget.go | 100 +++++++++++++++++++++++++++++++++++++ chat/schema_budget_test.go | 90 +++++++++++++++++++++++++++++++++ chat/session.go | 44 +++++++++------- 3 files changed, 215 insertions(+), 19 deletions(-) create mode 100644 chat/schema_budget.go create mode 100644 chat/schema_budget_test.go diff --git a/chat/schema_budget.go b/chat/schema_budget.go new file mode 100644 index 0000000..73ad850 --- /dev/null +++ b/chat/schema_budget.go @@ -0,0 +1,100 @@ +package chat + +import ( + "encoding/json" + + "github.com/mudler/cogito" + "github.com/mudler/xlog" +) + +// schemaWarnFraction is the share of the context budget the fixed prompt +// overhead (tool schemas + system prompt) may take before SchemaBudget warns. +const schemaWarnFraction = 0.60 + +// schemaErrorFraction is the share of the whole context window the fixed prompt +// overhead may take before SchemaBudget reports an error: past it, compaction +// cannot free enough room for a useful turn. +const schemaErrorFraction = 0.90 + +// SchemaBudget is a read-only observation of how much of the context the fixed +// per-request overhead takes. Floor is the estimated token count of the tool +// schemas plus the system prompt — the floor no compaction can go below. +type SchemaBudget struct { + Floor int + WarnThreshold float64 + ErrorThreshold float64 + Warn bool + Error bool +} + +// estimateSchemaTokens estimates the tokens the tool schemas and the system +// prompt take in every request, with the same byte/4 heuristic estimateTokens +// uses for messages. +func estimateSchemaTokens(tools []cogito.ToolDefinitionInterface, systemPrompt string) int { + n := len(systemPrompt) + for _, td := range tools { + if td == nil { + continue + } + if b, err := json.Marshal(td.Tool()); err == nil { + n += len(b) + } + } + return n / 4 +} + +// withTool registers td as a cogito tool and records it, so SchemaBudget can +// measure the tool definitions the session advertises. cogito keeps its tool +// list private, so this is the only place the session can see them. +func (s *Session) withTool(td cogito.ToolDefinitionInterface) cogito.Option { + s.schemaToolsMu.Lock() + s.schemaTools = append(s.schemaTools, td) + s.schemaToolsMu.Unlock() + return cogito.WithTools(td) +} + +// resetSchemaTools clears the recorded tool definitions; toolOptions calls it +// before it registers the tools again. +func (s *Session) resetSchemaTools() { + s.schemaToolsMu.Lock() + s.schemaTools = nil + s.schemaToolsMu.Unlock() +} + +// SchemaBudget estimates the fixed prompt overhead and compares it with the +// context budget and window. It does not change behavior; it only logs a +// warning when the overhead is large, for a caller to show the user. +// +// The tool definitions are the ones the last toolOptions call registered +// through withTool. MCP tools and the agent-spawning tools are added inside +// cogito and are not counted. +func (s *Session) SchemaBudget() SchemaBudget { + s.schemaToolsMu.Lock() + tools := append([]cogito.ToolDefinitionInterface(nil), s.schemaTools...) + s.schemaToolsMu.Unlock() + + floor := estimateSchemaTokens(tools, s.systemPrompt) + window := s.contextWindow() + budget := ContextBudget(s.compactionConfig(), window) + + sb := SchemaBudget{ + Floor: floor, + WarnThreshold: schemaWarnFraction, + ErrorThreshold: schemaErrorFraction, + } + if window <= 0 { + // Unknown window: nothing to compare against. + return sb + } + sb.Warn = floor > int(float64(budget)*schemaWarnFraction) + sb.Error = floor > int(float64(window)*schemaErrorFraction) + switch { + case sb.Error: + xlog.Warn("tool schemas and system prompt fill almost the whole context window; compaction cannot make room", + "schema_tokens", floor, "window", window, "budget", budget) + case sb.Warn: + xlog.Warn("tool schemas and system prompt take a large share of the context budget", + "schema_tokens", floor, "window", window, "budget", budget) + } + return sb +} diff --git a/chat/schema_budget_test.go b/chat/schema_budget_test.go new file mode 100644 index 0000000..47dec47 --- /dev/null +++ b/chat/schema_budget_test.go @@ -0,0 +1,90 @@ +package chat + +import ( + "encoding/json" + "strings" + "testing" + + "github.com/mudler/cogito" + openai "github.com/sashabaranov/go-openai" +) + +// schemaTestTool is a ToolDefinitionInterface whose schema size the test controls +// through the description length. +type schemaTestTool struct{ desc string } + +func (t schemaTestTool) Tool() openai.Tool { + return openai.Tool{Type: openai.ToolTypeFunction, Function: &openai.FunctionDefinition{Name: "t", Description: t.desc}} +} + +func (t schemaTestTool) Execute(map[string]any) (string, any, error) { return "", nil, nil } + +func TestSchemaBudgetEstimateMatchesByteQuarter(t *testing.T) { + tools := []cogito.ToolDefinitionInterface{schemaTestTool{desc: "one"}, schemaTestTool{desc: strings.Repeat("x", 100)}} + sys := strings.Repeat("s", 40) + + want := len(sys) + for _, td := range tools { + b, err := json.Marshal(td.Tool()) + if err != nil { + t.Fatal(err) + } + want += len(b) + } + want /= 4 + + if got := estimateSchemaTokens(tools, sys); got != want { + t.Fatalf("estimateSchemaTokens = %d, want %d", got, want) + } + if got := estimateSchemaTokens(nil, ""); got != 0 { + t.Fatalf("empty estimate = %d, want 0", got) + } +} + +func TestSchemaBudgetThresholds(t *testing.T) { + // Window 1000 → reserve min(4096, 500) = 500 → budget 500. + // Warn when floor > 0.60*500 = 300; Error when floor > 0.90*1000 = 900. + cases := []struct { + name string + sysBytes int + warn, error bool + }{ + {"small", 400, false, false}, // floor 100 + {"warn", 1400, true, false}, // floor 350 + {"error", 3800, true, true}, // floor 950 + {"at warn edge", 1200, false, false}, // floor 300, not strictly greater + } + for _, c := range cases { + t.Run(c.name, func(t *testing.T) { + s := newCompactTestSession(&fakeSummaryLLM{}, 2, nil, nil) + s.systemPrompt = strings.Repeat("p", c.sysBytes) + got := s.SchemaBudget() + if got.Floor != c.sysBytes/4 { + t.Fatalf("Floor = %d, want %d", got.Floor, c.sysBytes/4) + } + if got.Warn != c.warn || got.Error != c.error { + t.Fatalf("Warn/Error = %v/%v, want %v/%v", got.Warn, got.Error, c.warn, c.error) + } + if got.WarnThreshold != schemaWarnFraction || got.ErrorThreshold != schemaErrorFraction { + t.Fatalf("thresholds = %v/%v", got.WarnThreshold, got.ErrorThreshold) + } + }) + } +} + +func TestSchemaBudgetCountsRecordedTools(t *testing.T) { + s := newCompactTestSession(&fakeSummaryLLM{}, 2, nil, nil) + tool := schemaTestTool{desc: strings.Repeat("d", 2000)} + opt := s.withTool(tool) + if opt == nil { + t.Fatal("withTool returned nil option") + } + want := estimateSchemaTokens([]cogito.ToolDefinitionInterface{tool}, "") + if got := s.SchemaBudget().Floor; got != want { + t.Fatalf("Floor = %d, want %d (recorded tool not counted)", got, want) + } + s.resetSchemaTools() + if got := s.SchemaBudget().Floor; got != 0 { + t.Fatalf("Floor after reset = %d, want 0", got) + } +} diff --git a/chat/session.go b/chat/session.go index 6749051..b57176e 100644 --- a/chat/session.go +++ b/chat/session.go @@ -124,6 +124,11 @@ type Session struct { // sub-agents). May be nil (e.g. headless CLI without a job registry). shellJobs *wizmcp.ShellJobs + // schemaTools records the tool definitions toolOptions registers, for + // SchemaBudget. Guarded by schemaToolsMu. See schema_budget.go. + schemaTools []cogito.ToolDefinitionInterface + schemaToolsMu sync.Mutex + agentManager *cogito.AgentManager agentDefs []cogito.AgentDefinition agentModels map[string]bool // models configured per agent type (for the LLM-model guard) @@ -1361,6 +1366,7 @@ func buildUserFragment(f cogito.Fragment, text string, parts []ContentPart) cogi // the sub-agents a turn spawns resolve against the same model the turn itself // is using, even if SetModel lands halfway through. func (s *Session) toolOptions(turnCtx context.Context, goal, mainModel string) []cogito.Option { + s.resetSchemaTools() opts := []cogito.Option{ cogito.WithMCPs(s.allClients()...), // Disable cogito's sink-state "reply" tool so ExecuteTools is the whole @@ -1385,7 +1391,7 @@ func (s *Session) toolOptions(turnCtx context.Context, goal, mainModel string) [ opts = append(opts, cogito.EnableAgentSpawning) } if s.toolEnabled("ask_user") && !s.AutoApprove() { - opts = append(opts, cogito.WithTools(askUserToolDefinition(func(req AskRequest) string { + opts = append(opts, s.withTool(askUserToolDefinition(func(req AskRequest) string { if s.callbacks.OnAskUser != nil { return s.callbacks.OnAskUser(req) } @@ -1393,10 +1399,10 @@ func (s *Session) toolOptions(turnCtx context.Context, goal, mainModel string) [ }))) } if s.toolEnabled("agent_logs") { - opts = append(opts, cogito.WithTools(agentLogsToolDefinition(s.AgentLog))) + opts = append(opts, s.withTool(agentLogsToolDefinition(s.AgentLog))) } if s.toolEnabled("schedule_wakeup") { - opts = append(opts, cogito.WithTools(scheduleWakeupToolDefinition(func(req WakeupRequest) string { + opts = append(opts, s.withTool(scheduleWakeupToolDefinition(func(req WakeupRequest) string { if s.callbacks.OnScheduleWakeup != nil { return s.callbacks.OnScheduleWakeup(req) } @@ -1406,7 +1412,7 @@ func (s *Session) toolOptions(turnCtx context.Context, goal, mainModel string) [ }))) } if s.toolEnabled("cron") { - opts = append(opts, cogito.WithTools(cronToolDefinition(func(req CronRequest) string { + opts = append(opts, s.withTool(cronToolDefinition(func(req CronRequest) string { if s.callbacks.OnCronCreate != nil { return s.callbacks.OnCronCreate(req) } @@ -1414,7 +1420,7 @@ func (s *Session) toolOptions(turnCtx context.Context, goal, mainModel string) [ }))) } if s.toolEnabled("cron_list") { - opts = append(opts, cogito.WithTools(cronListToolDefinition(func() string { + opts = append(opts, s.withTool(cronListToolDefinition(func() string { if s.callbacks.OnCronList != nil { return s.callbacks.OnCronList() } @@ -1422,7 +1428,7 @@ func (s *Session) toolOptions(turnCtx context.Context, goal, mainModel string) [ }))) } if s.toolEnabled("cron_delete") { - opts = append(opts, cogito.WithTools(cronDeleteToolDefinition(func(id string) string { + opts = append(opts, s.withTool(cronDeleteToolDefinition(func(id string) string { if s.callbacks.OnCronDelete != nil { return s.callbacks.OnCronDelete(id) } @@ -1430,7 +1436,7 @@ func (s *Session) toolOptions(turnCtx context.Context, goal, mainModel string) [ }))) } if s.toolEnabled("cron_pause") { - opts = append(opts, cogito.WithTools(cronPauseToolDefinition(func(id string) string { + opts = append(opts, s.withTool(cronPauseToolDefinition(func(id string) string { if s.callbacks.OnCronPause != nil { return s.callbacks.OnCronPause(id) } @@ -1438,7 +1444,7 @@ func (s *Session) toolOptions(turnCtx context.Context, goal, mainModel string) [ }))) } if s.toolEnabled("cron_resume") { - opts = append(opts, cogito.WithTools(cronResumeToolDefinition(func(id string) string { + opts = append(opts, s.withTool(cronResumeToolDefinition(func(id string) string { if s.callbacks.OnCronResume != nil { return s.callbacks.OnCronResume(id) } @@ -1446,7 +1452,7 @@ func (s *Session) toolOptions(turnCtx context.Context, goal, mainModel string) [ }))) } if s.toolEnabled("cron_trigger") { - opts = append(opts, cogito.WithTools(cronTriggerToolDefinition(func(id string) string { + opts = append(opts, s.withTool(cronTriggerToolDefinition(func(id string) string { if s.callbacks.OnCronTrigger != nil { return s.callbacks.OnCronTrigger(id) } @@ -1458,21 +1464,21 @@ func (s *Session) toolOptions(turnCtx context.Context, goal, mainModel string) [ // specialist client with the tool's dedicated model and scopes the path to // the session working dir, mirroring how host file tools resolve paths. if s.toolEnabled("read_image") { - opts = append(opts, cogito.WithTools(readImageToolDefinition( + opts = append(opts, s.withTool(readImageToolDefinition( func(path, question string) (string, error) { return specialist.New(s.baseURL, s.apiKey).Describe( turnCtx, resolveWorkspacePath(s.workingDir, path), s.visionModel, question) }))) } if s.toolEnabled("transcribe_audio") { - opts = append(opts, cogito.WithTools(transcribeAudioToolDefinition( + opts = append(opts, s.withTool(transcribeAudioToolDefinition( func(path string) (string, error) { return specialist.New(s.baseURL, s.apiKey).Transcribe( turnCtx, resolveWorkspacePath(s.workingDir, path), s.transcribeModel) }))) } if s.toolEnabled("read_video") { - opts = append(opts, cogito.WithTools(readVideoToolDefinition( + opts = append(opts, s.withTool(readVideoToolDefinition( func(path, question string) (string, error) { return specialist.New(s.baseURL, s.apiKey).DescribeVideo( turnCtx, resolveWorkspacePath(s.workingDir, path), s.videoModel, question) @@ -1485,7 +1491,7 @@ func (s *Session) toolOptions(turnCtx context.Context, goal, mainModel string) [ // the goal so it does not re-arm on the next message. The callback body takes // runMu, which is correct: it fires during a run, not during option assembly. if goal != "" { - opts = append(opts, cogito.WithTools(goalDoneToolDefinition(func(justification string) string { + opts = append(opts, s.withTool(goalDoneToolDefinition(func(justification string) string { s.runMu.Lock() s.goalDone = true s.goal = "" @@ -1498,27 +1504,27 @@ func (s *Session) toolOptions(turnCtx context.Context, goal, mainModel string) [ // Wire the persistent memory tool so the assistant can save and retrieve // notes that survive compaction and model restarts. if s.toolEnabled("memory") { - opts = append(opts, cogito.WithTools(memoryToolDefinition(s.memoryStore))) + opts = append(opts, s.withTool(memoryToolDefinition(s.memoryStore))) } // Wire the ephemeral todo list so the assistant can plan and track // multi-step work within the current session (replace-all semantics, // like maki). if s.toolEnabled("todo_write") { - opts = append(opts, cogito.WithTools(todoWriteToolDefinition(s.todoList))) + opts = append(opts, s.withTool(todoWriteToolDefinition(s.todoList))) } // Wire the tree-sitter index tool so the assistant can skeletonize source // files — a compact structural overview before deciding what to read. if s.toolEnabled("index") { - opts = append(opts, cogito.WithTools(indexToolDefinition( + opts = append(opts, s.withTool(indexToolDefinition( func(p string) string { return resolveWorkspacePath(s.workingDir, p) }))) } // Wire the repo_map tool so the assistant can get a bird's-eye overview of // the whole codebase in one token-budgeted call. if s.toolEnabled("repo_map") { - opts = append(opts, cogito.WithTools(repoMapToolDefinition( + opts = append(opts, s.withTool(repoMapToolDefinition( s.workingDir, func(p string) string { return resolveWorkspacePath(s.workingDir, p) }, ))) @@ -1527,7 +1533,7 @@ func (s *Session) toolOptions(turnCtx context.Context, goal, mainModel string) [ // Wire the LSP tool so the assistant can do symbol-aware navigation // (definition, references, symbols) when a language server is configured. if s.lspManager != nil && s.toolEnabled("lsp") { - opts = append(opts, cogito.WithTools(lspToolDefinition( + opts = append(opts, s.withTool(lspToolDefinition( s.lspManager, func(p string) string { return resolveWorkspacePath(s.workingDir, p) }, ))) @@ -1538,7 +1544,7 @@ func (s *Session) toolOptions(turnCtx context.Context, goal, mainModel string) [ // session on the next turn after any mutating op. for _, d := range selfConfigToolDefs(s.configurator, s.requestReload) { if s.toolEnabled(d.name) { - opts = append(opts, cogito.WithTools(d.def)) + opts = append(opts, s.withTool(d.def)) } } From 74741dc739d6af954bd94b8eb51ee59721fe283d Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Fri, 25 Sep 2026 20:58:42 +0000 Subject: [PATCH 06/16] feat(chat): trim the history step by step when it overflows iterativeTrim runs compaction, compaction with a smaller tail, an aggressive tool-output prune and a hard truncation in turn, and validates each result with validateCompaction before it moves on. compactHistory now calls compactHistoryKeep, so a smaller tail does not write the shared config. Co-Authored-By: Claude Opus 5.5 --- chat/compact.go | 10 +- chat/overflow_trim.go | 233 +++++++++++++++++++++++++++++++++++++ chat/overflow_trim_test.go | 203 ++++++++++++++++++++++++++++++++ 3 files changed, 444 insertions(+), 2 deletions(-) create mode 100644 chat/overflow_trim.go create mode 100644 chat/overflow_trim_test.go diff --git a/chat/compact.go b/chat/compact.go index b0390e6..d3cd4c5 100644 --- a/chat/compact.go +++ b/chat/compact.go @@ -578,11 +578,17 @@ func (s *Session) CompactHistory() (before, after int, err error) { // passed-in ctx governs the summarization LLM call, allowing callers (e.g. // auto-compaction) to make it cancellable via a per-turn context. func (s *Session) compactHistory(ctx context.Context) (before, after int, err error) { + return s.compactHistoryKeep(ctx, s.compactionConfig().KeepRecent) +} + +// compactHistoryKeep is compactHistory with the tail length passed in rather +// than read from the config. Overflow recovery uses it to retry with a +// smaller tail without writing the shared CompactionConfig. +func (s *Session) compactHistoryKeep(ctx context.Context, keep int) (before, after int, err error) { msgs := s.fragment.Messages before = estimateTokens(msgs) - cfg := s.compactionConfig() s.prunedMu.Lock() pruned := make(map[string]string, len(s.prunedIDs)) for k, v := range s.prunedIDs { @@ -592,7 +598,7 @@ func (s *Session) compactHistory(ctx context.Context) (before, after int, err er view := func(m []openai.ChatCompletionMessage) []openai.ChatCompletionMessage { return stubbedView(m, pruned) } - summary, head, tail, err := s.summarizeFitting(ctx, msgs, cfg.KeepRecent, view) + summary, head, tail, err := s.summarizeFitting(ctx, msgs, keep, view) if err != nil { return before, before, err } diff --git a/chat/overflow_trim.go b/chat/overflow_trim.go new file mode 100644 index 0000000..681a01d --- /dev/null +++ b/chat/overflow_trim.go @@ -0,0 +1,233 @@ +package chat + +import ( + "context" + "errors" + "fmt" + "math" + + "github.com/mudler/cogito" + "github.com/mudler/nib/types" + "github.com/mudler/xlog" + openai "github.com/sashabaranov/go-openai" +) + +// trimStep records one step of the iterativeTrim escalation chain and whether +// it was the one that made the history fit. +type trimStep struct { + name string + applied bool +} + +// truncatedMarker is the text of the user-role message that stands in for the +// history a hard truncation drops. +const truncatedMarker = "Previous conversation truncated to fit context window." + +// iterativeTrim makes the fragment fit the context budget after an overflow. +// It runs an escalation chain and validates each step's result with +// validateCompaction before it moves on: +// +// 1. LLM compaction (compactHistory). +// 2. Compaction again with a halved KeepRecent, down to 1, when step 1 +// succeeded but its result did not fit (a large verbatim tail). It is +// skipped when the summary itself failed: a smaller keep makes the head +// larger, and summarizeFitting already moves the boundary for that case. +// 3. Aggressive prune: every tool output outside the trailing run is +// replaced by its stub, and the stubs are committed to the fragment. +// 4. Hard truncate: the history is saved as an artifact, and everything +// before a short tail is replaced by a user-role marker naming it. +// +// It returns nil when a step made the fragment fit, and the last validation +// error (errors.Is works on the Task 1 sentinels) when none did. +func (s *Session) iterativeTrim(ctx context.Context) error { + cfg := s.compactionConfig() + budget := s.trimBudget(cfg) + steps := []trimStep{{name: "compact"}, {name: "shrink-keep"}, {name: "prune"}, {name: "truncate"}} + done := func(i int) error { + steps[i].applied = true + xlog.Info("overflow recovery trimmed the history", "step", steps[i].name, "steps", steps) + return nil + } + + // Step 1: LLM compaction. + before := s.fragmentSnapshot() + _, _, cerr := s.compactHistory(ctx) + lastErr := cerr + if cerr == nil { + lastErr = s.checkCompaction(before, budget) + if lastErr == nil { + return done(0) + } + // Step 2: the summary worked but the result does not fit. Retry with + // a smaller tail, without touching the shared config. + for keep := max(cfg.KeepRecent/2, 1); ; keep /= 2 { + if ctx.Err() != nil { + return ctx.Err() + } + before = s.fragmentSnapshot() + if _, _, err := s.compactHistoryKeep(ctx, keep); err != nil { + lastErr = err + break + } + lastErr = s.checkCompaction(before, budget) + if lastErr == nil { + return done(1) + } + if keep <= 1 { + break + } + } + } + if ctx.Err() != nil { + return ctx.Err() + } + + // Step 3: stub every tool output the prune may drop. + if err := s.forcePrune(budget); err == nil { + return done(2) + } else { + lastErr = err + } + + // Step 4: hard truncate. + if err := s.hardTruncate(cfg, budget); err == nil { + return done(3) + } else { + lastErr = err + } + xlog.Warn("overflow recovery could not fit the history", "error", lastErr) + return fmt.Errorf("iterative trim: %w", lastErr) +} + +// trimBudget is the size the trimmed history must fit: the context budget +// less the prompt the tool schemas and system prompt take. With no known +// window there is nothing to measure against, and any shrink passes. +func (s *Session) trimBudget(cfg types.CompactionConfig) int { + window := s.contextWindow() + if window <= 0 { + return math.MaxInt + } + budget := ContextBudget(cfg, window) + if floor := s.SchemaBudget().Floor; floor > 0 { + budget -= floor + } + return budget +} + +// fragmentState is what a step needs to put back when its result is invalid. +type fragmentState struct { + frag cogito.Fragment + messages []openai.ChatCompletionMessage +} + +func (s *Session) fragmentSnapshot() fragmentState { + s.historyMu.Lock() + defer s.historyMu.Unlock() + return fragmentState{frag: s.fragment, messages: s.messages} +} + +// checkCompaction validates the fragment a compaction installed against the +// one before it. A result that is smaller but over budget stays in place, as +// the next step builds on it. A result that grew, came out empty or broke the +// tool pairing is worse than what it replaced, so the old state goes back. +func (s *Session) checkCompaction(before fragmentState, budget int) error { + s.historyMu.Lock() + defer s.historyMu.Unlock() + err := validateCompaction(before.frag.Messages, s.fragment.Messages, budget) + if err != nil && !errors.Is(err, ErrCompactionOverBudget) && !errors.Is(err, ErrCompactionNoOp) { + s.fragment = before.frag + s.messages = before.messages + } + return err +} + +// installFragment replaces the fragment with msgs, and the display copy with +// display when it is not nil. Like compactHistory, it keeps the running token +// counters but not LastUsage, which measured the history being replaced. +func (s *Session) installFragment(msgs, display []openai.ChatCompletionMessage) { + s.historyMu.Lock() + newFrag := cogito.NewFragment(msgs...) + if s.fragment.Status != nil { + statusCopy := *s.fragment.Status + statusCopy.LastUsage = cogito.LLMUsage{} + newFrag.Status = &statusCopy + } + s.fragment = newFrag + if display != nil { + s.messages = display + } + s.historyMu.Unlock() + s.live.reset() +} + +// forcePrune stubs every tool output except the trailing run and results +// below MinResultTokens, and commits the stubs to the fragment when the result +// fits the budget. +// +// It calls pruneToolOutputs directly rather than pruneMessages: a +// HighWaterTokens of 0 turns size pruning OFF, and effectivePruning raises +// the marks to a fraction of the window. A high mark of 1 and a low mark of 0 +// make the sweep take everything eligible. The new stubs are recorded in +// prunedIDs, so the next request's prune renders the same bytes instead of +// rewriting them. Outlines are not attached: a later prune would replace a stub +// with an outline by the plain stub, which changes an already-sent message. +func (s *Session) forcePrune(budget int) error { + before := s.fragmentSnapshot().frag.Messages + s.prunedMu.Lock() + if s.prunedIDs == nil { + s.prunedIDs = map[string]string{} + } + forced := types.ToolOutputPruningConfig{ + HighWaterTokens: 1, + LowWaterTokens: 0, + MinResultTokens: s.pruning.MinResultTokens, + } + out, newly, _ := pruneToolOutputs(before, forced, s.prunedIDs) + s.prunedMu.Unlock() + + if err := validateCompaction(before, out, budget); err != nil { + return err + } + s.prunedMu.Lock() + for _, n := range newly { + s.prunedIDs[n.id] = n.detail + } + s.prunedMu.Unlock() + s.installFragment(out, nil) + return nil +} + +// hardTruncate keeps the shortest pairing-safe tail that fits, starting from +// KeepRecent messages and halving down to 1, and replaces everything before it +// with a user-role marker. A mid-conversation system message is rejected or +// moved by some adapters, and renderMessages drops it from later summaries, so +// the marker is a user message like summaryMessage. The whole history is +// saved as an artifact first and the marker names it. +func (s *Session) hardTruncate(cfg types.CompactionConfig, budget int) error { + state := s.fragmentSnapshot() + msgs := state.frag.Messages + if len(msgs) < 2 { + return fmt.Errorf("%w: nothing to truncate", ErrCompactionNoOp) + } + + marker := truncatedMarker + if uri := s.spillCompactionHead(renderMessages(msgs)); uri != "" { + marker += " Use " + uri + " to page through the full history." + } + + var lastErr error = fmt.Errorf("%w: nothing to truncate", ErrCompactionNoOp) + for keep := max(cfg.KeepRecent, 1); ; keep /= 2 { + head, tail := splitForCompaction(msgs, keep) + if len(head) > 0 { + candidate := append([]openai.ChatCompletionMessage{{Role: "user", Content: marker}}, tail...) + lastErr = validateCompaction(msgs, candidate, budget) + if lastErr == nil { + s.installFragment(candidate, compactedDisplay(state.messages, tail)) + return nil + } + } + if keep <= 1 { + return lastErr + } + } +} diff --git a/chat/overflow_trim_test.go b/chat/overflow_trim_test.go new file mode 100644 index 0000000..7549f4b --- /dev/null +++ b/chat/overflow_trim_test.go @@ -0,0 +1,203 @@ +package chat + +import ( + "context" + "errors" + "strings" + "testing" + + "github.com/mudler/nib/mcp" + openai "github.com/sashabaranov/go-openai" +) + +// The test session has a 1000-token window, so the budget is 500 byte/4 +// tokens (2000 bytes). + +func trimFragment(s *Session) []openai.ChatCompletionMessage { + s.historyMu.Lock() + defer s.historyMu.Unlock() + return append([]openai.ChatCompletionMessage(nil), s.fragment.Messages...) +} + +func assertPaired(t *testing.T, msgs []openai.ChatCompletionMessage) { + t.Helper() + if err := validateToolPairing(msgs); err != nil { + t.Fatalf("tool pairing broken: %v", err) + } +} + +func TestIterativeTrimStep1CompactionFits(t *testing.T) { + big := strings.Repeat("x", 400) + frag := []openai.ChatCompletionMessage{ + {Role: "user", Content: "u1" + big}, + {Role: "assistant", Content: "a1" + big}, + {Role: "user", Content: "u2" + big}, + {Role: "assistant", Content: "a2" + big}, + {Role: "user", Content: "u3" + big}, + {Role: "assistant", Content: "a3" + big}, + } + llm := &fakeSummaryLLM{reply: "SUMMARY-TEXT"} + s := newCompactTestSession(llm, 2, frag, frag) + + if err := s.iterativeTrim(context.Background()); err != nil { + t.Fatalf("iterativeTrim: %v", err) + } + if llm.calls != 1 { + t.Fatalf("want 1 summary call, got %d", llm.calls) + } + got := trimFragment(s) + if len(got) != 3 || !strings.Contains(got[0].Content, "SUMMARY-TEXT") { + t.Fatalf("want [summary, u3, a3], got %d messages: %q", len(got), got[0].Content) + } + if estimateTokens(got) > ContextBudget(s.compactionConfig(), s.contextWindow()) { + t.Fatalf("result over budget: %d", estimateTokens(got)) + } + assertPaired(t, got) +} + +func TestIterativeTrimNoOpFallsThroughToShrinkKeep(t *testing.T) { + big := strings.Repeat("x", 600) + frag := []openai.ChatCompletionMessage{ + {Role: "user", Content: "u1" + big}, + {Role: "assistant", Content: "a1" + big}, + {Role: "user", Content: "u2" + big}, + {Role: "assistant", Content: "a2" + big}, + } + llm := &fakeSummaryLLM{reply: "SUMMARY-TEXT"} + // KeepRecent 10 > len(frag): compactHistory has nothing to compact. + s := newCompactTestSession(llm, 10, frag, frag) + + if err := s.iterativeTrim(context.Background()); err != nil { + t.Fatalf("iterativeTrim: %v", err) + } + if llm.calls != 1 { + t.Fatalf("want exactly 1 summary call (from the shrunk keep), got %d", llm.calls) + } + got := trimFragment(s) + if len(got) >= len(frag) || !strings.Contains(got[0].Content, "SUMMARY-TEXT") { + t.Fatalf("want a compacted fragment from step 2, got %d messages", len(got)) + } + if s.compaction.KeepRecent != 10 { + t.Fatalf("shared config mutated: KeepRecent=%d", s.compaction.KeepRecent) + } + assertPaired(t, got) +} + +func toolExchange(id, output string) []openai.ChatCompletionMessage { + return []openai.ChatCompletionMessage{ + {Role: "assistant", ToolCalls: []openai.ToolCall{{ID: id, Type: "function", + Function: openai.FunctionCall{Name: "bash", Arguments: `{"command":"ls"}`}}}}, + {Role: "tool", ToolCallID: id, Content: output}, + } +} + +func TestIterativeTrimFallsThroughToAggressivePrune(t *testing.T) { + out := strings.Repeat("o", 2400) + frag := []openai.ChatCompletionMessage{{Role: "user", Content: "u1"}} + frag = append(frag, toolExchange("c1", out)...) + frag = append(frag, toolExchange("c2", out)...) + frag = append(frag, openai.ChatCompletionMessage{Role: "assistant", Content: "done"}) + frag = append(frag, openai.ChatCompletionMessage{Role: "user", Content: "next"}) + + llm := &fakeSummaryLLM{err: errors.New("backend down")} + s := newCompactTestSession(llm, 2, frag, nil) + + if err := s.iterativeTrim(context.Background()); err != nil { + t.Fatalf("iterativeTrim: %v", err) + } + if llm.calls != 1 { + t.Fatalf("a failed summary must not be retried with a smaller keep; got %d calls", llm.calls) + } + got := trimFragment(s) + if len(got) != len(frag) { + t.Fatalf("prune must keep every message, got %d want %d", len(got), len(frag)) + } + for _, m := range got { + if m.Role == "tool" && m.Content == out { + t.Fatalf("tool output %s was not stubbed", m.ToolCallID) + } + } + if estimateTokens(got) > ContextBudget(s.compactionConfig(), s.contextWindow()) { + t.Fatalf("result over budget: %d", estimateTokens(got)) + } + assertPaired(t, got) + // The stubs are recorded, so the next request's prune sends the same bytes. + s.prunedMu.Lock() + _, ok1 := s.prunedIDs["c1"] + _, ok2 := s.prunedIDs["c2"] + s.prunedMu.Unlock() + if !ok1 || !ok2 { + t.Fatalf("stubbed ids not recorded in prunedIDs") + } + next := s.pruneMessages(got) + for i := range got { + if next[i].Content != got[i].Content { + t.Fatalf("message %d changes on the next request: %q -> %q", i, got[i].Content, next[i].Content) + } + } +} + +func TestIterativeTrimHardTruncates(t *testing.T) { + big := strings.Repeat("x", 600) + frag := []openai.ChatCompletionMessage{ + {Role: "user", Content: "u1" + big}, + {Role: "assistant", Content: "a1" + big}, + {Role: "user", Content: "u2" + big}, + {Role: "assistant", Content: "a2" + big}, + {Role: "user", Content: "u3" + big}, + {Role: "assistant", Content: "a3" + big}, + } + llm := &fakeSummaryLLM{err: errors.New("backend down")} + s := newCompactTestSession(llm, 2, frag, frag) + s.artifacts = mcp.NewArtifactStore() + + if err := s.iterativeTrim(context.Background()); err != nil { + t.Fatalf("iterativeTrim: %v", err) + } + got := trimFragment(s) + if got[0].Role != "user" || !strings.Contains(got[0].Content, "truncated to fit context window") || + !strings.Contains(got[0].Content, "artifact://1") { + t.Fatalf("want a user-role truncation marker naming artifact://1, got %s %q", got[0].Role, got[0].Content) + } + if last := got[len(got)-1]; last.Content != frag[len(frag)-1].Content { + t.Fatalf("the most recent message must be kept verbatim") + } + if estimateTokens(got) > ContextBudget(s.compactionConfig(), s.contextWindow()) { + t.Fatalf("result over budget: %d", estimateTokens(got)) + } + art := s.artifacts.Get(1) + if art == nil || !strings.Contains(art.Content, "u1") || !strings.Contains(art.Content, "a1") { + t.Fatalf("artifact must hold the dropped history") + } + assertPaired(t, got) + if len(s.messages) == 0 || s.messages[0].Role != "assistant" { + t.Fatalf("display copy not rebuilt: %+v", s.messages) + } +} + +func TestIterativeTrimAllStepsFail(t *testing.T) { + frag := []openai.ChatCompletionMessage{ + {Role: "user", Content: "u1" + strings.Repeat("x", 400)}, + {Role: "assistant", Content: "a1" + strings.Repeat("x", 400)}, + {Role: "user", Content: strings.Repeat("y", 4000)}, // alone over budget + } + llm := &fakeSummaryLLM{err: errors.New("backend down")} + s := newCompactTestSession(llm, 2, frag, frag) + + err := s.iterativeTrim(context.Background()) + if err == nil { + t.Fatal("want an error when no step fits the budget") + } + if !errors.Is(err, ErrCompactionOverBudget) { + t.Fatalf("want ErrCompactionOverBudget, got %v", err) + } + got := trimFragment(s) + if len(got) != len(frag) { + t.Fatalf("fragment changed on failure: %d messages", len(got)) + } + for i := range frag { + if got[i].Content != frag[i].Content { + t.Fatalf("message %d changed on failure", i) + } + } +} From e2e18effa04f92575c0579ec3c6a189eb4609d33 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Fri, 25 Sep 2026 21:04:31 +0000 Subject: [PATCH 07/16] feat(chat): compact the history without an LLM as a last resort Co-Authored-By: Claude Opus 5.5 --- chat/basic_compact.go | 154 +++++++++++++++++++++++++++++++ chat/basic_compact_test.go | 184 +++++++++++++++++++++++++++++++++++++ 2 files changed, 338 insertions(+) create mode 100644 chat/basic_compact.go create mode 100644 chat/basic_compact_test.go diff --git a/chat/basic_compact.go b/chat/basic_compact.go new file mode 100644 index 0000000..6101e68 --- /dev/null +++ b/chat/basic_compact.go @@ -0,0 +1,154 @@ +package chat + +import ( + "context" + "encoding/json" + "fmt" + "strings" + + openai "github.com/sashabaranov/go-openai" +) + +// basicCompactMarker opens the user-role message that stands in for the head +// basicCompact drops. +const basicCompactMarker = "Conversation compacted (basic fallback mode). Previous tool outputs and assistant responses were elided to fit the context window." + +const ( + basicAssistantChars = 200 // assistant text kept per line + basicArgsChars = 80 // tool arguments kept per line when there is no path +) + +// basicLine is one rendered line of the dropped head. User lines are the +// user's own words and go last when the result must shrink. +type basicLine struct { + text string + user bool +} + +// basicCompact compacts the fragment without an LLM call. It is the last +// resort when the summarizer is unavailable or iterativeTrim ran out of steps. +// +// The whole head is saved as an artifact first. The head is then replaced by +// one user-role message: the user messages verbatim, one line per assistant +// text (cut to 200 characters) and one stub line per tool call. The results +// are rendered as text, never as tool messages, so no tool result is left +// without its call. The tail, split by splitForCompaction, keeps its pairing. +// +// When the result does not fit, activity lines go first, oldest first, and +// then user lines, oldest first; they all stay in the artifact. When even the +// tail does not fit, the fragment is left alone and the validation error is +// returned. +func (s *Session) basicCompact(ctx context.Context) error { + if err := ctx.Err(); err != nil { + return err + } + cfg := s.compactionConfig() + budget := s.trimBudget(cfg) + state := s.fragmentSnapshot() + msgs := state.frag.Messages + + head, tail := splitForCompaction(msgs, max(cfg.KeepRecent, 1)) + if len(head) == 0 { + return fmt.Errorf("%w: nothing to compact", ErrCompactionNoOp) + } + + marker := basicCompactMarker + if uri := s.spillCompactionHead(renderMessages(head)); uri != "" { + marker += " Full history available at " + uri + "." + } + marker += " User messages are preserved verbatim." + + var system []openai.ChatCompletionMessage + for _, m := range head { + if m.Role == "system" { + system = append(system, m) + } + } + lines := basicRender(head) + + build := func() []openai.ChatCompletionMessage { + var b strings.Builder + b.WriteString(marker) + for _, l := range lines { + b.WriteString("\n") + b.WriteString(l.text) + } + out := append([]openai.ChatCompletionMessage(nil), system...) + out = append(out, openai.ChatCompletionMessage{Role: "user", Content: b.String()}) + return append(out, tail...) + } + + candidate := build() + for _, dropUser := range []bool{false, true} { + for estimateTokens(candidate) > budget { + i := oldestLine(lines, dropUser) + if i < 0 { + break + } + lines = append(lines[:i], lines[i+1:]...) + candidate = build() + } + } + if err := validateCompaction(msgs, candidate, budget); err != nil { + return fmt.Errorf("basic compact: %w", err) + } + s.installFragment(candidate, compactedDisplay(state.messages, tail)) + return nil +} + +// oldestLine is the index of the first line of the kind asked for, or -1. +func oldestLine(lines []basicLine, user bool) int { + for i, l := range lines { + if l.user == user { + return i + } + } + return -1 +} + +// basicRender renders the head as text lines: user messages verbatim, the +// first 200 characters of assistant text, and one stub per tool call. System +// messages are kept as messages by the caller, and tool results are covered by +// their call's stub. +func basicRender(head []openai.ChatCompletionMessage) []basicLine { + var lines []basicLine + for _, m := range head { + switch m.Role { + case "user": + if m.Content != "" { + lines = append(lines, basicLine{text: "User: " + m.Content, user: true}) + } + case "assistant": + if text := strings.TrimSpace(m.Content); text != "" { + lines = append(lines, basicLine{text: "Assistant: " + cutRunes(text, basicAssistantChars)}) + } + for _, tc := range m.ToolCalls { + lines = append(lines, basicLine{text: fmt.Sprintf("[tool %s(%s) ran, result elided]", tc.Function.Name, basicToolTarget(tc.Function.Arguments))}) + } + } + } + return lines +} + +// basicToolTarget is what a tool stub names: the path argument when there is +// one, or else the first 80 characters of the arguments. +func basicToolTarget(args string) string { + var parsed map[string]any + if json.Unmarshal([]byte(args), &parsed) == nil { + for _, key := range []string{"path", "file_path"} { + if p, ok := parsed[key].(string); ok && p != "" { + return p + } + } + } + return cutRunes(args, basicArgsChars) +} + +// cutRunes returns the first n characters of s. +func cutRunes(s string, n int) string { + r := []rune(s) + if len(r) <= n { + return s + } + return string(r[:n]) +} diff --git a/chat/basic_compact_test.go b/chat/basic_compact_test.go new file mode 100644 index 0000000..d9f0d5b --- /dev/null +++ b/chat/basic_compact_test.go @@ -0,0 +1,184 @@ +package chat + +import ( + "context" + "errors" + "fmt" + "strings" + "testing" + + "github.com/mudler/nib/mcp" + openai "github.com/sashabaranov/go-openai" +) + +// The test session has a 1000-token window, so the budget is 500 byte/4 +// tokens (2000 bytes). + +func basicToolCall(id, name, args string) openai.ToolCall { + return openai.ToolCall{ID: id, Type: openai.ToolTypeFunction, Function: openai.FunctionCall{Name: name, Arguments: args}} +} + +func TestBasicCompactRendersHeadAndKeepsTail(t *testing.T) { + bigResult := "RESULT-BODY" + strings.Repeat("r", 3000) + longAssistant := "a1" + strings.Repeat("z", 400) + frag := []openai.ChatCompletionMessage{ + {Role: "system", Content: "sys"}, + {Role: "user", Content: "u1 please do X\n exactly like this"}, + {Role: "assistant", ToolCalls: []openai.ToolCall{basicToolCall("c1", "read", `{"path":"/a/b.go"}`)}}, + {Role: "tool", ToolCallID: "c1", Content: bigResult}, + {Role: "assistant", Content: longAssistant}, + {Role: "user", Content: "u2 now run ls"}, + {Role: "assistant", ToolCalls: []openai.ToolCall{basicToolCall("c2", "bash", `{"command":"ls"}`)}}, + {Role: "tool", ToolCallID: "c2", Content: "ok"}, + {Role: "assistant", Content: "done"}, + } + llm := &fakeSummaryLLM{reply: "SUMMARY-TEXT"} + s := newCompactTestSession(llm, 3, frag, frag) + s.artifacts = mcp.NewArtifactStore() + + if err := s.basicCompact(context.Background()); err != nil { + t.Fatalf("basicCompact: %v", err) + } + if llm.calls != 0 { + t.Fatalf("basicCompact must not call the LLM, got %d calls", llm.calls) + } + got := trimFragment(s) + if len(got) != 5 { + t.Fatalf("want [system, marker, 3 tail], got %d: %+v", len(got), got) + } + if got[0].Role != "system" || got[0].Content != "sys" { + t.Fatalf("head system message must be kept first, got %s %q", got[0].Role, got[0].Content) + } + body := got[1].Content + if got[1].Role != "user" || !strings.Contains(body, "basic fallback mode") || !strings.Contains(body, "artifact://1") { + t.Fatalf("want a user-role marker naming artifact://1, got %s %q", got[1].Role, body) + } + for _, want := range []string{ + "u1 please do X\n exactly like this", + "u2 now run ls", + "[tool read(/a/b.go) ran, result elided]", + "Assistant: " + longAssistant[:200], + } { + if !strings.Contains(body, want) { + t.Fatalf("marker missing %q:\n%s", want, body) + } + } + if strings.Contains(body, longAssistant[:201]) || strings.Contains(body, "RESULT-BODY") { + t.Fatalf("assistant text must be cut to 200 chars and tool output elided:\n%s", body) + } + for i, m := range got[2:] { + orig := frag[6+i] + if m.Role != orig.Role || m.Content != orig.Content || m.ToolCallID != orig.ToolCallID || len(m.ToolCalls) != len(orig.ToolCalls) { + t.Fatalf("tail message %d not kept verbatim: %+v", i, m) + } + } + for _, m := range got { + if m.Role == "tool" && m.ToolCallID == "c1" { + t.Fatalf("head tool result must not survive as a tool message") + } + } + assertPaired(t, got) + if estimateTokens(got) > s.trimBudget(s.compactionConfig()) { + t.Fatalf("result over budget: %d", estimateTokens(got)) + } + art := s.artifacts.Get(1) + if art == nil || !strings.Contains(art.Content, "RESULT-BODY") || !strings.Contains(art.Content, "u1 please do X") { + t.Fatalf("artifact must hold the full untruncated head, got %+v", art) + } + if len(s.messages) == 0 || s.messages[0].Role != "assistant" || s.messages[len(s.messages)-1].Content != "done" { + t.Fatalf("display copy not rebuilt: %+v", s.messages) + } +} + +func TestBasicCompactDropsActivityLinesWhenTight(t *testing.T) { + var frag []openai.ChatCompletionMessage + for i := range 10 { + frag = append(frag, + openai.ChatCompletionMessage{Role: "user", Content: fmt.Sprintf("user-%d", i)}, + openai.ChatCompletionMessage{Role: "assistant", Content: fmt.Sprintf("asst-%d-", i) + strings.Repeat("q", 250)}, + ) + } + frag = append(frag, + openai.ChatCompletionMessage{Role: "user", Content: "last-user" + strings.Repeat("t", 400)}, + openai.ChatCompletionMessage{Role: "assistant", Content: "last-asst" + strings.Repeat("t", 400)}, + ) + s := newCompactTestSession(&fakeSummaryLLM{}, 2, frag, frag) + s.artifacts = mcp.NewArtifactStore() + + if err := s.basicCompact(context.Background()); err != nil { + t.Fatalf("basicCompact: %v", err) + } + got := trimFragment(s) + if estimateTokens(got) > s.trimBudget(s.compactionConfig()) { + t.Fatalf("result over budget: %d", estimateTokens(got)) + } + body := got[0].Content + for i := range 10 { + if !strings.Contains(body, fmt.Sprintf("user-%d\n", i)) { + t.Fatalf("user-%d must be preserved while activity lines can go:\n%s", i, body) + } + } + if strings.Contains(body, "asst-0-") { + t.Fatalf("the oldest activity line must be dropped first:\n%s", body) + } + if !strings.Contains(body, "asst-9-") { + t.Fatalf("the newest activity line should survive:\n%s", body) + } + if got[len(got)-1].Content != frag[len(frag)-1].Content { + t.Fatalf("tail must be kept verbatim") + } +} + +func TestBasicCompactDropsOldUsersAsLastResort(t *testing.T) { + var frag []openai.ChatCompletionMessage + for i := range 8 { + frag = append(frag, + openai.ChatCompletionMessage{Role: "user", Content: fmt.Sprintf("user-%d-", i) + strings.Repeat("w", 400)}, + openai.ChatCompletionMessage{Role: "assistant", Content: fmt.Sprintf("asst-%d", i)}, + ) + } + frag = append(frag, + openai.ChatCompletionMessage{Role: "user", Content: "last-user"}, + openai.ChatCompletionMessage{Role: "assistant", Content: "last-asst"}, + ) + s := newCompactTestSession(&fakeSummaryLLM{}, 2, frag, frag) + s.artifacts = mcp.NewArtifactStore() + + if err := s.basicCompact(context.Background()); err != nil { + t.Fatalf("basicCompact: %v", err) + } + got := trimFragment(s) + if estimateTokens(got) > s.trimBudget(s.compactionConfig()) { + t.Fatalf("result over budget: %d", estimateTokens(got)) + } + body := got[0].Content + if strings.Contains(body, "Assistant: ") { + t.Fatalf("activity lines must all go before any user line:\n%s", body) + } + if strings.Contains(body, "user-0-") { + t.Fatalf("the oldest user line must be dropped when nothing else fits:\n%s", body) + } + if !strings.Contains(body, "user-7-"+strings.Repeat("w", 400)) { + t.Fatalf("the newest user line must survive verbatim:\n%s", body) + } + if art := s.artifacts.Get(1); art == nil || !strings.Contains(art.Content, "user-0-") { + t.Fatalf("dropped user lines must stay in the artifact") + } +} + +func TestBasicCompactFailsWhenTailAloneOverBudget(t *testing.T) { + frag := []openai.ChatCompletionMessage{ + {Role: "user", Content: "u1" + strings.Repeat("x", 800)}, + {Role: "assistant", Content: "a1"}, + {Role: "user", Content: strings.Repeat("y", 4000)}, + } + s := newCompactTestSession(&fakeSummaryLLM{}, 1, frag, frag) + + err := s.basicCompact(context.Background()) + if !errors.Is(err, ErrCompactionOverBudget) { + t.Fatalf("want ErrCompactionOverBudget, got %v", err) + } + if got := trimFragment(s); len(got) != len(frag) { + t.Fatalf("a failed basic compact must leave the fragment alone, got %d messages", len(got)) + } +} From e6d8a9d028babcdfef7aef3055e5217db0db9b71 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Fri, 25 Sep 2026 21:19:33 +0000 Subject: [PATCH 08/16] feat(chat): compress old tool output in sticky, graduated levels Add compressToolOutput with four levels (full, truncated, outline, elided) and Session.progressivePrune, which runs pruneMessages and then assigns levels by position and context pressure. The turn's messages manipulator now calls progressivePrune, so every request carries the levels. Levels are recorded by tool_call_id with their rendered text, only go up, and are assigned in one batch only when pressure rises into a higher band. The prompt prefix changes at most once per threshold crossing. Compaction and the forced prune of the overflow recovery render the same levels, so they see what the requests sent. An elided result with no artifact names the tool that produced it. The mid-turn compaction tests now use 7500-token results. At 9000 the request passed 90% of the budget, the older result was elided and the request fell under the trigger, so no compaction ran to test. Co-Authored-By: Claude Opus 5.5 --- chat/compact.go | 8 +- chat/midturn.go | 7 +- chat/midturn_compact_test.go | 23 +- chat/overflow_trim.go | 14 +- chat/progressive_compress.go | 293 +++++++++++++++++++++++ chat/progressive_compress_test.go | 383 ++++++++++++++++++++++++++++++ chat/prune.go | 6 +- chat/session.go | 7 + 8 files changed, 725 insertions(+), 16 deletions(-) create mode 100644 chat/progressive_compress.go create mode 100644 chat/progressive_compress_test.go diff --git a/chat/compact.go b/chat/compact.go index d3cd4c5..c6bf645 100644 --- a/chat/compact.go +++ b/chat/compact.go @@ -594,9 +594,15 @@ func (s *Session) compactHistoryKeep(ctx context.Context, keep int) (before, aft for k, v := range s.prunedIDs { pruned[k] = v } + compressed := s.copyCompressedLocked() s.prunedMu.Unlock() + // The summary sees what the requests saw: stubs, then the progressive + // compression levels (see progressivePrune). view := func(m []openai.ChatCompletionMessage) []openai.ChatCompletionMessage { - return stubbedView(m, pruned) + if len(compressed) == 0 { + return stubbedView(m, pruned) + } + return requestView(m, pruned, compressed) } summary, head, tail, err := s.summarizeFitting(ctx, msgs, keep, view) if err != nil { diff --git a/chat/midturn.go b/chat/midturn.go index 1b7dccd..80460d6 100644 --- a/chat/midturn.go +++ b/chat/midturn.go @@ -80,8 +80,9 @@ func sameMessage(a, b openai.ChatCompletionMessage) bool { } // manipulate is the cogito.WithMessagesManipulator body for a turn: it puts the -// summary in place of the history it covers, prunes tool output, and compacts -// when the result crosses the auto-compaction trigger. +// summary in place of the history it covers, prunes and progressively +// compresses tool output (progressivePrune), and compacts when the result +// crosses the auto-compaction trigger. func (c *turnCompactor) manipulate(msgs []openai.ChatCompletionMessage) []openai.ChatCompletionMessage { base, repl := 0, 0 view := msgs @@ -90,7 +91,7 @@ func (c *turnCompactor) manipulate(msgs []openai.ChatCompletionMessage) []openai view = make([]openai.ChatCompletionMessage, 0, repl+len(msgs)-base) view = append(append(view, c.replacement...), msgs[base:]...) } - out := c.s.pruneMessages(view) + out := c.s.progressivePrune(view) overhead := c.overhead() if !c.failed && c.s.shouldCompactNow(estimateTokens(out)+overhead) { diff --git a/chat/midturn_compact_test.go b/chat/midturn_compact_test.go index 401606e..45ca915 100644 --- a/chat/midturn_compact_test.go +++ b/chat/midturn_compact_test.go @@ -223,11 +223,18 @@ func answer(marker string, tokens int) string { // The real failure: a turn that starts well under the trigger grows past the // window through its own tool results, because auto-compaction only ran after -// the turn. Two 9000-token results take the third request to about 18000 -// tokens, over the 15200 trigger; the turn must compact before sending it. +// the turn. Two 7500-token results and the ~1300-token system prompt take the +// third request to about 16300 tokens, over the 15200 trigger; the turn must +// compact before sending it. +// +// The results are sized so the request lands between 80% and 90% of the +// 19000 budget. There progressivePrune leaves a one-line result as it is, so +// the request still crosses the trigger. Above 90% it elides the older result +// and the request drops under the trigger without a summary, which is the +// compression working, not the compaction under test. func TestMidTurnCompactionKeepsEveryRequestUnderTheTrigger(t *testing.T) { _, backend, events := runMidTurn(t, []string{ - answer("FIRST_ANSWER", 9000), answer("SECOND_ANSWER", 9000), answer("THIRD_ANSWER", 1000), + answer("FIRST_ANSWER", 7500), answer("SECOND_ANSWER", 7500), answer("THIRD_ANSWER", 1000), }) turns := backend.turns() @@ -273,7 +280,7 @@ func TestMidTurnCompactionKeepsEveryRequestUnderTheTrigger(t *testing.T) { // the compaction carry the same one. func TestMidTurnCompactionReusesTheSummary(t *testing.T) { _, backend, _ := runMidTurn(t, []string{ - answer("FIRST_ANSWER", 9000), answer("SECOND_ANSWER", 9000), answer("THIRD_ANSWER", 1000), + answer("FIRST_ANSWER", 7500), answer("SECOND_ANSWER", 7500), answer("THIRD_ANSWER", 1000), }) if n := backend.summaryCount(); n != 1 { @@ -301,7 +308,7 @@ func TestMidTurnCompactionReusesTheSummary(t *testing.T) { // and every tool result follows the assistant message that called it. func TestMidTurnCompactionKeepsSystemPromptAndToolPairs(t *testing.T) { _, backend, _ := runMidTurn(t, []string{ - answer("FIRST_ANSWER", 9000), answer("SECOND_ANSWER", 9000), answer("THIRD_ANSWER", 1000), + answer("FIRST_ANSWER", 7500), answer("SECOND_ANSWER", 7500), answer("THIRD_ANSWER", 1000), }) turns := backend.turns() @@ -341,7 +348,7 @@ func TestMidTurnCompactionKeepsSystemPromptAndToolPairs(t *testing.T) { // copy records the compaction the way compactHistory does. func TestMidTurnCompactionLeavesTheFragmentCompacted(t *testing.T) { s, _, _ := runMidTurn(t, []string{ - answer("FIRST_ANSWER", 9000), answer("SECOND_ANSWER", 9000), answer("THIRD_ANSWER", 1000), + answer("FIRST_ANSWER", 7500), answer("SECOND_ANSWER", 7500), answer("THIRD_ANSWER", 1000), }) s.historyMu.Lock() @@ -411,8 +418,8 @@ func TestNoMidTurnCompactionUnderTheTrigger(t *testing.T) { // off by the length of the first replacement, it would drop or repeat a step. func TestMidTurnCompactionTwiceInOneTurn(t *testing.T) { s, backend, _ := runMidTurn(t, []string{ - answer("FIRST_ANSWER", 9000), answer("SECOND_ANSWER", 9000), - answer("THIRD_ANSWER", 9000), answer("FOURTH_ANSWER", 1000), + answer("FIRST_ANSWER", 7500), answer("SECOND_ANSWER", 7500), + answer("THIRD_ANSWER", 7500), answer("FOURTH_ANSWER", 1000), }) if n := backend.summaryCount(); n != 2 { diff --git a/chat/overflow_trim.go b/chat/overflow_trim.go index 681a01d..0aa1f31 100644 --- a/chat/overflow_trim.go +++ b/chat/overflow_trim.go @@ -183,9 +183,21 @@ func (s *Session) forcePrune(budget int) error { MinResultTokens: s.pruning.MinResultTokens, } out, newly, _ := pruneToolOutputs(before, forced, s.prunedIDs) + // Measure what the requests send: a result the prune keeps may already + // be compressed to a recorded level (see progressivePrune). The new stubs + // take precedence over a level, as they do in the requests. + pruned := make(map[string]string, len(s.prunedIDs)+len(newly)) + for k, v := range s.prunedIDs { + pruned[k] = v + } + for _, n := range newly { + pruned[n.id] = n.detail + } + beforeView := requestView(before, s.prunedIDs, s.compressed) + outView := requestView(out, pruned, s.compressed) s.prunedMu.Unlock() - if err := validateCompaction(before, out, budget); err != nil { + if err := validateCompaction(beforeView, outView, budget); err != nil { return err } s.prunedMu.Lock() diff --git a/chat/progressive_compress.go b/chat/progressive_compress.go new file mode 100644 index 0000000..a3b6805 --- /dev/null +++ b/chat/progressive_compress.go @@ -0,0 +1,293 @@ +package chat + +import ( + "fmt" + "regexp" + "strings" + + openai "github.com/sashabaranov/go-openai" +) + +// CompressionLevel is how much of a tool result a request carries. Higher +// levels carry less. A result's level only ever goes up (see progressivePrune). +type CompressionLevel int + +const ( + // CompressionFull is the original content. + CompressionFull CompressionLevel = iota + // CompressionTruncated keeps the first truncHeadLines and the last + // truncTailLines lines, with a marker for the lines between. + CompressionTruncated + // CompressionOutline keeps the file outline for a read or index result, + // and the first and last line for any other result. + CompressionOutline + // CompressionElided keeps a one-line stub that points to the artifact + // holding the full output, when there is one. + CompressionElided +) + +const ( + truncHeadLines = 20 + truncTailLines = 10 +) + +// The pressure bands at which progressivePrune assigns levels. Below +// compressStartPressure nothing is compressed, so a short session pays no +// prefix-cache re-prefill. +const ( + compressStartPressure = 0.5 + compressEscalatePressure = 0.8 + compressElidePressure = 0.9 +) + +var artifactRefRe = regexp.MustCompile(`artifact://\d+`) + +// compressToolOutput returns content compressed to level. +// +// path is the tool's path argument, or "". outlineOf returns the outline of a +// file, or "" when it has none; the caller passes nil for a result that is not +// a read or index result, and then the outline level keeps the first and last +// line. +// +// The result is a function of its arguments and of outlineOf alone, and the +// Full and Elided levels are idempotent. +func compressToolOutput(content string, path string, level CompressionLevel, outlineOf func(string) string) string { + switch level { + case CompressionTruncated: + lines := strings.Split(content, "\n") + if len(lines) <= truncHeadLines+truncTailLines { + return content + } + elided := len(lines) - truncHeadLines - truncTailLines + out := make([]string, 0, truncHeadLines+truncTailLines+1) + out = append(out, lines[:truncHeadLines]...) + out = append(out, fmt.Sprintf("[...%d lines elided...]", elided)) + out = append(out, lines[len(lines)-truncTailLines:]...) + return strings.Join(out, "\n") + case CompressionOutline: + if path != "" && outlineOf != nil { + if outline := capOutline(outlineOf(path)); outline != "" { + stub := fmt.Sprintf("[outline of %s; read one part back with offset and limit]\n%s", path, outline) + if len(stub) < len(content) { + return stub + } + } + } + lines := strings.Split(content, "\n") + if len(lines) <= 2 { + return content + } + return lines[0] + "\n[...elided...]\n" + lines[len(lines)-1] + case CompressionElided: + return elidedStub("", path, content) + default: + return content + } +} + +// elidedStub is the CompressionElided text for content. It points to the +// artifact the content names, when there is one. Otherwise it names the tool +// and the path, so the model can tell what it lost and fetch it again. +// +// compressToolOutput has no tool name to give and passes "". progressivePrune +// knows the tool from the assistant message that issued the call, and passes +// it. The text is a function of its arguments alone, so it is idempotent. +func elidedStub(name, path, content string) string { + if ref := artifactRefRe.FindString(content); ref != "" { + return "[output elided — see " + ref + " for full content]" + } + what := strings.TrimSpace(name + " " + path) + if what == "" { + return "[output elided]" + } + return "[" + what + " — output elided; re-run or re-read if needed]" +} + +// compressedResult is the level a tool result was compressed to and the text +// it was rendered as. The text is stored, not re-rendered: an outline is read +// from disk, and a file that changed would move the prompt prefix. +type compressedResult struct { + level CompressionLevel + content string +} + +// pressureBand maps context pressure to a band: 0 (no compression), 1 +// (positional levels), 2 (escalated), 3 (everything old elided). +func pressureBand(pressure float64) int { + switch { + case pressure >= compressElidePressure: + return 3 + case pressure >= compressEscalatePressure: + return 2 + case pressure >= compressStartPressure: + return 1 + } + return 0 +} + +// assignLevel is the level the message at index i of n gets in band. The last +// keep messages are always Full. +func assignLevel(i, n, keep, band int) CompressionLevel { + if band <= 0 || i >= n-keep { + return CompressionFull + } + if band >= 3 { + return CompressionElided + } + var lvl CompressionLevel + switch { + case i < n/3: + lvl = CompressionOutline + case i < 2*n/3: + lvl = CompressionTruncated + default: + return CompressionFull + } + if band == 2 { + lvl++ + } + return lvl +} + +// applyCompressed renders every recorded level onto out, in place. A result +// with a stub in pruned keeps its stub. +func applyCompressed(out []openai.ChatCompletionMessage, compressed map[string]compressedResult, pruned map[string]string) { + if len(compressed) == 0 { + return + } + for i := range out { + if out[i].Role != "tool" { + continue + } + c, ok := compressed[out[i].ToolCallID] + if !ok { + continue + } + if _, stubbed := pruned[out[i].ToolCallID]; stubbed { + continue + } + out[i].Content = c.content + } +} + +// requestView returns a copy of msgs as a request renders them: the stubs in +// pruned, then the levels in compressed. Compaction and the overflow recovery +// measure and summarize this view, so they see what the requests saw. +func requestView(msgs []openai.ChatCompletionMessage, pruned map[string]string, compressed map[string]compressedResult) []openai.ChatCompletionMessage { + out := make([]openai.ChatCompletionMessage, len(msgs)) + copy(out, stubbedView(msgs, pruned)) + applyCompressed(out, compressed, pruned) + return out +} + +// compressedViewLocked is requestView with the session's recorded stubs and +// levels. Callers hold prunedMu. +func (s *Session) compressedViewLocked(msgs []openai.ChatCompletionMessage) []openai.ChatCompletionMessage { + return requestView(msgs, s.prunedIDs, s.compressed) +} + +// copyCompressedLocked returns a copy of s.compressed. Callers hold prunedMu. +func (s *Session) copyCompressedLocked() map[string]compressedResult { + out := make(map[string]compressedResult, len(s.compressed)) + for k, v := range s.compressed { + out[k] = v + } + return out +} + +// progressivePrune is pruneMessages plus graduated compression of the tool +// results pruning leaves in full. It is what every request goes through: the +// turn's manipulator (turnCompactor.manipulate) calls it in place of +// pruneMessages. +// +// pruneMessages still runs first. Its stubs are for results that are wrong +// (a later edit) or redundant (a later read), and for the size sweep; a +// result it stubbed keeps that stub and gets no level. +// +// Levels are sticky. Each result's level and rendered text are recorded by +// ToolCallID and re-applied on every call, and a level only goes up. New +// levels are assigned in one batch, only when the context pressure (the +// request as sent, against ContextBudget) rises into a higher band than on +// the previous call. The prefix therefore changes at most once per threshold +// crossing, as with the high/low-water sweep. A result stubbed by +// pruneMessages keeps its stub. +// +// msgs is never written through; the result is a new slice. +func (s *Session) progressivePrune(msgs []openai.ChatCompletionMessage) []openai.ChatCompletionMessage { + out := s.pruneMessages(msgs) + + // modelMu reads first: modelMu must not nest under prunedMu. + cfg := s.compactionConfig() + budget := ContextBudget(cfg, s.contextWindow()) + keep := max(cfg.KeepRecent, 1) + + s.prunedMu.Lock() + defer s.prunedMu.Unlock() + if s.pruning.Disabled { + return out + } + if s.compressed == nil { + s.compressed = map[string]compressedResult{} + } + forgetAbsentCompressed(s.compressed, msgs) + applyCompressed(out, s.compressed, s.prunedIDs) + + band := 0 + if budget > 0 { + band = pressureBand(float64(estimateTokens(out)) / float64(budget)) + } + if band > s.compressBand { + calls := indexToolCalls(msgs) + protected := trailingToolRun(msgs) + for i, m := range msgs { + if m.Role != "tool" || m.ToolCallID == "" || protected[m.ToolCallID] { + continue + } + if _, stubbed := s.prunedIDs[m.ToolCallID]; stubbed { + continue + } + lvl := assignLevel(i, len(msgs), keep, band) + if prev, ok := s.compressed[m.ToolCallID]; lvl == CompressionFull || (ok && prev.level >= lvl) { + continue + } + info := calls[m.ToolCallID] + var outlineOf func(string) string + if pathObservingTools[info.name] { + outlineOf = s.stubOutline + } + text := compressToolOutput(m.Content, info.path, lvl, outlineOf) + if lvl == CompressionElided { + // Only here is the tool's name known; see elidedStub. + name := info.name + if name == "" { + name = "tool" + } + text = elidedStub(name, info.path, m.Content) + } + c := compressedResult{level: lvl, content: text} + s.compressed[m.ToolCallID] = c + out[i].Content = c.content + } + } + s.compressBand = band + return out +} + +// forgetAbsentCompressed drops levels whose tool result is no longer in msgs, +// for the same reason as forgetAbsentIDs. +func forgetAbsentCompressed(c map[string]compressedResult, msgs []openai.ChatCompletionMessage) { + if len(c) == 0 { + return + } + present := make(map[string]bool, len(msgs)) + for _, m := range msgs { + if m.Role == "tool" && m.ToolCallID != "" { + present[m.ToolCallID] = true + } + } + for id := range c { + if !present[id] { + delete(c, id) + } + } +} diff --git a/chat/progressive_compress_test.go b/chat/progressive_compress_test.go new file mode 100644 index 0000000..7c0b212 --- /dev/null +++ b/chat/progressive_compress_test.go @@ -0,0 +1,383 @@ +package chat + +import ( + "context" + "fmt" + "strings" + "testing" + + "github.com/mudler/nib/types" + openai "github.com/sashabaranov/go-openai" +) + +func numberedLines(n int) string { + lines := make([]string, n) + for i := range lines { + lines[i] = fmt.Sprintf("line %d", i+1) + } + return strings.Join(lines, "\n") +} + +func TestProgressiveCompressFullUnchanged(t *testing.T) { + in := numberedLines(100) + if got := compressToolOutput(in, "a.go", CompressionFull, nil); got != in { + t.Fatalf("full level changed the content") + } + if got := compressToolOutput(compressToolOutput(in, "", CompressionFull, nil), "", CompressionFull, nil); got != in { + t.Fatalf("full level not idempotent") + } +} + +func TestProgressiveCompressTruncated(t *testing.T) { + in := numberedLines(100) + got := compressToolOutput(in, "", CompressionTruncated, nil) + lines := strings.Split(got, "\n") + // 20 head + marker + 10 tail + if len(lines) != 31 { + t.Fatalf("got %d lines, want 31:\n%s", len(lines), got) + } + if lines[0] != "line 1" || lines[19] != "line 20" { + t.Fatalf("head wrong: %q .. %q", lines[0], lines[19]) + } + if lines[20] != "[...70 lines elided...]" { + t.Fatalf("marker wrong: %q", lines[20]) + } + if lines[21] != "line 91" || lines[30] != "line 100" { + t.Fatalf("tail wrong: %q .. %q", lines[21], lines[30]) + } + // Short content has nothing to elide. + short := numberedLines(30) + if got := compressToolOutput(short, "", CompressionTruncated, nil); got != short { + t.Fatalf("short content changed: %q", got) + } +} + +func TestProgressiveCompressOutline(t *testing.T) { + in := numberedLines(100) + called := "" + outlineOf := func(p string) string { called = p; return "func Foo [1-10]" } + got := compressToolOutput(in, "a.go", CompressionOutline, outlineOf) + if called != "a.go" { + t.Fatalf("outline function not called with the path, got %q", called) + } + if !strings.Contains(got, "func Foo [1-10]") || strings.Contains(got, "line 50") { + t.Fatalf("outline level did not use the outline: %q", got) + } + + // No outline function (not a read/index result): first + last line. + got = compressToolOutput(in, "", CompressionOutline, nil) + want := "line 1\n[...elided...]\nline 100" + if got != want { + t.Fatalf("got %q, want %q", got, want) + } + + // Outline unavailable: fall back to first + last line. + got = compressToolOutput(in, "a.go", CompressionOutline, func(string) string { return "" }) + if got != want { + t.Fatalf("fallback got %q, want %q", got, want) + } +} + +func TestProgressiveCompressElided(t *testing.T) { + in := "head\n... output spilled; full content at artifact://7 ...\ntail" + got := compressToolOutput(in, "", CompressionElided, nil) + if got != "[output elided — see artifact://7 for full content]" { + t.Fatalf("got %q", got) + } + if again := compressToolOutput(got, "", CompressionElided, nil); again != got { + t.Fatalf("elided not idempotent: %q -> %q", got, again) + } + + got = compressToolOutput(numberedLines(50), "a.go", CompressionElided, nil) + if strings.Contains(got, "artifact://") || !strings.Contains(got, "elided") || !strings.Contains(got, "a.go") { + t.Fatalf("no-artifact stub wrong: %q", got) + } + if again := compressToolOutput(got, "a.go", CompressionElided, nil); again != got { + t.Fatalf("elided not idempotent: %q -> %q", got, again) + } +} + +// progressiveSession builds a session whose budget is window-reserve tokens. +func progressiveSession(window, keep int) *Session { + return &Session{ + compaction: types.CompactionConfig{KeepRecent: keep, MaxContextTokens: window, ReserveTokens: 1, Threshold: 0.8}, + // Size pruning off: only the progressive levels are under test. + pruning: types.ToolOutputPruningConfig{DisableStaleReads: true}, + } +} + +// progressiveHistory is n exchanges of a bash call + 100-line result, then a +// user message. Results sit at indices 1, 3, 5, ... +func progressiveHistory(n int) []openai.ChatCompletionMessage { + var msgs []openai.ChatCompletionMessage + for i := range n { + id := fmt.Sprintf("c%d", i) + msgs = append(msgs, callMsg(id, "bash", `{"command":"ls"}`), resultMsg(id, numberedLines(100))) + } + return append(msgs, openai.ChatCompletionMessage{Role: "user", Content: "next"}) +} + +func levelsOf(t *testing.T, in, out []openai.ChatCompletionMessage) []string { + t.Helper() + var got []string + for i := range out { + if out[i].Role != "tool" { + continue + } + switch { + case out[i].Content == in[i].Content: + got = append(got, "full") + case strings.Contains(out[i].Content, "output elided"): + got = append(got, "elided") + case strings.Contains(out[i].Content, "[...elided...]"): + got = append(got, "outline") + case strings.Contains(out[i].Content, "lines elided...]"): + got = append(got, "truncated") + default: + got = append(got, "?") + } + } + return got +} + +func checkPairing(t *testing.T, in, out []openai.ChatCompletionMessage) { + t.Helper() + if len(in) != len(out) { + t.Fatalf("message count changed %d -> %d", len(in), len(out)) + } + for i := range in { + if in[i].Role != out[i].Role || in[i].ToolCallID != out[i].ToolCallID || len(in[i].ToolCalls) != len(out[i].ToolCalls) { + t.Fatalf("message %d changed shape", i) + } + } +} + +func TestProgressiveCompressAssignment(t *testing.T) { + msgs := progressiveHistory(9) // 19 messages, ~6.3k tokens + // ~60% pressure: positional levels, no escalation. + s := progressiveSession(estimateTokens(msgs)*10/6, 4) + out := s.progressivePrune(msgs) + checkPairing(t, msgs, out) + got := strings.Join(levelsOf(t, msgs, out), ",") + // indices 1,3,5 < 19/3=6 → outline; 7,9,11 < 12 → truncated; 13 → full + // (last third); 15,17 in the last 4 → full. + want := "outline,outline,outline,truncated,truncated,truncated,full,full,full" + if got != want { + t.Fatalf("levels\n got %s\nwant %s", got, want) + } +} + +func TestProgressiveCompressLowPressureLeavesAlone(t *testing.T) { + msgs := progressiveHistory(9) + s := progressiveSession(estimateTokens(msgs)*10, 4) + out := s.progressivePrune(msgs) + for i := range msgs { + if out[i].Content != msgs[i].Content { + t.Fatalf("message %d compressed at low pressure", i) + } + } +} + +func TestProgressiveCompressPressureEscalation(t *testing.T) { + msgs := progressiveHistory(9) + s := progressiveSession(estimateTokens(msgs)*100/85, 4) + out := s.progressivePrune(msgs) + checkPairing(t, msgs, out) + got := strings.Join(levelsOf(t, msgs, out), ",") + want := "elided,elided,elided,outline,outline,outline,full,full,full" + if got != want { + t.Fatalf("80%% levels\n got %s\nwant %s", got, want) + } + + s = progressiveSession(estimateTokens(msgs)*100/95, 4) + out = s.progressivePrune(msgs) + checkPairing(t, msgs, out) + got = strings.Join(levelsOf(t, msgs, out), ",") + want = "elided,elided,elided,elided,elided,elided,elided,full,full" + if got != want { + t.Fatalf("90%% levels\n got %s\nwant %s", got, want) + } +} + +func TestProgressiveCompressReadUsesOutline(t *testing.T) { + msgs := []openai.ChatCompletionMessage{ + callMsg("r1", "read", `{"path":"missing-file.go"}`), resultMsg("r1", numberedLines(100)), + callMsg("b1", "bash", `{}`), resultMsg("b1", numberedLines(100)), + callMsg("b2", "bash", `{}`), resultMsg("b2", numberedLines(100)), + {Role: "user", Content: "next"}, + } + s := progressiveSession(estimateTokens(msgs)*10/6, 2) + out := s.progressivePrune(msgs) + // missing-file.go cannot be indexed, so the read falls back to first+last. + if out[1].Content != "line 1\n[...elided...]\nline 100" { + t.Fatalf("read result: %q", out[1].Content) + } +} + +// Stickiness: once a result is compressed, dropping pressure does not bring it +// back, and consecutive requests below the next threshold keep the prefix. +func TestProgressiveCompressSticky(t *testing.T) { + msgs := progressiveHistory(9) + s := progressiveSession(estimateTokens(msgs)*100/85, 4) + first := s.progressivePrune(msgs) + + // Pressure drops well below every threshold (bigger window). + s.compaction.MaxContextTokens = estimateTokens(msgs) * 10 + second := s.progressivePrune(msgs) + for i := range first { + if first[i].Content != second[i].Content { + t.Fatalf("message %d changed after pressure dropped:\n%q\n->\n%q", i, first[i].Content, second[i].Content) + } + } + + // Growing within the same band: the existing prefix is byte-identical, and + // newly old results are not compressed one per request. + s.compaction.MaxContextTokens = estimateTokens(msgs) * 100 / 55 + third := s.progressivePrune(msgs) + grown := append(append([]openai.ChatCompletionMessage{}, msgs...), + callMsg("n1", "bash", `{}`), resultMsg("n1", numberedLines(100)), + openai.ChatCompletionMessage{Role: "user", Content: "more"}) + fourth := s.progressivePrune(grown) + for i := range third { + if third[i].Content != fourth[i].Content { + t.Fatalf("prefix message %d changed between requests in the same band", i) + } + } + + // Compaction's view renders the same levels the requests saw. + s.prunedMu.Lock() + view := s.compressedViewLocked(grown) + s.prunedMu.Unlock() + for i := range fourth { + if view[i].Content != fourth[i].Content { + t.Fatalf("compaction view differs at %d", i) + } + } +} + +// A result already stubbed by the size sweep keeps its stub bytes. +func TestProgressiveCompressKeepsExistingStubs(t *testing.T) { + msgs := progressiveHistory(9) + s := progressiveSession(estimateTokens(msgs)*100/85, 4) + s.prunedIDs = map[string]string{"c0": detailBudget} + out := s.progressivePrune(msgs) + if want := prunedStub("bash", "", detailBudget); out[1].Content != want { + t.Fatalf("existing stub rewritten: %q", out[1].Content) + } +} + +func TestProgressiveCompressDisabled(t *testing.T) { + msgs := progressiveHistory(9) + s := progressiveSession(estimateTokens(msgs)*100/95, 4) + s.pruning.Disabled = true + out := s.progressivePrune(msgs) + for i := range msgs { + if out[i].Content != msgs[i].Content { + t.Fatalf("message %d compressed with pruning disabled", i) + } + } +} + +// With no artifact to point to, the elided stub progressivePrune renders names +// the tool that produced the output. +func TestProgressiveCompressElidedNamesTool(t *testing.T) { + msgs := progressiveHistory(9) + s := progressiveSession(estimateTokens(msgs)*100/95, 4) + out := s.progressivePrune(msgs) + if !strings.Contains(out[1].Content, "output elided") || !strings.Contains(out[1].Content, "bash") { + t.Fatalf("elided stub does not name the tool: %q", out[1].Content) + } + if strings.Contains(out[1].Content, "artifact://") { + t.Fatalf("elided stub names an artifact that does not exist: %q", out[1].Content) + } +} + +// The real request path compresses: a turn's manipulator applies the levels, +// and two consecutive requests in the same band send byte-identical prefixes. +func TestProgressiveCompressThroughManipulate(t *testing.T) { + msgs := progressiveHistory(9) + s := progressiveSession(estimateTokens(msgs)*10/6, 4) + c := s.newTurnCompactor(context.Background()) + + first := c.manipulate(msgs) + checkPairing(t, msgs, first) + got := strings.Join(levelsOf(t, msgs, first), ",") + want := "outline,outline,outline,truncated,truncated,truncated,full,full,full" + if got != want { + t.Fatalf("manipulate levels\n got %s\nwant %s", got, want) + } + + // The next step of the same turn: one more tool exchange, same band. + grown := append(append([]openai.ChatCompletionMessage{}, msgs...), + callMsg("n1", "bash", `{"command":"ls"}`), resultMsg("n1", "short")) + second := c.manipulate(grown) + for i := range first { + if first[i].Content != second[i].Content { + t.Fatalf("prefix message %d changed between consecutive requests:\n%q\n->\n%q", i, first[i].Content, second[i].Content) + } + } + // The same request again is byte-identical as a whole. + third := c.manipulate(grown) + for i := range second { + if second[i].Content != third[i].Content { + t.Fatalf("message %d changed on a repeated request", i) + } + } +} + +// Compaction summarizes what the requests sent: the recorded levels, not the +// raw fragment. +func TestProgressiveCompressCompactionViewMatchesRequests(t *testing.T) { + var msgs []openai.ChatCompletionMessage + for i := range 9 { + id := fmt.Sprintf("c%d", i) + lines := make([]string, 100) + for j := range lines { + lines[j] = fmt.Sprintf("%s-row-%d", id, j+1) + } + msgs = append(msgs, callMsg(id, "bash", `{"command":"ls"}`), resultMsg(id, strings.Join(lines, "\n"))) + } + msgs = append(msgs, openai.ChatCompletionMessage{Role: "user", Content: "next"}) + + llm := &fakeSummaryLLM{reply: "SUMMARY"} + s := newCompactTestSession(llm, 4, msgs, append([]openai.ChatCompletionMessage(nil), msgs...)) + s.compaction.MaxContextTokens = estimateTokens(msgs) * 10 / 6 + s.compaction.ReserveTokens = 1 + s.pruning = types.ToolOutputPruningConfig{DisableStaleReads: true} + + sent := s.progressivePrune(msgs) + if sent[1].Content == msgs[1].Content { + t.Fatalf("setup: c0 not compressed") + } + if _, _, err := s.compactHistory(context.Background()); err != nil { + t.Fatalf("compactHistory: %v", err) + } + var prompt strings.Builder + for _, m := range llm.lastReq.Messages { + prompt.WriteString(m.Content) + } + if strings.Contains(prompt.String(), "c0-row-50") { + t.Fatalf("summary saw c0's raw body, which the requests never sent") + } + if !strings.Contains(prompt.String(), "c0-row-1") { + t.Fatalf("summary did not see c0's compressed form") + } +} + +// The overflow recovery's forced prune measures the request as sent: a result +// the prune keeps, but a recorded level already compresses, counts at its +// compressed size. +func TestProgressiveCompressForcePruneMeasuresTheRequestView(t *testing.T) { + frag := []openai.ChatCompletionMessage{{Role: "user", Content: "u1"}} + frag = append(frag, toolExchange("a", strings.Repeat("a", 4*5000))...) + frag = append(frag, toolExchange("b", numberedLines(500))...) // ~900 tokens + frag = append(frag, toolExchange("c", "short")...) + s := newCompactTestSession(&fakeSummaryLLM{}, 2, frag, nil) + // b is under MinResultTokens, so the forced prune keeps it. + s.pruning = types.ToolOutputPruningConfig{MinResultTokens: 2000} + s.compressed = map[string]compressedResult{"b": {level: CompressionElided, content: "[bash — output elided; re-run or re-read if needed]"}} + + if err := s.forcePrune(500); err != nil { + t.Fatalf("forcePrune measured the raw fragment, not the request: %v", err) + } +} diff --git a/chat/prune.go b/chat/prune.go index 3315e21..e5c0bc4 100644 --- a/chat/prune.go +++ b/chat/prune.go @@ -439,9 +439,9 @@ func pruneToolOutputs(msgs []openai.ChatCompletionMessage, cfg types.ToolOutputP return out, newly, freed } -// pruneMessages is the cogito.WithMessagesManipulator body: it rewrites the -// messages about to be sent, and records what it stubbed so the next call makes -// the same decisions. +// pruneMessages rewrites the messages about to be sent, and records what it +// stubbed so the next call makes the same decisions. The request path reaches +// it through progressivePrune, which compresses what it leaves in full. // // It runs on the turn goroutine but the state is guarded anyway — the // manipulator is called from inside cogito's loop, and nothing here should diff --git a/chat/session.go b/chat/session.go index b57176e..c00a244 100644 --- a/chat/session.go +++ b/chat/session.go @@ -258,6 +258,13 @@ type Session struct { // a stub whose wording changed between calls would move the prompt prefix // just as un-stubbing it would. prunedIDs map[string]string + // compressed maps each tool_call_id progressivePrune compressed to its + // level and the text it was rendered as. Like prunedIDs it only grows in + // level, so an already-sent result keeps its bytes. compressBand is the + // pressure band of the previous call; levels are assigned only when the + // band rises above it. + compressed map[string]compressedResult + compressBand int // outputLimitsMu guards the tool-output limits policy (budget, per-line // truncation, artifact spill). SetToolOutputLimits writes it from the UI From 009ae0d63fe52ce59ed0dbaed834b12557b684e0 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Fri, 25 Sep 2026 21:38:22 +0000 Subject: [PATCH 09/16] feat(chat): classify overflow errors from a catalog of patterns Replace the substring markers and the "first two N tokens" guesser with an ordered list of regex rows. Named groups (window, total, input, output) say what each figure means. classifyOverflow applies the structured error code, the rows, an output-cap guard and a 413 rule, and returns one of KindContext, KindBudget or KindOutputCap. Fixtures in chat/testdata/overflow include the two regolo captures on both cogito paths, and the typed streaming text form cogito is moving to. Each row is marked strong or weak. The status gate (400, 413 or unknown) applies only to weak evidence: the generic phrase rows and the 413 rule. A structured code, a provider row or a config row classifies at any status, because LocalAI and llama.cpp report an overflow as HTTP 500 (gRPC Internal). isContextOverflow is now KindContext only. isBudgetOverflow and isOutputCapError are added for the recovery work. learnedWindowFrom, overflowFigures and overflowDetail read the catalog. compaction.overflow_patterns adds user rows ahead of the built-ins. Changes to existing tests and callers: - summary_shift_test: TestSummaryGivesUpAfterMaxAttempts expects isWindowOverflow. Its vLLM rejections state an input that fits, so they are budget overflows now. - config_test: compare Compaction with reflect.DeepEqual, because the struct now holds a slice. - The summary retry (summarizeFitting, rollingCover) shrinks on context and budget overflows through canShrinkSummary. - classifyBackendError reads classifyOverflow: a LocalAI overflow at HTTP 500 stays fatal; a generic phrase in a 5xx is retried. - No test learned a window from a wording that no row names. Co-Authored-By: Claude Opus 5.5 --- README.md | 19 + chat/compact.go | 46 ++- chat/errors.go | 147 +++---- chat/errors_window_test.go | 8 +- chat/overflow_patterns.go | 364 ++++++++++++++++++ chat/overflow_patterns_test.go | 337 ++++++++++++++++ chat/retry.go | 5 +- chat/session.go | 11 +- chat/settings.go | 1 + chat/summary_shift_test.go | 4 +- chat/testdata/overflow/anthropic-budget.json | 10 + .../overflow/anthropic-prompt-too-long.json | 10 + .../code-context-length-exceeded.json | 10 + .../code-context-window-exceeded.json | 10 + chat/testdata/overflow/gemini.json | 10 + .../generic-maximum-context-length.json | 10 + .../llamacpp-localai-500-typed-stream.json | 11 + .../overflow/llamacpp-localai-500.json | 11 + chat/testdata/overflow/llamacpp-localai.json | 10 + .../overflow/llamacpp-stream-chunk.json | 10 + chat/testdata/overflow/negative-auth.json | 10 + .../negative-generic-server-error.json | 11 + .../negative-max-tokens-validation.json | 10 + .../overflow/negative-rate-limit.json | 11 + chat/testdata/overflow/openai.json | 10 + .../regolo-vllm-budget-nonstream.json | 10 + .../overflow/regolo-vllm-budget-stream.json | 10 + .../regolo-vllm-output-cap-nonstream.json | 10 + .../regolo-vllm-output-cap-stream.json | 10 + .../overflow/vllm-current-budget.json | 10 + .../overflow/vllm-current-typed-stream.json | 11 + chat/testdata/overflow/vllm-current.json | 10 + chat/testdata/overflow/vllm-older.json | 10 + config/config.go | 24 ++ config/config_test.go | 15 +- selfdoc/README.md | 19 + types/config.go | 16 + 37 files changed, 1122 insertions(+), 129 deletions(-) create mode 100644 chat/overflow_patterns.go create mode 100644 chat/overflow_patterns_test.go create mode 100644 chat/testdata/overflow/anthropic-budget.json create mode 100644 chat/testdata/overflow/anthropic-prompt-too-long.json create mode 100644 chat/testdata/overflow/code-context-length-exceeded.json create mode 100644 chat/testdata/overflow/code-context-window-exceeded.json create mode 100644 chat/testdata/overflow/gemini.json create mode 100644 chat/testdata/overflow/generic-maximum-context-length.json create mode 100644 chat/testdata/overflow/llamacpp-localai-500-typed-stream.json create mode 100644 chat/testdata/overflow/llamacpp-localai-500.json create mode 100644 chat/testdata/overflow/llamacpp-localai.json create mode 100644 chat/testdata/overflow/llamacpp-stream-chunk.json create mode 100644 chat/testdata/overflow/negative-auth.json create mode 100644 chat/testdata/overflow/negative-generic-server-error.json create mode 100644 chat/testdata/overflow/negative-max-tokens-validation.json create mode 100644 chat/testdata/overflow/negative-rate-limit.json create mode 100644 chat/testdata/overflow/openai.json create mode 100644 chat/testdata/overflow/regolo-vllm-budget-nonstream.json create mode 100644 chat/testdata/overflow/regolo-vllm-budget-stream.json create mode 100644 chat/testdata/overflow/regolo-vllm-output-cap-nonstream.json create mode 100644 chat/testdata/overflow/regolo-vllm-output-cap-stream.json create mode 100644 chat/testdata/overflow/vllm-current-budget.json create mode 100644 chat/testdata/overflow/vllm-current-typed-stream.json create mode 100644 chat/testdata/overflow/vllm-current.json create mode 100644 chat/testdata/overflow/vllm-older.json diff --git a/README.md b/README.md index 1505974..1c5f36d 100644 --- a/README.md +++ b/README.md @@ -575,8 +575,27 @@ log_level: error # prompt is fitted into the window minus that reservation. If the backend # still rejects it as too large, nib summarizes fewer of the older turns and # keeps the rest verbatim, instead of cutting the text it summarizes. +# +# nib reads a backend's "request too large" error to decide what to do: +# compact when the prompt itself does not fit (kind "context"), lower the +# requested output when only prompt + output does not fit ("budget"), or +# lower the output cap when the output alone is above the model's maximum +# ("output_cap"). It knows the wording of llama.cpp/LocalAI, vLLM, OpenAI, +# Anthropic and Gemini. overflow_patterns teaches it another backend without a +# rebuild: the rows are tried before the built-in ones, first match wins. A +# regex states its figures with the named groups window, total, input and +# output. A row with an invalid regex or kind is logged and skipped. A +# backend-specific wording (built-in or from this list) counts at any HTTP +# status, because LocalAI reports an overflow as HTTP 500; a generic phrase +# such as "context window" counts only with status 400, 413 or none. An error +# with HTTP status 400 or 413 that no pattern recognises is logged at debug +# level, so you can copy its text from the log. compaction: summary_max_tokens: 16384 + overflow_patterns: + - name: my-backend + kind: context + regex: 'prompt has (?P\d+) tokens, limit is (?P\d+)' # Optional: external MCP servers mcp_servers: diff --git a/chat/compact.go b/chat/compact.go index c6bf645..27c63f5 100644 --- a/chat/compact.go +++ b/chat/compact.go @@ -2,9 +2,7 @@ package chat import ( "context" - "errors" "fmt" - "strconv" "strings" "github.com/mudler/nib/types" @@ -97,8 +95,8 @@ func shouldAutoCompact(cfg types.CompactionConfig, window, promptTokens int) boo // 4096-context model is exactly nib's audience, so the floor sits below any // real served model and above any plausible misparse of an output-token count // (the largest default max_tokens backends print in overflow errors is in the -// hundreds — vLLM's "512 output tokens" is the known example, and it is only -// skipped today because tokenCountRe's \s* cannot span the word "output"). +// hundreds — vLLM's "512 output tokens" is the known example; the overflow +// catalog's named groups keep it out of the window, see overflow_patterns.go). // // The ceiling exists because strconv.Atoi clamps an absurdly long digit run to // MaxInt instead of erroring, so garbage lands far ABOVE any floor rather than @@ -219,6 +217,16 @@ func canRecoverFromOverflow(turnCtx context.Context, err error) bool { return turnCtx.Err() == nil && isContextOverflow(err) } +// canShrinkSummary reports whether a rejected summary request is worth +// retrying with a smaller prompt: a context overflow, or a budget overflow +// (the prompt fits but prompt + reserved output does not), the user did not +// cancel. The summary's output reservation is already capped, so shrinking +// its prompt is what fits a budget overflow here; this is not compaction of +// the conversation, which a budget overflow must never trigger. +func canShrinkSummary(ctx context.Context, err error) bool { + return ctx.Err() == nil && isWindowOverflow(err) +} + // overflowRetries reports how many context-overflow recoveries the current turn // performed. Read by tests; the cap itself is enforced in SendMessage. func (s *Session) overflowRetries() int { @@ -543,22 +551,20 @@ func summaryRetryTarget(err error, sent int) int { } // overflowFigures reads the request size and the window from an overflow -// error, applying the same rule as learnedWindowFrom: the larger figure is the -// request and the smaller is the limit, whatever order the backend used. +// error for the subtractive summary retry: (Total, Window) from the catalog +// row, with Input + Output (or Input alone) standing in for a missing total. +// A context or a budget overflow qualifies; a generic wording states no +// figures and reports ok == false, so the caller halves instead. func overflowFigures(err error) (needs, allows int, ok bool) { - for e := err; e != nil; e = errors.Unwrap(e) { - m := tokenCountRe.FindAllStringSubmatch(e.Error(), -1) - if !hasOverflowMarker(e.Error()) || len(m) < 2 { - continue - } - a, _ := strconv.Atoi(m[0][1]) - b, _ := strconv.Atoi(m[1][1]) - needs, allows = max(a, b), min(a, b) - if allows > 0 { - return needs, allows, true - } + info := classifyOverflow(err) + if info.Kind != KindContext && info.Kind != KindBudget { + return 0, 0, false + } + needs, allows = info.overflowNeeds(), info.Window + if needs <= 0 || allows <= 0 { + return 0, 0, false } - return 0, 0, false + return needs, allows, true } // CompactHistory summarizes the older portion of the conversation via the LLM @@ -710,7 +716,7 @@ func (s *Session) summarizeFitting(ctx context.Context, msgs []openai.ChatComple if serr == nil { return summary, head, tail, nil } - if maxSummaryAttempts <= 1 || !canRecoverFromOverflow(ctx, serr) { + if maxSummaryAttempts <= 1 || !canShrinkSummary(ctx, serr) { return "", nil, nil, serr } target := summaryRetryTarget(serr, sent) @@ -820,7 +826,7 @@ func (s *Session) rollingCover(ctx context.Context, head []openai.ChatCompletion continue } failures++ - if failures >= maxSummaryAttempts || !canRecoverFromOverflow(ctx, serr) { + if failures >= maxSummaryAttempts || !canShrinkSummary(ctx, serr) { return "", 0, serr } if target = summaryRetryTarget(serr, sent); target <= 0 { diff --git a/chat/errors.go b/chat/errors.go index 43f16ee..0ed57a4 100644 --- a/chat/errors.go +++ b/chat/errors.go @@ -3,42 +3,9 @@ package chat import ( "errors" "fmt" - "regexp" - "strconv" "strings" ) -// contextOverflowMarkers are substrings different backends emit when a request -// exceeds the model's context window. Only a fragment of each message is -// matched here, because the wording around it varies. The full messages, which -// learnedWindowFrom parses for the model's stated limit, are: -// -// - llama.cpp / LocalAI (tools/server/server-context.cpp): -// "request (9739 tokens) exceeds the available context size (8192 tokens), try increasing it" -// - OpenAI: -// "This model's maximum context length is 8192 tokens. However, your messages resulted in 9739 tokens. Please reduce the length of the messages." -// - vLLM, current (vllm/renderers/params.py): -// "This model's maximum context length is 8192 tokens. However, you requested 512 output tokens and your prompt contains 9739 input tokens, for a total of 10251 tokens. Please reduce the length of the input prompt or the number of requested output tokens." -// - vLLM, older OpenAI-compatible server: -// "This model's maximum context length is 8192 tokens. However, you requested 10251 tokens (9739 in the messages, 512 in the completion). Please reduce the length of the messages or completion." -// -// This comment used to abbreviate vLLM's message to "maximum context length is -// 8192 tokens". That fragment is enough to trip a marker but carries only ONE -// figure, so it teaches learnedWindowFrom nothing — an abbreviation that reads -// as a complete message invites tests asserting coverage nib does not have. -// Record the whole thing. -var contextOverflowMarkers = []string{ - "exceeds the available context size", - "maximum context length", - "context length", - "context window", - "context size", -} - -// tokenCountRe pulls token counts (e.g. "8192 tokens") out of a backend error -// so we can report how far over the limit the request was. -var tokenCountRe = regexp.MustCompile(`(\d{3,})\s*tokens`) - // FriendlyError wraps a noisy backend error with a short, actionable message // while preserving the original via Unwrap (so errors.Is/As still work). type FriendlyError struct { @@ -56,7 +23,7 @@ func humanizeError(err error) error { if err == nil { return nil } - if isContextOverflow(err) { + if isWindowOverflow(err) { return &FriendlyError{err: err, msg: contextOverflowMessage(err.Error())} } if isRateLimitError(err) { @@ -106,31 +73,28 @@ const emptyReplyMessage = "the model returned an empty reply (no text and no too // A bool parameter would leave one function whose every caller passes a // constant, and both branches would still have to be tested separately. func humanizeTurnError(err error, compacted bool) error { - if compacted && isContextOverflow(err) { + if compacted && isWindowOverflow(err) { return &FriendlyError{err: err, msg: contextOverflowRetriedMessage(err.Error())} } return humanizeError(err) } // isContextOverflow reports whether err is a backend complaining that the -// request did not fit the model's context window. Factored out of -// humanizeError so the recovery path and the message path cannot drift apart -// by consulting different marker lists. +// prompt itself did not fit the model's context window: the one rejection +// compaction can fix. A budget overflow (prompt + requested output) and an +// output-cap error are not context overflows and must not trigger compaction. +// Every overflow question in the package reads classifyOverflow, so the +// recovery path and the message path cannot drift apart. func isContextOverflow(err error) bool { - return err != nil && hasOverflowMarker(err.Error()) + return classifyOverflow(err).Kind == KindContext } -// hasOverflowMarker holds the single walk over contextOverflowMarkers. Both -// isContextOverflow and learnedWindowFrom go through it, so there is exactly -// one marker list and one matching rule in the package. -func hasOverflowMarker(msg string) bool { - low := strings.ToLower(msg) - for _, marker := range contextOverflowMarkers { - if strings.Contains(low, marker) { - return true - } - } - return false +// isWindowOverflow reports a context or a budget overflow: the request, as +// sent, did not fit the window. It decides the user-facing message, which is +// true for both; recovery must still tell them apart. +func isWindowOverflow(err error) bool { + k := classifyOverflow(err).Kind + return k == KindContext || k == KindBudget } // learnedWindowFrom extracts the model's real context window from an overflow @@ -139,67 +103,52 @@ func hasOverflowMarker(msg string) bool { // so inconsistently, so this error is nib's only trustworthy source; inferring // a window from a model name would be a guess presented as a fact. // -// Two figures are required. tokenCountRe can match a single number, and with -// one number there is no way to tell whether it is the limit or the request -// size — treating a request size as the window would raise the compaction -// trigger above the real limit and suppress compaction exactly when it is most -// needed. Fewer than two figures teaches nothing. -// -// The window is the SMALLER figure regardless of the order the backend printed -// them, matching contextOverflowMessage. +// A window is learned only from a catalog row that names it (the "window" +// group): a generic wording, or a row that states only a request size, teaches +// nothing, because mistaking a request size for the window would raise the +// compaction trigger above the real limit. Context and budget overflows both +// state the real window; an output-cap error states the maximum OUTPUT, which +// is not the window. // -// The whole unwrap chain is tried, not just err.Error(). humanizeError wraps -// the backend error in a FriendlyError whose own text drops the second figure -// — contextOverflowMessage prints "model allows 262144)" with no "tokens" for -// tokenCountRe to anchor on — so reading only the outermost message would -// learn nothing from a humanized error. Recovery runs before humanizeError -// today, but that ordering is one edit away from silently never learning a -// window again, and no test would catch it. Parsing every level makes the -// order irrelevant. +// classifyOverflow walks the whole unwrap chain, so a humanized error (whose +// own text drops the figures) still teaches the window of the error it wraps. func learnedWindowFrom(err error) (int, bool) { - for e := err; e != nil; e = errors.Unwrap(e) { - if n, ok := windowFromMessage(e.Error()); ok { - return n, true - } + info := classifyOverflow(err) + if (info.Kind == KindContext || info.Kind == KindBudget) && info.Window > 0 { + return info.Window, true } return 0, false } -// windowFromMessage applies the two-figure rule to one error message. -func windowFromMessage(msg string) (int, bool) { - if !hasOverflowMarker(msg) { - return 0, false - } - m := tokenCountRe.FindAllStringSubmatch(msg, -1) - if len(m) < 2 { - return 0, false +// overflowNeeds is the request size an overflow states: the total, else input +// plus output, else the input alone (llama.cpp and OpenAI state only the +// prompt, and reserve no output in the count). +func (o overflowInfo) overflowNeeds() int { + switch { + case o.Total > 0: + return o.Total + case o.Input > 0 && o.Output > 0: + return o.Input + o.Output + default: + return o.Input } - a, _ := strconv.Atoi(m[0][1]) - b, _ := strconv.Atoi(m[1][1]) - if b < a { - a = b - } - if a <= 0 { - return 0, false - } - return a, true } // overflowDetail renders the parenthesised token figures for an overflow -// message, or "" when the backend did not report two of them. The larger count -// is the request size and the smaller is the model's limit, regardless of the -// order the backend printed them — the same rule learnedWindowFrom applies. +// message, or "" when the backend stated none. Which figure is the request +// size and which the limit comes from the catalog row, not number order. func overflowDetail(raw string) string { - m := tokenCountRe.FindAllStringSubmatch(raw, -1) - if len(m) < 2 { - return "" - } - needs, _ := strconv.Atoi(m[0][1]) - allows, _ := strconv.Atoi(m[1][1]) - if allows > needs { - needs, allows = allows, needs + info := classifyOverflow(errors.New(raw)) + needs := info.overflowNeeds() + switch { + case needs > 0 && info.Window > 0: + return fmt.Sprintf(" (needs ~%d tokens, model allows %d)", needs, info.Window) + case info.Window > 0: + return fmt.Sprintf(" (model allows %d)", info.Window) + case needs > 0: + return fmt.Sprintf(" (needs ~%d tokens)", needs) } - return fmt.Sprintf(" (needs ~%d tokens, model allows %d)", needs, allows) + return "" } // contextOverflowMessage builds the user-facing text for the FIRST context- diff --git a/chat/errors_window_test.go b/chat/errors_window_test.go index 1d1bb72..5b9082c 100644 --- a/chat/errors_window_test.go +++ b/chat/errors_window_test.go @@ -33,11 +33,9 @@ func TestLearnedWindowFromRealBackendPhrasings(t *testing.T) { want: 8192, }, { - // vLLM current, vllm/renderers/params.py. Note the near-miss: the - // "512 output tokens" and "9739 input tokens" figures do NOT match - // tokenCountRe, because the intervening word breaks `\s*tokens`. - // The two that do match are the limit and the total, so the - // smaller is still the window. + // vLLM current, vllm/renderers/params.py. The catalog row names + // each figure (window, output, input, total), so the window is + // read by name, not guessed from number order. name: "vLLM, current", msg: "This model's maximum context length is 8192 tokens. However, you requested 512 output tokens and your prompt contains 9739 input tokens, for a total of 10251 tokens. Please reduce the length of the input prompt or the number of requested output tokens.", want: 8192, diff --git a/chat/overflow_patterns.go b/chat/overflow_patterns.go new file mode 100644 index 0000000..eafe242 --- /dev/null +++ b/chat/overflow_patterns.go @@ -0,0 +1,364 @@ +package chat + +import ( + "errors" + "regexp" + "strconv" + "strings" + "sync/atomic" + + "github.com/mudler/nib/types" + "github.com/mudler/xlog" + openai "github.com/sashabaranov/go-openai" +) + +// OverflowKind says what a backend rejection means for recovery. +type OverflowKind int + +const ( + // KindNone: the error is not about size. + KindNone OverflowKind = iota + // KindContext: the prompt itself does not fit the window. Compact. + KindContext + // KindBudget: the prompt fits, but prompt + requested output does not. + // Lower max_tokens and retry. Never compact. + KindBudget + // KindOutputCap: the requested output alone is above the model's maximum. + // Lower the output cap. Do not compact. + KindOutputCap +) + +// String returns the name fixtures and config use for the kind. +func (k OverflowKind) String() string { + switch k { + case KindContext: + return "context" + case KindBudget: + return "budget" + case KindOutputCap: + return "output_cap" + default: + return "none" + } +} + +func parseOverflowKind(s string) (OverflowKind, bool) { + switch s { + case "context": + return KindContext, true + case "budget": + return KindBudget, true + case "output_cap": + return KindOutputCap, true + } + return KindNone, false +} + +// overflowInfo is what classifyOverflow read from an error. A zero figure +// means the message did not state it. +type overflowInfo struct { + Kind OverflowKind + Pattern string + Status int + Window int + Total int + Input int + Output int +} + +// hasFigures reports whether the match stated any number at all. +func (o overflowInfo) hasFigures() bool { + return o.Window > 0 || o.Total > 0 || o.Input > 0 || o.Output > 0 +} + +// overflowPattern is one row of the catalog. Figures come from the named +// capture groups window, total, input and output, so a row states what each +// number means instead of relying on the order the backend printed them. +type overflowPattern struct { + name string + kind OverflowKind + re *regexp.Regexp + strong bool +} + +// Evidence strength of a row, for the status gate (see classifyOverflowWith). +// A strong row (a structured code or a specific provider wording) classifies +// at any HTTP status; a weak row (a generic phrase) only with 400, 413 or an +// unknown status. +const ( + strongEvidence = true + weakEvidence = false +) + +// builtinOverflowPatterns are tried in order, most specific first; the first +// match wins. Each row with figures is backed by a fixture in +// testdata/overflow. The generic phrases at the end detect an overflow in a +// wording no row names, but state no figures, so no window is learned from +// them and the summary retry falls back to halving. +var builtinOverflowPatterns = []overflowPattern{ + // vLLM, newest (seen through regolo). + {"vllm_total", KindContext, regexp.MustCompile(`maximum context length of (?P\d+) tokens\. You requested a total of (?P\d+) tokens: (?P\d+) tokens from the input messages and (?P\d+) tokens for the completion`), strongEvidence}, + // vLLM, current (vllm/renderers/params.py). + {"vllm_current", KindContext, regexp.MustCompile(`maximum context length is (?P\d+) tokens\. However, you requested (?P\d+) output tokens and your prompt contains (?P\d+) input tokens, for a total of (?P\d+) tokens`), strongEvidence}, + // vLLM, older OpenAI-compatible server. + {"vllm_older", KindContext, regexp.MustCompile(`maximum context length is (?P\d+) tokens\. However, you requested (?P\d+) tokens \((?P\d+) in the messages, (?P\d+) in the completion\)`), strongEvidence}, + // Anthropic, prompt + max_tokens. + {"anthropic_budget", KindBudget, regexp.MustCompile("input length and `?max_tokens`? exceed context limit: (?P\\d+) \\+ (?P\\d+) > (?P\\d+)"), strongEvidence}, + // vLLM output cap (seen through regolo; note no space after the period). + {"vllm_output_cap", KindOutputCap, regexp.MustCompile(`max_(?:completion_)?tokens is too large: (?P\d+)\.\s*This model supports at most (?P\d+) completion tokens`), strongEvidence}, + // OpenAI. + {"openai", KindContext, regexp.MustCompile(`maximum context length is (?P\d+) tokens\. However, your messages resulted in (?P\d+) tokens`), strongEvidence}, + // llama.cpp / LocalAI (tools/server/server-context.cpp). + {"llamacpp", KindContext, regexp.MustCompile(`request \((?P\d+) tokens\) exceeds the available context size \((?P\d+) tokens\)`), strongEvidence}, + // Anthropic, prompt alone. + {"anthropic", KindContext, regexp.MustCompile(`prompt is too long: (?P\d+) tokens > (?P\d+) maximum`), strongEvidence}, + // Gemini / Vertex. + {"gemini", KindContext, regexp.MustCompile(`(?i)input token count \((?P\d+)\) exceeds the maximum number of tokens allowed \((?P\d+)\)`), strongEvidence}, + // Bare codes. + {"code_context_length_exceeded", KindContext, regexp.MustCompile(`context_length_exceeded|ContextWindowExceededError`), strongEvidence}, + // Generic phrases: detection only, no figures. + {"generic_available_context_size", KindContext, regexp.MustCompile(`(?i)exceeds the available context size`), weakEvidence}, + {"generic_maximum_context_length", KindContext, regexp.MustCompile(`(?i)maximum context length`), weakEvidence}, + {"generic_context_window", KindContext, regexp.MustCompile(`(?i)context window`), weakEvidence}, + {"generic_context_size", KindContext, regexp.MustCompile(`(?i)context size`), weakEvidence}, + {"generic_context_length", KindContext, regexp.MustCompile(`(?i)context length`), weakEvidence}, +} + +// configOverflowPatterns holds the compiled compaction.overflow_patterns rows. +// It is package state rather than a session field because classifyOverflow is +// a free function called from paths that hold no session (humanizeError, +// classifyBackendError); every session of one process loads the same config. +var configOverflowPatterns atomic.Pointer[[]overflowPattern] + +// setConfigOverflowPatterns compiles the user's rows. An invalid regex or kind +// is logged and skipped; it never fails the session. +func setConfigOverflowPatterns(in []types.OverflowPattern) { + var rows []overflowPattern + for _, p := range in { + kind, ok := parseOverflowKind(p.Kind) + if !ok { + xlog.Warn("overflow pattern skipped: unknown kind", "name", p.Name, "kind", p.Kind) + continue + } + re, err := regexp.Compile(p.Regex) + if err != nil { + xlog.Warn("overflow pattern skipped: invalid regex", "name", p.Name, "error", err) + continue + } + name := p.Name + if name == "" { + name = "config" + } + // A config row names one backend's wording: strong evidence. + rows = append(rows, overflowPattern{name: name, kind: kind, re: re, strong: strongEvidence}) + } + configOverflowPatterns.Store(&rows) +} + +func overflowRows() []overflowPattern { + if p := configOverflowPatterns.Load(); p != nil && len(*p) > 0 { + return append(append([]overflowPattern(nil), *p...), builtinOverflowPatterns...) + } + return builtinOverflowPatterns +} + +// errorStatus returns the HTTP status carried by the error chain, or 0 when +// unknown. The typed go-openai errors are read first; then the text, because +// cogito's streaming path flattens the status into "localai stream: status +// 400: " and the OpenAI SDK prints "status code: 400". +func errorStatus(err error) int { + if err == nil { + return 0 + } + var re *openai.RequestError + if errors.As(err, &re) && re.HTTPStatusCode > 0 { + return re.HTTPStatusCode + } + var ae *openai.APIError + if errors.As(err, &ae) && ae.HTTPStatusCode > 0 { + return ae.HTTPStatusCode + } + if m := statusRe.FindStringSubmatch(strings.ToLower(err.Error())); m != nil { + n, _ := strconv.Atoi(m[1]) + return n + } + return 0 +} + +// classifyOverflow reports what kind of size rejection err is, with the +// figures the backend stated. It knows nothing about the request that failed, +// so a 413 without token wording is KindNone (see classifyOverflowWith). +func classifyOverflow(err error) overflowInfo { + return classifyOverflowWith(err, 0, 0) +} + +// classifyOverflowWith is classifyOverflow with the local estimate of the last +// request and the window it was sent against. They decide only one case: a 413 +// whose text names no tokens, which is an overflow when the request was at +// least 80% of the window and otherwise a byte or media limit that compaction +// cannot fix. Zeros mean "unknown", and such a 413 is KindNone. +// +// The status gate applies to weak evidence only: a generic phrase row and the +// 413 rule count when the status is 400, 413 or unknown (0: an error chunk +// mid-stream carries none, and that must not disable recovery). Strong +// evidence (a structured code, a specific provider row, a config row) +// classifies at any status: LocalAI and llama.cpp report an overflow as HTTP +// 500, and proxies rewrap statuses. +func classifyOverflowWith(err error, lastRequestTokens, window int) overflowInfo { + if err == nil { + return overflowInfo{} + } + status := errorStatus(err) + weakOK := status == 0 || status == 400 || status == 413 + m := matchOverflow(err, weakOK) + m.Status = status + if m.Kind != KindNone { + return m + } + // e. 413 without token wording (weak evidence; the status is 413 here, so + // the gate always admits it). + if status == 413 { + if window > 0 && lastRequestTokens > 0 && float64(lastRequestTokens) >= 0.8*float64(window) { + return overflowInfo{Kind: KindContext, Pattern: "http_413", Status: status} + } + xlog.Debug("HTTP 413 not classified as a context overflow: no token wording and the last request was not near the window", + "error", err.Error(), "last_request_tokens", lastRequestTokens, "window", window) + } + return overflowInfo{Status: status} +} + +// matchOverflow applies steps b-d (structured code, regex rows, output-cap +// guard) to err. weakOK says whether the status gate admits weak rows. +func matchOverflow(err error, weakOK bool) overflowInfo { + var info overflowInfo + if err == nil { + return info + } + // b. Structured error code: strong evidence. + var ae *openai.APIError + if errors.As(err, &ae) { + if code, ok := ae.Code.(string); ok { + switch code { + case "context_length_exceeded": + info = overflowInfo{Kind: KindContext, Pattern: "code:" + code} + case "max_tokens_exceeded": + info = overflowInfo{Kind: KindBudget, Pattern: "code:" + code} + } + } + } + // c. Regex rows, over every level of the chain: humanizeError wraps the + // backend error in a FriendlyError whose own text keeps the generic + // wording but not every figure. A row with figures wins over one without, + // and a strong row over a weak one. + var row overflowInfo + rowStrong := false + for e := err; e != nil; e = errors.Unwrap(e) { + r, strong := matchRows(e.Error(), weakOK) + if r.Kind == KindNone { + continue + } + if row.Kind == KindNone || (strong && !rowStrong) { + row, rowStrong = r, strong + } + if r.hasFigures() { + row, rowStrong = r, strong + break + } + } + if info.Kind != KindNone { + // The code decides the kind; the rows only add figures. + if row.hasFigures() { + info.Window, info.Total, info.Input, info.Output = row.Window, row.Total, row.Input, row.Output + } + return info + } + // d. Output-cap guard: a message naming the output parameter that no row + // with figures explained is an invalid-parameter error, not an overflow. + if !row.hasFigures() && namesOutputParam(err.Error()) { + return overflowInfo{} + } + return row +} + +var outputParamRe = regexp.MustCompile(`(?i)max_(?:completion_|output_)?tokens`) + +func namesOutputParam(msg string) bool { return outputParamRe.MatchString(msg) } + +// matchRows returns the first row matching msg, with its figures, the kind +// they imply and whether the row is strong evidence. Weak rows are skipped +// unless weakOK. +func matchRows(msg string, weakOK bool) (overflowInfo, bool) { + for _, p := range overflowRows() { + if !p.strong && !weakOK { + continue + } + sub := p.re.FindStringSubmatch(msg) + if sub == nil { + continue + } + info := overflowInfo{Kind: p.kind, Pattern: p.name} + for i, g := range p.re.SubexpNames() { + if g == "" || sub[i] == "" { + continue + } + n, err := strconv.Atoi(sub[i]) + if err != nil { + continue + } + switch g { + case "window": + info.Window = n + case "total": + info.Total = n + case "input": + info.Input = n + case "output": + info.Output = n + } + } + // Budget vs Context comes from the figures, not the row: when the + // input fits the window and output is reserved, lowering the output + // fixes it; otherwise the prompt alone does not fit. + if p.kind != KindOutputCap && info.Input > 0 && info.Output > 0 && info.Window > 0 { + if info.Input < info.Window { + info.Kind = KindBudget + } else { + info.Kind = KindContext + } + } + return info, p.strong + } + return overflowInfo{}, false +} + +// isBudgetOverflow reports whether err says the prompt fits but prompt plus +// the requested output does not. +func isBudgetOverflow(err error) bool { + return classifyOverflow(err).Kind == KindBudget +} + +// isOutputCapError reports whether err says the requested output alone is +// above the model's maximum. +func isOutputCapError(err error) bool { + return classifyOverflow(err).Kind == KindOutputCap +} + +// logUnclassifiedRejection logs a 400/413 that no pattern recognised, so an +// unknown overflow wording can be turned into a pattern. +func logUnclassifiedRejection(err error, info overflowInfo) { + if err == nil || info.Kind != KindNone { + return + } + if s := errorStatus(err); s == 400 || s == 413 { + xlog.Debug("request rejected; if this is a context overflow, add a pattern (see compaction.overflow_patterns)", + "status", s, "error", err.Error()) + } +} + +// classifyTurnOverflow is classifyOverflowWith for the turn that just failed: +// the last request's size is the in-flight usage figure, or the conversation's +// estimate when the backend has reported none (ContextTokens), measured +// against the window the session budgets for. +func (s *Session) classifyTurnOverflow(err error) overflowInfo { + return classifyOverflowWith(err, s.ContextTokens(), s.contextWindow()) +} diff --git a/chat/overflow_patterns_test.go b/chat/overflow_patterns_test.go new file mode 100644 index 0000000..6e2f187 --- /dev/null +++ b/chat/overflow_patterns_test.go @@ -0,0 +1,337 @@ +package chat + +import ( + "encoding/json" + "errors" + "fmt" + "os" + "path/filepath" + "strings" + "testing" + + "github.com/mudler/nib/types" + openai "github.com/sashabaranov/go-openai" +) + +// overflowFixture is one file of chat/testdata/overflow: an error string as +// nib's error chain prints it, and what classifyOverflow must make of it. +type overflowFixture struct { + Provider string `json:"provider"` + Source string `json:"source"` + Raw string `json:"raw"` + Kind string `json:"kind"` + Window int `json:"window"` + Total int `json:"total"` + Input int `json:"input"` + Output int `json:"output"` + // Status, when set, is the HTTP status classifyOverflow must report. + Status int `json:"status"` +} + +func loadOverflowFixtures(t *testing.T) map[string]overflowFixture { + t.Helper() + paths, err := filepath.Glob(filepath.Join("testdata", "overflow", "*.json")) + if err != nil { + t.Fatal(err) + } + if len(paths) == 0 { + t.Fatal("no fixtures in testdata/overflow") + } + out := make(map[string]overflowFixture, len(paths)) + for _, p := range paths { + b, err := os.ReadFile(p) + if err != nil { + t.Fatal(err) + } + var f overflowFixture + if err := json.Unmarshal(b, &f); err != nil { + t.Fatalf("%s: %v", p, err) + } + out[strings.TrimSuffix(filepath.Base(p), ".json")] = f + } + return out +} + +// Every fixture classifies to its kind and its figures. A new provider is one +// new file in the folder. +func TestOverflowPatternFixtures(t *testing.T) { + for name, f := range loadOverflowFixtures(t) { + t.Run(name, func(t *testing.T) { + got := classifyOverflow(errors.New(f.Raw)) + if got.Kind.String() != f.Kind { + t.Fatalf("kind = %s (pattern %q), want %s", got.Kind, got.Pattern, f.Kind) + } + if got.Window != f.Window || got.Total != f.Total || got.Input != f.Input || got.Output != f.Output { + t.Fatalf("figures = window %d total %d input %d output %d, want %d %d %d %d", + got.Window, got.Total, got.Input, got.Output, f.Window, f.Total, f.Input, f.Output) + } + if f.Status != 0 && got.Status != f.Status { + t.Fatalf("status = %d, want %d", got.Status, f.Status) + } + if got.Kind != KindNone && got.Pattern == "" { + t.Fatal("a classified error must name the pattern that matched") + } + }) + } +} + +// The two real regolo captures, verbatim, on both of cogito's paths. +func TestOverflowPatternRegoloCaptures(t *testing.T) { + fx := loadOverflowFixtures(t) + for _, name := range []string{"regolo-vllm-budget-stream", "regolo-vllm-budget-nonstream"} { + f, ok := fx[name] + if !ok { + t.Fatalf("fixture %s missing", name) + } + err := errors.New(f.Raw) + info := classifyOverflow(err) + if info.Kind != KindBudget { + t.Fatalf("%s: kind = %s, want budget: 133 input tokens fit, the 209990 reservation does not", name, info.Kind) + } + if isContextOverflow(err) { + t.Fatalf("%s: a budget overflow must not trigger compaction", name) + } + if !isBudgetOverflow(err) { + t.Fatalf("%s: isBudgetOverflow = false", name) + } + if w, ok := learnedWindowFrom(err); !ok || w != 210000 { + t.Fatalf("%s: learned window = %d, %v, want 210000", name, w, ok) + } + needs, allows, ok := overflowFigures(err) + if !ok || needs != 210123 || allows != 210000 { + t.Fatalf("%s: overflowFigures = %d, %d, %v, want 210123, 210000", name, needs, allows, ok) + } + } + for _, name := range []string{"regolo-vllm-output-cap-stream", "regolo-vllm-output-cap-nonstream"} { + f := fx[name] + err := errors.New(f.Raw) + if got := classifyOverflow(err); got.Kind != KindOutputCap || got.Output != 250000 || got.Window != 210000 { + t.Fatalf("%s: %+v, want output_cap 250000 / 210000", name, got) + } + if isContextOverflow(err) { + t.Fatalf("%s: an output-cap error is not a context overflow", name) + } + if !isOutputCapError(err) { + t.Fatalf("%s: isOutputCapError = false", name) + } + if w, ok := learnedWindowFrom(err); ok { + t.Fatalf("%s: learned window %d from an output-cap error", name, w) + } + } +} + +// Budget versus Context is decided from the figures: the same row with an +// input that alone fills the window is a context overflow. +func TestClassifyOverflowBudgetVersusContext(t *testing.T) { + budget := errors.New("Requested token count exceeds the model's maximum context length of 210000 tokens. You requested a total of 210123 tokens: 133 tokens from the input messages and 209990 tokens for the completion.") + if got := classifyOverflow(budget); got.Kind != KindBudget { + t.Fatalf("kind = %s, want budget", got.Kind) + } + ctx := errors.New("Requested token count exceeds the model's maximum context length of 210000 tokens. You requested a total of 212000 tokens: 210000 tokens from the input messages and 2000 tokens for the completion.") + if got := classifyOverflow(ctx); got.Kind != KindContext { + t.Fatalf("kind = %s, want context: the prompt alone does not fit", got.Kind) + } + if !isContextOverflow(ctx) { + t.Fatal("isContextOverflow = false for a prompt that alone fills the window") + } +} + +// The status gate applies to weak evidence only. A specific provider row is +// strong and classifies at any status: LocalAI and llama.cpp report an +// overflow as HTTP 500 over gRPC. A generic phrase counts only with 400, 413 +// or an unknown status. +func TestClassifyOverflowStatusGate(t *testing.T) { + const localai = "rpc error: code = Internal desc = request (9739 tokens) exceeds the available context size (8192 tokens), try increasing it" + strong := []struct { + name string + err error + }{ + {"localai stream text", errors.New("localai stream: status 500: " + localai)}, + {"localai typed stream text", errors.New("localai stream: error, status code: 500, status: 500 Internal Server Error, message: " + localai)}, + {"RequestError", fmt.Errorf("wrapped: %w", &openai.RequestError{HTTPStatusCode: 500, Err: errors.New(localai)})}, + {"APIError", fmt.Errorf("wrapped: %w", &openai.APIError{HTTPStatusCode: 500, Message: localai})}, + } + for _, tc := range strong { + t.Run("strong 500 "+tc.name, func(t *testing.T) { + got := classifyOverflow(tc.err) + if got.Kind != KindContext || got.Status != 500 || got.Window != 8192 || got.Input != 9739 { + t.Fatalf("%+v, want context with status 500, window 8192, input 9739", got) + } + if !isContextOverflow(tc.err) { + t.Fatal("isContextOverflow = false for a LocalAI overflow reported as HTTP 500") + } + if classifyBackendError(tc.err) != errFatal { + t.Fatal("a LocalAI overflow reported as HTTP 500 must not be retried") + } + }) + } + for _, status := range []int{429, 503, 401} { + err := fmt.Errorf("localai stream: status %d: %s", status, localai) + if got := classifyOverflow(err); got.Kind != KindContext || got.Window != 8192 { + t.Fatalf("strong row, status %d: %+v, want context", status, got) + } + } + + const generic = "context window exceeded" + for _, status := range []int{500, 429, 401, 503} { + err := fmt.Errorf("localai stream: status %d: %s", status, generic) + if got := classifyOverflow(err); got.Kind != KindNone { + t.Fatalf("generic row, status %d: kind = %s (pattern %q), want none", status, got.Kind, got.Pattern) + } + if got := classifyOverflow(&openai.RequestError{HTTPStatusCode: status, Err: errors.New(generic)}); got.Kind != KindNone { + t.Fatalf("generic row, RequestError %d: kind = %s, want none", status, got.Kind) + } + if got := classifyOverflow(&openai.APIError{HTTPStatusCode: status, Message: generic}); got.Kind != KindNone { + t.Fatalf("generic row, APIError %d: kind = %s, want none", status, got.Kind) + } + } + if got := classifyOverflow(errors.New(generic)); got.Kind != KindContext || got.Status != 0 { + t.Fatalf("generic row, no status (stream error chunk): %+v, want context with status 0", got) + } + + cases := []struct { + name string + err error + }{ + {"RequestError", fmt.Errorf("wrapped: %w", &openai.RequestError{HTTPStatusCode: 400, Err: errors.New(generic)})}, + {"APIError", fmt.Errorf("wrapped: %w", &openai.APIError{HTTPStatusCode: 400, Message: generic})}, + {"localai stream text", errors.New("localai stream: status 400: " + generic)}, + {"localai typed stream text", errors.New("localai stream: error, status code: 400, status: 400 Bad Request, message: " + generic)}, + {"openai sdk text", errors.New("error, status code: 400, status: 400 Bad Request, message: " + generic)}, + } + for _, tc := range cases { + t.Run("generic 400 "+tc.name, func(t *testing.T) { + if got := errorStatus(tc.err); got != 400 { + t.Fatalf("errorStatus = %d, want 400", got) + } + if got := classifyOverflow(tc.err); got.Kind != KindContext || got.Status != 400 { + t.Fatalf("%+v, want context with status 400", got) + } + }) + } + if got := errorStatus(errors.New("dial tcp: connection refused")); got != 0 { + t.Fatalf("errorStatus of a status-less error = %d, want 0", got) + } + if got := errorStatus(nil); got != 0 { + t.Fatalf("errorStatus(nil) = %d", got) + } +} + +// Every built-in row says whether it is strong or weak evidence, and only the +// generic phrases are weak. +func TestOverflowPatternStrength(t *testing.T) { + for _, p := range builtinOverflowPatterns { + generic := strings.HasPrefix(p.name, "generic") + if p.strong == generic { + t.Fatalf("row %q: strong = %v, want %v", p.name, p.strong, !generic) + } + } +} + +// A structured error code decides the kind even when no row knows the text. +func TestClassifyOverflowStructuredCode(t *testing.T) { + err := fmt.Errorf("wrapped: %w", &openai.APIError{Code: "context_length_exceeded", Message: "your request is too big for this model"}) + if got := classifyOverflow(err); got.Kind != KindContext { + t.Fatalf("kind = %s, want context", got.Kind) + } + err = &openai.APIError{Code: "max_tokens_exceeded", Message: "the reservation does not fit"} + if got := classifyOverflow(err); got.Kind != KindBudget { + t.Fatalf("kind = %s, want budget", got.Kind) + } + // The rows still run for the figures. + err = &openai.APIError{Code: "context_length_exceeded", Message: "This model's maximum context length is 8192 tokens. However, your messages resulted in 9739 tokens."} + if got := classifyOverflow(err); got.Kind != KindContext || got.Window != 8192 || got.Input != 9739 { + t.Fatalf("%+v, want context with window 8192, input 9739", got) + } +} + +// The output-cap guard: a max_tokens validation error is not an overflow. +func TestClassifyOverflowMaxTokensValidation(t *testing.T) { + for _, msg := range []string{ + "Invalid 'max_tokens': integer above maximum value", + "max_tokens must be less than the context length", + "max_completion_tokens: exceeds the context window", + } { + if got := classifyOverflow(errors.New(msg)); got.Kind != KindNone { + t.Fatalf("%q: kind = %s (pattern %q), want none", msg, got.Kind, got.Pattern) + } + } +} + +// A 413 whose text names no tokens may be a byte or media limit, which +// compaction cannot fix: it is an overflow only when the last request was +// near the window. +func TestClassifyOverflow413WithoutTokenWording(t *testing.T) { + err := &openai.RequestError{HTTPStatusCode: 413, Err: errors.New("Request Entity Too Large")} + if got := classifyOverflow(err); got.Kind != KindNone { + t.Fatalf("no estimate: kind = %s, want none", got.Kind) + } + if got := classifyOverflowWith(err, 7000, 10000); got.Kind != KindNone { + t.Fatalf("70%% of the window: kind = %s, want none", got.Kind) + } + if got := classifyOverflowWith(err, 8000, 10000); got.Kind != KindContext || got.Status != 413 { + t.Fatalf("80%% of the window: %+v, want context with status 413", got) + } + // A 413 with token wording is classified by its row, estimate or not. + worded := &openai.RequestError{HTTPStatusCode: 413, Err: errors.New("prompt is too long: 210000 tokens > 200000 maximum")} + if got := classifyOverflow(worded); got.Kind != KindContext || got.Window != 200000 { + t.Fatalf("%+v, want context with window 200000", got) + } +} + +// Config patterns are tried before the built-ins; an invalid regex is skipped. +func TestOverflowPatternConfigPrecedence(t *testing.T) { + t.Cleanup(func() { setConfigOverflowPatterns(nil) }) + msg := errors.New("request (9739 tokens) exceeds the available context size (8192 tokens), try increasing it") + if got := classifyOverflow(msg); got.Pattern != "llamacpp" { + t.Fatalf("built-in pattern = %q, want llamacpp", got.Pattern) + } + setConfigOverflowPatterns([]types.OverflowPattern{ + {Name: "broken", Kind: "context", Regex: `(unclosed`}, + {Name: "bad-kind", Kind: "sideways", Regex: `request`}, + {Name: "mine", Kind: "context", Regex: `request \((?P\d+) tokens\) exceeds the available context size \((?P\d+) tokens\)`}, + {Name: "new-backend", Kind: "context", Regex: `prompt of (?P\d+) is over (?P\d+)`}, + }) + got := classifyOverflow(msg) + if got.Pattern != "mine" || got.Window != 8192 || got.Input != 9739 { + t.Fatalf("%+v, want the config row \"mine\" with its figures", got) + } + if got := classifyOverflow(errors.New("prompt of 5000 is over 4096")); got.Kind != KindContext || got.Pattern != "new-backend" || got.Window != 4096 { + t.Fatalf("%+v, want the config row for a wording nib does not know", got) + } + if got := classifyOverflow(errors.New("request timed out")); got.Kind != KindNone { + t.Fatalf("an invalid-kind row matched: %+v", got) + } +} + +// A generic-row match is an overflow with no figures, and no window is +// learned from it. +func TestClassifyOverflowGenericRow(t *testing.T) { + err := errors.New("the prompt exceeds the context window of this deployment") + got := classifyOverflow(err) + if got.Kind != KindContext || !strings.HasPrefix(got.Pattern, "generic") { + t.Fatalf("%+v, want a generic context match", got) + } + if got.Window != 0 || got.Total != 0 || got.Input != 0 || got.Output != 0 { + t.Fatalf("a generic row stated figures: %+v", got) + } + if w, ok := learnedWindowFrom(err); ok { + t.Fatalf("learned window %d from a generic row", w) + } + if _, _, ok := overflowFigures(err); ok { + t.Fatal("overflowFigures reported figures for a generic row") + } +} + +// A figure-carrying row nested under a humanized FriendlyError still wins +// over the generic wording of the outer message. +func TestClassifyOverflowPrefersFiguresInTheChain(t *testing.T) { + raw := errors.New("Requested token count exceeds the model's maximum context length of 210000 tokens. You requested a total of 210123 tokens: 133 tokens from the input messages and 209990 tokens for the completion.") + humanized := humanizeError(raw) + if humanized == raw { + t.Fatal("a budget overflow was not humanized") + } + if got := classifyOverflow(humanized); got.Kind != KindBudget || got.Window != 210000 { + t.Fatalf("%+v, want the budget figures from the wrapped error", got) + } +} diff --git a/chat/retry.go b/chat/retry.go index 80011d1..397e48e 100644 --- a/chat/retry.go +++ b/chat/retry.go @@ -135,7 +135,10 @@ var fatalMarkers = []string{ // it. Only the small set above — client-side failures that never succeed on // retry — stops immediately. func classifyBackendError(err error) backendErrorClass { - if err == nil || isContextOverflow(err) { + // Any size rejection the catalog recognises is fatal. Specific rows + // classify at any status, so a LocalAI overflow reported as a 500 is not + // repeated; a generic phrase in a 5xx is not an overflow and is retried. + if err == nil || classifyOverflow(err).Kind != KindNone { return errFatal } low := strings.ToLower(err.Error()) diff --git a/chat/session.go b/chat/session.go index c00a244..a3eae77 100644 --- a/chat/session.go +++ b/chat/session.go @@ -684,6 +684,7 @@ func NewSession(ctx context.Context, cfg types.Config, callbacks Callbacks, tran s.compactionAutoDetected = true s.compaction.MaxContextTokens = defaultContextTokens } + setConfigOverflowPatterns(cfg.Compaction.OverflowPatterns) return s, nil } @@ -1899,7 +1900,13 @@ func (s *Session) SendMessage(text string, parts ...ContentPart) (string, error) // Never after an interrupt: a cancelled turnCtx means the user // pressed Ctrl+C, and re-sending is the opposite of what they asked // for. See canRecoverFromOverflow. - if canRecoverFromOverflow(turnCtx, err) { + // + // The turn's own classification, not the free isContextOverflow: + // only the session knows how large the failed request was, which + // decides whether a 413 with no token wording is an overflow. + turnOverflow := s.classifyTurnOverflow(err) + logUnclassifiedRejection(err, turnOverflow) + if turnCtx.Err() == nil && turnOverflow.Kind == KindContext { if w, ok := learnedWindowFrom(err); ok { s.rememberWindow(w, mainModel) } @@ -1998,7 +2005,7 @@ func (s *Session) SendMessage(text string, parts ...ContentPart) (string, error) // clearing the conversation, because compaction just did the // equivalent and the retry has already been made — see // contextOverflowRetriedMessage. - overflow := isContextOverflow(err) + overflow := turnOverflow.Kind == KindContext err = humanizeTurnError(err, s.overflowRetries() > 0) if s.callbacks.OnError != nil { s.callbacks.OnError(err) diff --git a/chat/settings.go b/chat/settings.go index 3693085..47b4625 100644 --- a/chat/settings.go +++ b/chat/settings.go @@ -65,6 +65,7 @@ func (s *Session) SetCompaction(c types.CompactionConfig) { s.compactionAutoDetected = false } s.compaction = c + setConfigOverflowPatterns(c.OverflowPatterns) } // compactionConfig returns a copy of the compaction policy, read under the diff --git a/chat/summary_shift_test.go b/chat/summary_shift_test.go index 6c0ddc3..194bd6b 100644 --- a/chat/summary_shift_test.go +++ b/chat/summary_shift_test.go @@ -164,7 +164,9 @@ func TestSummaryGivesUpAfterMaxAttempts(t *testing.T) { s.artifacts = mcp.NewArtifactStore() _, _, err := s.CompactHistory() - if !isContextOverflow(err) { + // The prompt fits beside the reservation's window, so the catalog reads + // these rejections as budget overflows, not context overflows. + if !isWindowOverflow(err) { t.Fatalf("want the last overflow, got %v", err) } if len(llm.reqs) != maxSummaryAttempts { diff --git a/chat/testdata/overflow/anthropic-budget.json b/chat/testdata/overflow/anthropic-budget.json new file mode 100644 index 0000000..f4d4088 --- /dev/null +++ b/chat/testdata/overflow/anthropic-budget.json @@ -0,0 +1,10 @@ +{ + "provider": "Anthropic", + "source": "Anthropic Messages API 400 invalid_request_error text (maki-providers/src/error.rs fixture)", + "raw": "input length and `max_tokens` exceed context limit: 199773 + 8192 > 200000, decrease input length or `max_tokens` and try again", + "kind": "budget", + "window": 200000, + "total": 0, + "input": 199773, + "output": 8192 +} diff --git a/chat/testdata/overflow/anthropic-prompt-too-long.json b/chat/testdata/overflow/anthropic-prompt-too-long.json new file mode 100644 index 0000000..e4f60f2 --- /dev/null +++ b/chat/testdata/overflow/anthropic-prompt-too-long.json @@ -0,0 +1,10 @@ +{ + "provider": "Anthropic", + "source": "Anthropic Messages API 400 invalid_request_error text", + "raw": "prompt is too long: 210000 tokens > 200000 maximum", + "kind": "context", + "window": 200000, + "total": 0, + "input": 210000, + "output": 0 +} diff --git a/chat/testdata/overflow/code-context-length-exceeded.json b/chat/testdata/overflow/code-context-length-exceeded.json new file mode 100644 index 0000000..1d64b56 --- /dev/null +++ b/chat/testdata/overflow/code-context-length-exceeded.json @@ -0,0 +1,10 @@ +{ + "provider": "OpenAI-compatible", + "source": "bare error code, no figures", + "raw": "error, status code: 400, status: 400 Bad Request, message: context_length_exceeded", + "kind": "context", + "window": 0, + "total": 0, + "input": 0, + "output": 0 +} diff --git a/chat/testdata/overflow/code-context-window-exceeded.json b/chat/testdata/overflow/code-context-window-exceeded.json new file mode 100644 index 0000000..527f130 --- /dev/null +++ b/chat/testdata/overflow/code-context-window-exceeded.json @@ -0,0 +1,10 @@ +{ + "provider": "LiteLLM", + "source": "LiteLLM exception class name, no figures", + "raw": "litellm.ContextWindowExceededError: litellm.BadRequestError: the prompt does not fit", + "kind": "context", + "window": 0, + "total": 0, + "input": 0, + "output": 0 +} diff --git a/chat/testdata/overflow/gemini.json b/chat/testdata/overflow/gemini.json new file mode 100644 index 0000000..4ac880b --- /dev/null +++ b/chat/testdata/overflow/gemini.json @@ -0,0 +1,10 @@ +{ + "provider": "Gemini / Vertex", + "source": "Gemini API 400 INVALID_ARGUMENT message", + "raw": "The input token count (210000) exceeds the maximum number of tokens allowed (200000).", + "kind": "context", + "window": 200000, + "total": 0, + "input": 210000, + "output": 0 +} diff --git a/chat/testdata/overflow/generic-maximum-context-length.json b/chat/testdata/overflow/generic-maximum-context-length.json new file mode 100644 index 0000000..ddd4604 --- /dev/null +++ b/chat/testdata/overflow/generic-maximum-context-length.json @@ -0,0 +1,10 @@ +{ + "provider": "unknown", + "source": "generic wording with one figure: detected, but no window can be attributed", + "raw": "maximum context length is 8192 tokens", + "kind": "context", + "window": 0, + "total": 0, + "input": 0, + "output": 0 +} diff --git a/chat/testdata/overflow/llamacpp-localai-500-typed-stream.json b/chat/testdata/overflow/llamacpp-localai-500-typed-stream.json new file mode 100644 index 0000000..7bbe191 --- /dev/null +++ b/chat/testdata/overflow/llamacpp-localai-500-typed-stream.json @@ -0,0 +1,11 @@ +{ + "provider": "llama.cpp / LocalAI", + "source": "the same LocalAI HTTP 500 overflow in the text form of cogito's typed streaming error (in progress: status in *openai.APIError / *openai.RequestError)", + "raw": "failed to make a streaming decision after 3 attempts: localai stream: error, status code: 500, status: 500 Internal Server Error, message: rpc error: code = Internal desc = request (9739 tokens) exceeds the available context size (8192 tokens), try increasing it", + "kind": "context", + "status": 500, + "window": 8192, + "total": 0, + "input": 9739, + "output": 0 +} diff --git a/chat/testdata/overflow/llamacpp-localai-500.json b/chat/testdata/overflow/llamacpp-localai-500.json new file mode 100644 index 0000000..2e37789 --- /dev/null +++ b/chat/testdata/overflow/llamacpp-localai-500.json @@ -0,0 +1,11 @@ +{ + "provider": "llama.cpp / LocalAI", + "source": "LocalAI returns a llama.cpp overflow as HTTP 500 (gRPC Internal) on the streaming path, as cogito's streaming client reports it: a specific row is strong evidence and classifies at any status", + "raw": "failed to make a streaming decision after 3 attempts: localai stream: status 500: rpc error: code = Internal desc = request (9739 tokens) exceeds the available context size (8192 tokens), try increasing it", + "kind": "context", + "status": 500, + "window": 8192, + "total": 0, + "input": 9739, + "output": 0 +} diff --git a/chat/testdata/overflow/llamacpp-localai.json b/chat/testdata/overflow/llamacpp-localai.json new file mode 100644 index 0000000..292bedc --- /dev/null +++ b/chat/testdata/overflow/llamacpp-localai.json @@ -0,0 +1,10 @@ +{ + "provider": "llama.cpp / LocalAI", + "source": "llama.cpp tools/server/server-context.cpp, surfaced by LocalAI over gRPC and wrapped by cogito's tool-selection loop", + "raw": "failed to select tool: failed to pick tool: tool selection failed: failed to make a decision after 3 attempts: rpc error: code = Internal desc = request (9739 tokens) exceeds the available context size (8192 tokens), try increasing it", + "kind": "context", + "window": 8192, + "total": 0, + "input": 9739, + "output": 0 +} diff --git a/chat/testdata/overflow/llamacpp-stream-chunk.json b/chat/testdata/overflow/llamacpp-stream-chunk.json new file mode 100644 index 0000000..ef94e25 --- /dev/null +++ b/chat/testdata/overflow/llamacpp-stream-chunk.json @@ -0,0 +1,10 @@ +{ + "provider": "llama.cpp / LocalAI", + "source": "llama.cpp error delivered as an SSE error chunk mid-stream: no HTTP status in the chain", + "raw": "failed to make a streaming decision after 3 attempts: localai stream: request (368203 tokens) exceeds the available context size (262144 tokens), try increasing it", + "kind": "context", + "window": 262144, + "total": 0, + "input": 368203, + "output": 0 +} diff --git a/chat/testdata/overflow/negative-auth.json b/chat/testdata/overflow/negative-auth.json new file mode 100644 index 0000000..f4e2992 --- /dev/null +++ b/chat/testdata/overflow/negative-auth.json @@ -0,0 +1,10 @@ +{ + "provider": "OpenAI-compatible", + "source": "authentication failure", + "raw": "error, status code: 401, status: 401 Unauthorized, message: Incorrect API key provided", + "kind": "none", + "window": 0, + "total": 0, + "input": 0, + "output": 0 +} diff --git a/chat/testdata/overflow/negative-generic-server-error.json b/chat/testdata/overflow/negative-generic-server-error.json new file mode 100644 index 0000000..ae888b6 --- /dev/null +++ b/chat/testdata/overflow/negative-generic-server-error.json @@ -0,0 +1,11 @@ +{ + "provider": "unknown", + "source": "HTTP 500 whose text only has a generic context phrase (no provider row): the status gate on weak evidence makes it KindNone", + "raw": "localai stream: status 500: internal error while computing the context window", + "kind": "none", + "status": 500, + "window": 0, + "total": 0, + "input": 0, + "output": 0 +} diff --git a/chat/testdata/overflow/negative-max-tokens-validation.json b/chat/testdata/overflow/negative-max-tokens-validation.json new file mode 100644 index 0000000..622c766 --- /dev/null +++ b/chat/testdata/overflow/negative-max-tokens-validation.json @@ -0,0 +1,10 @@ +{ + "provider": "OpenAI", + "source": "max_tokens parameter validation, not about size (output-cap guard)", + "raw": "Invalid 'max_tokens': integer above maximum value. Expected a value <= 16384, but got 100000 instead.", + "kind": "none", + "window": 0, + "total": 0, + "input": 0, + "output": 0 +} diff --git a/chat/testdata/overflow/negative-rate-limit.json b/chat/testdata/overflow/negative-rate-limit.json new file mode 100644 index 0000000..d51a219 --- /dev/null +++ b/chat/testdata/overflow/negative-rate-limit.json @@ -0,0 +1,11 @@ +{ + "provider": "OpenAI-compatible", + "source": "rate limit (429) with a generic context wording in the body: weak evidence at a status other than 400/413 is not an overflow", + "raw": "localai stream: status 429: Rate limit exceeded for this context window, retry in 20s", + "kind": "none", + "status": 429, + "window": 0, + "total": 0, + "input": 0, + "output": 0 +} diff --git a/chat/testdata/overflow/openai.json b/chat/testdata/overflow/openai.json new file mode 100644 index 0000000..f7aa355 --- /dev/null +++ b/chat/testdata/overflow/openai.json @@ -0,0 +1,10 @@ +{ + "provider": "OpenAI", + "source": "OpenAI chat/completions 400 message, as decoded into *openai.APIError", + "raw": "This model's maximum context length is 8192 tokens. However, your messages resulted in 9739 tokens. Please reduce the length of the messages.", + "kind": "context", + "window": 8192, + "total": 0, + "input": 9739, + "output": 0 +} diff --git a/chat/testdata/overflow/regolo-vllm-budget-nonstream.json b/chat/testdata/overflow/regolo-vllm-budget-nonstream.json new file mode 100644 index 0000000..713151b --- /dev/null +++ b/chat/testdata/overflow/regolo-vllm-budget-nonstream.json @@ -0,0 +1,10 @@ +{ + "provider": "regolo (LiteLLM + vLLM)", + "source": "live capture from regolo (LiteLLM in front of vLLM, model glm5.2), HTTP 400; raw is the body decoded by cogito's non-streaming client into *openai.APIError (its Error() is the message; HTTPStatusCode is not set), wrapped by cogito's decision loop", + "raw": "failed to select tool: failed to pick tool: tool selection failed: failed to make a decision after 3 attempts: regolo.BadRequestError: Hosted_vllmException - {\"object\":\"error\",\"message\":\"Requested token count exceeds the model's maximum context length of 210000 tokens. You requested a total of 210123 tokens: 133 tokens from the input messages and 209990 tokens for the completion. Please reduce the number of tokens in the input messages or the completion to fit within the limit.\",\"type\":\"BadRequestError\",\"param\":null,\"code\":400}No fallback model group found for original model_group=glm5.2.. Received Model Group=glm5.2\nAvailable Model Group Fallbacks=None\nregolo.BadRequestError: Hosted_vllmException - {\"object\":\"error\",\"message\":\"Requested token count exceeds the model's maximum context length of 210000 tokens. You requested a total of 210123 tokens: 133 tokens from the input messages and 209990 tokens for the completion. Please reduce the number of tokens in the input messages or the completion to fit within the limit.\",\"type\":\"BadRequestError\",\"param\":null,\"code\":400}No fallback model group found for original model_group=glm5.2. Retried: 2 times", + "kind": "budget", + "window": 210000, + "total": 210123, + "input": 133, + "output": 209990 +} diff --git a/chat/testdata/overflow/regolo-vllm-budget-stream.json b/chat/testdata/overflow/regolo-vllm-budget-stream.json new file mode 100644 index 0000000..628e0f7 --- /dev/null +++ b/chat/testdata/overflow/regolo-vllm-budget-stream.json @@ -0,0 +1,10 @@ +{ + "provider": "regolo (LiteLLM + vLLM)", + "source": "live capture from regolo (LiteLLM in front of vLLM, model glm5.2), HTTP 400; raw is the body as cogito's streaming client reports it (fmt.Errorf(\"localai stream: status %d: %s\")), wrapped by cogito's tool-selection loop", + "raw": "failed to select tool: failed to pick tool: tool selection failed: failed to make a streaming decision after 3 attempts: localai stream: status 400: {\"error\":{\"message\":\"regolo.BadRequestError: Hosted_vllmException - {\\\"object\\\":\\\"error\\\",\\\"message\\\":\\\"Requested token count exceeds the model's maximum context length of 210000 tokens. You requested a total of 210123 tokens: 133 tokens from the input messages and 209990 tokens for the completion. Please reduce the number of tokens in the input messages or the completion to fit within the limit.\\\",\\\"type\\\":\\\"BadRequestError\\\",\\\"param\\\":null,\\\"code\\\":400}No fallback model group found for original model_group=glm5.2.. Received Model Group=glm5.2\\nAvailable Model Group Fallbacks=None\\nregolo.BadRequestError: Hosted_vllmException - {\\\"object\\\":\\\"error\\\",\\\"message\\\":\\\"Requested token count exceeds the model's maximum context length of 210000 tokens. You requested a total of 210123 tokens: 133 tokens from the input messages and 209990 tokens for the completion. Please reduce the number of tokens in the input messages or the completion to fit within the limit.\\\",\\\"type\\\":\\\"BadRequestError\\\",\\\"param\\\":null,\\\"code\\\":400}No fallback model group found for original model_group=glm5.2. Retried: 2 times\",\"type\":null,\"param\":null,\"code\":\"400\"}}", + "kind": "budget", + "window": 210000, + "total": 210123, + "input": 133, + "output": 209990 +} diff --git a/chat/testdata/overflow/regolo-vllm-output-cap-nonstream.json b/chat/testdata/overflow/regolo-vllm-output-cap-nonstream.json new file mode 100644 index 0000000..08c31dd --- /dev/null +++ b/chat/testdata/overflow/regolo-vllm-output-cap-nonstream.json @@ -0,0 +1,10 @@ +{ + "provider": "regolo (LiteLLM + vLLM)", + "source": "live capture from regolo (LiteLLM in front of vLLM, model glm5.2), HTTP 400; raw is the body decoded by cogito's non-streaming client into *openai.APIError (its Error() is the message; HTTPStatusCode is not set), wrapped by cogito's decision loop", + "raw": "failed to select tool: failed to pick tool: tool selection failed: failed to make a decision after 3 attempts: regolo.BadRequestError: Hosted_vllmException - {\"object\":\"error\",\"message\":\"max_completion_tokens is too large: 250000.This model supports at most 210000 completion tokens.\",\"type\":\"BadRequestError\",\"param\":null,\"code\":400}No fallback model group found for original model_group=glm5.2.. Received Model Group=glm5.2\nAvailable Model Group Fallbacks=None\nregolo.BadRequestError: Hosted_vllmException - {\"object\":\"error\",\"message\":\"max_completion_tokens is too large: 250000.This model supports at most 210000 completion tokens.\",\"type\":\"BadRequestError\",\"param\":null,\"code\":400}No fallback model group found for original model_group=glm5.2. Retried: 2 times", + "kind": "output_cap", + "window": 210000, + "total": 0, + "input": 0, + "output": 250000 +} diff --git a/chat/testdata/overflow/regolo-vllm-output-cap-stream.json b/chat/testdata/overflow/regolo-vllm-output-cap-stream.json new file mode 100644 index 0000000..dc7836d --- /dev/null +++ b/chat/testdata/overflow/regolo-vllm-output-cap-stream.json @@ -0,0 +1,10 @@ +{ + "provider": "regolo (LiteLLM + vLLM)", + "source": "live capture from regolo (LiteLLM in front of vLLM, model glm5.2), HTTP 400; raw is the body as cogito's streaming client reports it (fmt.Errorf(\"localai stream: status %d: %s\")), wrapped by cogito's tool-selection loop", + "raw": "failed to select tool: failed to pick tool: tool selection failed: failed to make a streaming decision after 3 attempts: localai stream: status 400: {\"error\":{\"message\":\"regolo.BadRequestError: Hosted_vllmException - {\\\"object\\\":\\\"error\\\",\\\"message\\\":\\\"max_completion_tokens is too large: 250000.This model supports at most 210000 completion tokens.\\\",\\\"type\\\":\\\"BadRequestError\\\",\\\"param\\\":null,\\\"code\\\":400}No fallback model group found for original model_group=glm5.2.. Received Model Group=glm5.2\\nAvailable Model Group Fallbacks=None\\nregolo.BadRequestError: Hosted_vllmException - {\\\"object\\\":\\\"error\\\",\\\"message\\\":\\\"max_completion_tokens is too large: 250000.This model supports at most 210000 completion tokens.\\\",\\\"type\\\":\\\"BadRequestError\\\",\\\"param\\\":null,\\\"code\\\":400}No fallback model group found for original model_group=glm5.2. Retried: 2 times\",\"type\":null,\"param\":null,\"code\":\"400\"}}", + "kind": "output_cap", + "window": 210000, + "total": 0, + "input": 0, + "output": 250000 +} diff --git a/chat/testdata/overflow/vllm-current-budget.json b/chat/testdata/overflow/vllm-current-budget.json new file mode 100644 index 0000000..1aa541e --- /dev/null +++ b/chat/testdata/overflow/vllm-current-budget.json @@ -0,0 +1,10 @@ +{ + "provider": "vLLM", + "source": "vLLM source, vllm/renderers/params.py: the same wording when the prompt fits but prompt + max_tokens does not", + "raw": "This model's maximum context length is 8192 tokens. However, you requested 4096 output tokens and your prompt contains 5000 input tokens, for a total of 9096 tokens. Please reduce the length of the input prompt or the number of requested output tokens.", + "kind": "budget", + "window": 8192, + "total": 9096, + "input": 5000, + "output": 4096 +} diff --git a/chat/testdata/overflow/vllm-current-typed-stream.json b/chat/testdata/overflow/vllm-current-typed-stream.json new file mode 100644 index 0000000..6ff7901 --- /dev/null +++ b/chat/testdata/overflow/vllm-current-typed-stream.json @@ -0,0 +1,11 @@ +{ + "provider": "vLLM", + "source": "vLLM current wording in the text form of cogito's typed streaming error (in progress): \"localai stream: error, status code: 400, status: 400 Bad Request, message: ...\"", + "raw": "localai stream: error, status code: 400, status: 400 Bad Request, message: This model's maximum context length is 8192 tokens. However, you requested 512 output tokens and your prompt contains 9739 input tokens, for a total of 10251 tokens. Please reduce the length of the input prompt or the number of requested output tokens.", + "kind": "context", + "status": 400, + "window": 8192, + "total": 10251, + "input": 9739, + "output": 512 +} diff --git a/chat/testdata/overflow/vllm-current.json b/chat/testdata/overflow/vllm-current.json new file mode 100644 index 0000000..701861a --- /dev/null +++ b/chat/testdata/overflow/vllm-current.json @@ -0,0 +1,10 @@ +{ + "provider": "vLLM", + "source": "vLLM source, vllm/renderers/params.py (current OpenAI-compatible server), as surfaced by cogito's streaming client", + "raw": "localai stream: status 400: This model's maximum context length is 8192 tokens. However, you requested 512 output tokens and your prompt contains 9739 input tokens, for a total of 10251 tokens. Please reduce the length of the input prompt or the number of requested output tokens.", + "kind": "context", + "window": 8192, + "total": 10251, + "input": 9739, + "output": 512 +} diff --git a/chat/testdata/overflow/vllm-older.json b/chat/testdata/overflow/vllm-older.json new file mode 100644 index 0000000..8ee48e5 --- /dev/null +++ b/chat/testdata/overflow/vllm-older.json @@ -0,0 +1,10 @@ +{ + "provider": "vLLM", + "source": "vLLM source, older OpenAI-compatible server (serving_engine.py)", + "raw": "This model's maximum context length is 8192 tokens. However, you requested 10251 tokens (9739 in the messages, 512 in the completion). Please reduce the length of the messages or completion.", + "kind": "context", + "window": 8192, + "total": 10251, + "input": 9739, + "output": 512 +} diff --git a/config/config.go b/config/config.go index 0e21925..673a7fe 100644 --- a/config/config.go +++ b/config/config.go @@ -5,6 +5,7 @@ import ( "os" "path/filepath" "reflect" + "regexp" "time" "dario.cat/mergo" @@ -218,6 +219,7 @@ func LoadWith(o LoadOptions) types.Config { applyOverrides(&cfg, o.Overrides) cfg = withDefaults(cfg) + cfg.Compaction.OverflowPatterns = validOverflowPatterns(cfg.Compaction.OverflowPatterns) // Carry the override (not the resolved root) so consumers keep resolving it // through plugin.BaseDirIn / config.WritablePathIn: with no override those @@ -453,3 +455,25 @@ func withDefaults(cfg types.Config) types.Config { } return cfg } + +// validOverflowPatterns drops the compaction.overflow_patterns rows that cannot +// work: a regex that does not compile, or a kind nib does not know. Each one is +// reported on stderr and skipped. A bad pattern must never fail startup: the +// built-in patterns still cover the backend, only the user's addition is lost. +func validOverflowPatterns(in []types.OverflowPattern) []types.OverflowPattern { + var out []types.OverflowPattern + for _, p := range in { + switch p.Kind { + case "context", "budget", "output_cap": + default: + fmt.Fprintf(os.Stderr, "nib: compaction.overflow_patterns: %q: unknown kind %q (want context, budget or output_cap), skipped\n", p.Name, p.Kind) + continue + } + if _, err := regexp.Compile(p.Regex); err != nil { + fmt.Fprintf(os.Stderr, "nib: compaction.overflow_patterns: %q: %v, skipped\n", p.Name, err) + continue + } + out = append(out, p) + } + return out +} diff --git a/config/config_test.go b/config/config_test.go index 641f4dc..c1d8edd 100644 --- a/config/config_test.go +++ b/config/config_test.go @@ -35,7 +35,7 @@ func TestWithDefaultsKeepsUserValues(t *testing.T) { ReserveTokens: 512, SummaryMaxTokens: 2048, }} got := withDefaults(in) - if got.Compaction != in.Compaction { + if !reflect.DeepEqual(got.Compaction, in.Compaction) { t.Fatalf("withDefaults overrode user values: %+v", got.Compaction) } } @@ -660,3 +660,16 @@ func TestWritablePathInPrefersInjectedRoot(t *testing.T) { t.Fatalf("WritablePathIn = %q, want %q", got, want) } } + +// An invalid compaction.overflow_patterns row is skipped at load time; it +// never fails startup, and the valid rows survive. +func TestValidOverflowPatternsSkipsInvalidRows(t *testing.T) { + got := validOverflowPatterns([]types.OverflowPattern{ + {Name: "broken", Kind: "context", Regex: `(unclosed`}, + {Name: "bad-kind", Kind: "sideways", Regex: `x`}, + {Name: "ok", Kind: "budget", Regex: `(?P\d+) \+ (?P\d+) > (?P\d+)`}, + }) + if len(got) != 1 || got[0].Name != "ok" { + t.Fatalf("valid rows = %+v, want only \"ok\"", got) + } +} diff --git a/selfdoc/README.md b/selfdoc/README.md index 1505974..1c5f36d 100644 --- a/selfdoc/README.md +++ b/selfdoc/README.md @@ -575,8 +575,27 @@ log_level: error # prompt is fitted into the window minus that reservation. If the backend # still rejects it as too large, nib summarizes fewer of the older turns and # keeps the rest verbatim, instead of cutting the text it summarizes. +# +# nib reads a backend's "request too large" error to decide what to do: +# compact when the prompt itself does not fit (kind "context"), lower the +# requested output when only prompt + output does not fit ("budget"), or +# lower the output cap when the output alone is above the model's maximum +# ("output_cap"). It knows the wording of llama.cpp/LocalAI, vLLM, OpenAI, +# Anthropic and Gemini. overflow_patterns teaches it another backend without a +# rebuild: the rows are tried before the built-in ones, first match wins. A +# regex states its figures with the named groups window, total, input and +# output. A row with an invalid regex or kind is logged and skipped. A +# backend-specific wording (built-in or from this list) counts at any HTTP +# status, because LocalAI reports an overflow as HTTP 500; a generic phrase +# such as "context window" counts only with status 400, 413 or none. An error +# with HTTP status 400 or 413 that no pattern recognises is logged at debug +# level, so you can copy its text from the log. compaction: summary_max_tokens: 16384 + overflow_patterns: + - name: my-backend + kind: context + regex: 'prompt has (?P\d+) tokens, limit is (?P\d+)' # Optional: external MCP servers mcp_servers: diff --git a/types/config.go b/types/config.go index 3cc774a..fca9c83 100644 --- a/types/config.go +++ b/types/config.go @@ -59,6 +59,22 @@ type CompactionConfig struct { // the head is saved and the summary includes an artifact://N reference. // Set to true to disable — compaction runs as before, nothing saved. DisableArtifactSpill bool `yaml:"disable_artifact_spill"` + // OverflowPatterns teach nib a backend's context-overflow wording without a + // rebuild. They are tried before the built-in patterns, first match wins. + // An invalid regex or kind is logged and skipped; it never fails startup. + OverflowPatterns []OverflowPattern `yaml:"overflow_patterns"` +} + +// OverflowPattern is one user-supplied row of the overflow error catalog. +// Regex figures come from the named groups window, total, input and output, +// so the pattern states what each number means. +type OverflowPattern struct { + Name string `yaml:"name"` + // Kind is "context" (the prompt does not fit: compact), "budget" (prompt + // plus requested output does not fit: lower max_tokens) or "output_cap" + // (the requested output alone is above the model's maximum). + Kind string `yaml:"kind"` + Regex string `yaml:"regex"` } // ToolOutputPruningConfig controls replacing stale or oversized tool results From fd573612e01548d80669e0c86564f2a43cb7100a Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Fri, 25 Sep 2026 21:55:50 +0000 Subject: [PATCH 10/16] feat(chat): recover from overflow with trim, fallback and kind-specific retries Overflow recovery in SendMessage now acts on the overflow kind: - context: run iterativeTrim, then basicCompact when trimming cannot make the history fit. The whole chain is one recovery attempt, and the retry is announced only when a step changed the history. - budget: do not compact. Retry once with the output reservation set to window - input - outputSafetyMargin, from the backend's figures, through a one-turn override that requestLimits reads. When that is below minOutputTokens, or the lowered retry overflows again, only a smaller prompt helps, so the turn goes through the context chain. - output_cap: do not compact. Lower the session's output cap (and the client's) to the stated maximum, never raising it, and retry once. When the retry after recovery still overflows, the rollback keeps the compaction and removes only the failed turn's own messages. Before, it restored the pre-turn history, so the next turn overflowed again and the session was stuck. Changes to existing tests: the overflow fixtures seeded a history of a few bytes, which a summary cannot shrink, so the validated chain rejected it. They now seed longer replies. The tests that need the whole chain to fail keep the tiny history (tinyHistory). The regolo regression test runs over HTTP at 1:10 scale: vLLM window 21000, LiteLLM /model/info 9600/9600. Co-Authored-By: Claude Opus 5.5 --- README.md | 7 +- chat/modellimits.go | 60 ++- chat/overflow_e2e_test.go | 12 +- chat/overflow_retry_test.go | 25 +- chat/overflow_trim_integration_test.go | 534 +++++++++++++++++++++++++ chat/session.go | 212 ++++++++-- selfdoc/README.md | 7 +- 7 files changed, 814 insertions(+), 43 deletions(-) create mode 100644 chat/overflow_trim_integration_test.go diff --git a/README.md b/README.md index 1c5f36d..53e655e 100644 --- a/README.md +++ b/README.md @@ -580,7 +580,12 @@ log_level: error # compact when the prompt itself does not fit (kind "context"), lower the # requested output when only prompt + output does not fit ("budget"), or # lower the output cap when the output alone is above the model's maximum -# ("output_cap"). It knows the wording of llama.cpp/LocalAI, vLLM, OpenAI, +# ("output_cap"), and then retry the turn once. When the summary alone does +# not make the history fit, nib stubs old tool outputs, then drops older turns +# (saved as an artifact:// the model can read), and as a last resort compacts +# without the model. If the retry still does not fit, the turn fails, but the +# compaction is kept, so the next turn does not overflow again. It knows the +# wording of llama.cpp/LocalAI, vLLM, OpenAI, # Anthropic and Gemini. overflow_patterns teaches it another backend without a # rebuild: the rows are tried before the built-in ones, first match wins. A # regex states its figures with the named groups window, total, input and diff --git a/chat/modellimits.go b/chat/modellimits.go index 76ea31a..8c648d8 100644 --- a/chat/modellimits.go +++ b/chat/modellimits.go @@ -95,9 +95,65 @@ func factoryOutputCap(llm cogito.LLM, provider types.ModelProviderConfig, store // requestLimits reports the output cap and context window the turn's requests // are clamped against (see clampOutputTokens), read under one lock so a model -// switch cannot pair one model's cap with another's window. +// switch cannot pair one model's cap with another's window. A one-turn +// override (turnOutputCap) lowers the cap, and never raises it. func (s *Session) requestLimits() (cap, window int) { s.modelMu.RLock() defer s.modelMu.RUnlock() - return s.outputCap, s.windowLocked() + cap = s.outputCap + if o := s.turnOutputCap; o > 0 && (cap <= 0 || o < cap) { + cap = o + } + return cap, s.windowLocked() +} + +// setTurnOutputCap sets (or, with 0, clears) the current turn's output cap +// override. +func (s *Session) setTurnOutputCap(v int) { + s.modelMu.Lock() + s.turnOutputCap = v + s.modelMu.Unlock() +} + +// lowerOutputCap lowers the output cap of model to max after the backend +// said the requested output is above what it allows. It never raises the +// cap. It reports whether the cap changed, so the caller retries only a +// request that will differ. The client's own cap is lowered too, because it +// is what a request carries when the window is unknown and nothing clamps. +func (s *Session) lowerOutputCap(max int, model string) bool { + if max <= 0 { + return false + } + s.modelMu.Lock() + if s.llmModel != model || (s.outputCap > 0 && s.outputCap <= max) { + s.modelMu.Unlock() + return false + } + s.outputCap = max + llm := s.llm + s.modelMu.Unlock() + if setter, ok := llm.(maxTokensSetter); ok { + setter.SetMaxTokens(max) + } + return true +} + +// budgetRetryOutput is the output reservation that fits a budget overflow: +// the window less the input the backend counted, less outputSafetyMargin. The +// backend's count is exact, which nib's estimate is not. It reports false when +// the figures are missing or the result is below minOutputTokens: then the +// prompt itself leaves no useful room, and only a smaller prompt helps. +func budgetRetryOutput(info overflowInfo) (int, bool) { + input := info.Input + if input <= 0 && info.Total > 0 && info.Output > 0 { + input = info.Total - info.Output + } + if info.Window <= 0 || input <= 0 { + return 0, false + } + out := info.Window - input - outputSafetyMargin + if out < minOutputTokens { + return 0, false + } + return out, true } diff --git a/chat/overflow_e2e_test.go b/chat/overflow_e2e_test.go index cae229c..69f0550 100644 --- a/chat/overflow_e2e_test.go +++ b/chat/overflow_e2e_test.go @@ -84,8 +84,8 @@ func TestOverflowRecoveryOverHTTP(t *testing.T) { // Seed enough history that compaction has a head to summarise. s.fragment = s.fragment. - AddMessage("user", "u1").AddMessage("assistant", "a1"). - AddMessage("user", "u2").AddMessage("assistant", "a2") + AddMessage("user", "u1").AddMessage("assistant", "a1 "+overflowFiller). + AddMessage("user", "u2").AddMessage("assistant", "a2 "+overflowFiller) got, err := s.SendMessage("what changed?") if err != nil { @@ -249,8 +249,8 @@ func TestOverflowRetryKeepsTheSystemPrompt(t *testing.T) { // FIRST (the first turn's own add), then history accumulates after it. s.ensureSystemPrompt() s.fragment = s.fragment. - AddMessage("user", "u1").AddMessage("assistant", "a1"). - AddMessage("user", "u2").AddMessage("assistant", "a2") + AddMessage("user", "u1").AddMessage("assistant", "a1 "+overflowFiller). + AddMessage("user", "u2").AddMessage("assistant", "a2 "+overflowFiller) if got := s.fragment.Messages[0].Role; got != "system" { t.Fatalf("seeded fragment starts with %q, not the system prompt: the bug would be hidden", got) } @@ -336,8 +336,8 @@ func TestOverflowRecoveryOverStreamedErrorChunk(t *testing.T) { defer s.Close() s.fragment = s.fragment. - AddMessage("user", "u1").AddMessage("assistant", "a1"). - AddMessage("user", "u2").AddMessage("assistant", "a2") + AddMessage("user", "u1").AddMessage("assistant", "a1 "+overflowFiller). + AddMessage("user", "u2").AddMessage("assistant", "a2 "+overflowFiller) got, err := s.SendMessage("what changed?") if err != nil { diff --git a/chat/overflow_retry_test.go b/chat/overflow_retry_test.go index 8280a6c..177cb9f 100644 --- a/chat/overflow_retry_test.go +++ b/chat/overflow_retry_test.go @@ -100,14 +100,32 @@ func newOverflowSession(t *testing.T, llm cogito.LLM) *Session { } // Enough history that compaction has something to summarise; without a // head to compact the retry is correctly skipped and the test would prove - // nothing. + // nothing. The replies are long because overflow recovery validates its + // result: a summary larger than the history it replaces is rejected, and + // a real overflow never comes from a few bytes of history. + s.fragment = cogito.NewFragment( + openai.ChatCompletionMessage{Role: "user", Content: "u1"}, + openai.ChatCompletionMessage{Role: "assistant", Content: "a1 " + overflowFiller}, + openai.ChatCompletionMessage{Role: "user", Content: "u2"}, + openai.ChatCompletionMessage{Role: "assistant", Content: "a2 " + overflowFiller}, + ) + return s +} + +// overflowFiller makes a seeded reply large enough that compacting it shrinks +// the history. +var overflowFiller = strings.Repeat("earlier detail ", 150) + +// tinyHistory seeds a history no step of overflow recovery can shrink: a +// summary, a truncation marker or a basic-compaction marker is larger than +// it. With the summariser failing too, the whole recovery chain fails. +func tinyHistory(s *Session) { s.fragment = cogito.NewFragment( openai.ChatCompletionMessage{Role: "user", Content: "u1"}, openai.ChatCompletionMessage{Role: "assistant", Content: "a1"}, openai.ChatCompletionMessage{Role: "user", Content: "u2"}, openai.ChatCompletionMessage{Role: "assistant", Content: "a2"}, ) - return s } // The headline: an overflow no longer ends the turn. @@ -198,6 +216,7 @@ func TestOverflowErrorReflectsWhetherCompactionWasTried(t *testing.T) { // retrying. Nothing has been done about the size, so the original // advice is still the right advice. s := newOverflowSession(t, &summaryFailingLLM{overflowLLM: overflowLLM{failures: 99}}) + tinyHistory(s) _, err := s.SendMessage("what changed?") if err == nil { @@ -432,6 +451,7 @@ func TestOverflowDoesNotAnnounceARetryItWillNotMake(t *testing.T) { t.Run("compaction fails", func(t *testing.T) { rec := &retryStatusRecorder{} s := newOverflowSession(t, &summaryFailingLLM{overflowLLM: overflowLLM{failures: 99}}) + tinyHistory(s) s.callbacks = rec.callbacks() if _, err := s.SendMessage("what changed?"); err == nil { @@ -488,6 +508,7 @@ func (f *summaryFailingLLM) CreateChatCompletion(ctx context.Context, req openai func TestOverflowReportsOriginalErrorWhenCompactionFails(t *testing.T) { llm := &summaryFailingLLM{overflowLLM: overflowLLM{failures: 99}} s := newOverflowSession(t, llm) + tinyHistory(s) _, err := s.SendMessage("what changed?") if err == nil { diff --git a/chat/overflow_trim_integration_test.go b/chat/overflow_trim_integration_test.go new file mode 100644 index 0000000..b0532aa --- /dev/null +++ b/chat/overflow_trim_integration_test.go @@ -0,0 +1,534 @@ +package chat + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "io" + "net/http" + "net/http/httptest" + "os" + "strings" + "sync" + "sync/atomic" + "testing" + + "github.com/mudler/cogito" + "github.com/mudler/nib/types" + "github.com/mudler/xlog" + openai "github.com/sashabaranov/go-openai" +) + +// withOverflowTrim replaces the iterativeTrim step of overflow recovery for +// one test, so the basicCompact fallback can be reached on demand. +func withOverflowTrim(t *testing.T, f func(*Session, context.Context) error) { + t.Helper() + prev := overflowTrim + overflowTrim = f + t.Cleanup(func() { overflowTrim = prev }) +} + +// bigHistorySession is newOverflowSession with a history that basicCompact can +// shrink: short user messages and long assistant replies, which it elides. +func bigHistorySession(t *testing.T, llm cogito.LLM) *Session { + t.Helper() + s := newOverflowSession(t, llm) + long := strings.Repeat("assistant detail ", 200) + s.fragment = cogito.NewFragment( + openai.ChatCompletionMessage{Role: "user", Content: "u1"}, + openai.ChatCompletionMessage{Role: "assistant", Content: "a1 " + long}, + openai.ChatCompletionMessage{Role: "user", Content: "u2"}, + openai.ChatCompletionMessage{Role: "assistant", Content: "a2 " + long}, + openai.ChatCompletionMessage{Role: "user", Content: "u3"}, + openai.ChatCompletionMessage{Role: "assistant", Content: "a3 " + long}, + ) + return s +} + +func (r *retryStatusRecorder) saw(sub string) bool { + r.mu.Lock() + defer r.mu.Unlock() + for _, l := range r.lines { + if strings.Contains(l, sub) { + return true + } + } + return false +} + +func fragmentHas(s *Session, pred func(openai.ChatCompletionMessage) bool) bool { + s.historyMu.Lock() + defer s.historyMu.Unlock() + for _, m := range s.fragment.Messages { + if pred(m) { + return true + } + } + return false +} + +// Step 1 of iterativeTrim (LLM compaction) is enough: the turn retries. +func TestOverflowChainTrimSucceedsAtStepOne(t *testing.T) { + rec := &retryStatusRecorder{} + llm := &overflowLLM{failures: 1} + s := newOverflowSession(t, llm) + s.callbacks = rec.callbacks() + + if _, err := s.SendMessage("what changed?"); err != nil { + t.Fatalf("SendMessage: %v", err) + } + if s.overflowRetries() != 1 { + t.Fatalf("overflow retries = %d, want 1", s.overflowRetries()) + } + llm.mu.Lock() + asks := llm.asks + llm.mu.Unlock() + if asks == 0 { + t.Fatal("no summary was requested: the retry did not come from LLM compaction") + } + if !rec.saw("compacting and retrying") || rec.saw("basic fallback") { + t.Fatalf("want the compaction announcement, not the fallback one; statuses: %v", rec.lines) + } +} + +// iterativeTrim fails; basicCompact runs and succeeds; the turn retries with +// the fallback announcement. +func TestOverflowChainFallsBackToBasicCompact(t *testing.T) { + var trims atomic.Int32 + withOverflowTrim(t, func(*Session, context.Context) error { + trims.Add(1) + return errors.New("iterative trim: nothing fits") + }) + rec := &retryStatusRecorder{} + s := bigHistorySession(t, &overflowLLM{failures: 1}) + s.callbacks = rec.callbacks() + + if _, err := s.SendMessage("what changed?"); err != nil { + t.Fatalf("SendMessage: %v", err) + } + if trims.Load() != 1 { + t.Fatalf("iterativeTrim ran %d times, want 1", trims.Load()) + } + if s.overflowRetries() != 1 { + t.Fatalf("overflow retries = %d, want 1", s.overflowRetries()) + } + if !rec.saw("using basic fallback compaction and retrying") { + t.Fatalf("the fallback retry was not announced; statuses: %v", rec.lines) + } + if !fragmentHas(s, func(m openai.ChatCompletionMessage) bool { return strings.HasPrefix(m.Content, basicCompactMarker) }) { + t.Fatal("the fragment carries no basic-compaction marker") + } + rec.mu.Lock() + last := rec.lines[len(rec.lines)-1] + rec.mu.Unlock() + if last != retryResumeStatus { + t.Fatalf("last status = %q, want %q", last, retryResumeStatus) + } +} + +// Both steps fail: the turn is rolled back and the overflow surfaces. +func TestOverflowChainBothFailRollsBack(t *testing.T) { + s := newOverflowSession(t, &summaryFailingLLM{overflowLLM: overflowLLM{failures: 99}}) + tinyHistory(s) + s.ensureSystemPrompt() + before := append([]openai.ChatCompletionMessage(nil), s.fragment.Messages...) + + _, err := s.SendMessage("what changed?") + if err == nil || !isContextOverflow(err) { + t.Fatalf("err = %v, want the overflow", err) + } + if s.overflowRetries() != 0 { + t.Fatalf("overflow retries = %d, want 0", s.overflowRetries()) + } + got := s.fragment.Messages + if len(got) != len(before) { + t.Fatalf("fragment has %d messages, want the %d from before the turn", len(got), len(before)) + } + for i := range got { + if got[i].Role != before[i].Role || got[i].Content != before[i].Content { + t.Fatalf("message %d changed: %+v", i, got[i]) + } + } +} + +// The whole chain is one recovery attempt. +func TestOverflowChainCountsAsOneRetry(t *testing.T) { + var trims atomic.Int32 + withOverflowTrim(t, func(*Session, context.Context) error { + trims.Add(1) + return errors.New("iterative trim: nothing fits") + }) + s := bigHistorySession(t, &overflowLLM{failures: 99}) + + if _, err := s.SendMessage("what changed?"); err == nil { + t.Fatal("expected the turn to fail") + } + if s.overflowRetries() != 1 { + t.Fatalf("overflow retries = %d, want 1 for the whole chain", s.overflowRetries()) + } + if trims.Load() != 1 { + t.Fatalf("iterativeTrim ran %d times, want 1", trims.Load()) + } +} + +// A cancelled turn never enters the chain. +func TestOverflowChainSkipsCancelledTurn(t *testing.T) { + var trims atomic.Int32 + withOverflowTrim(t, func(*Session, context.Context) error { + trims.Add(1) + return nil + }) + llm := &interruptingOverflowLLM{overflowLLM: overflowLLM{failures: 99}} + s := newOverflowSession(t, llm) + llm.s = s + + if _, err := s.SendMessage("what changed?"); err == nil { + t.Fatal("expected the interrupted turn to fail") + } + if trims.Load() != 0 || s.overflowRetries() != 0 { + t.Fatalf("trim ran %d times, retries = %d on a cancelled turn", trims.Load(), s.overflowRetries()) + } +} + +// pairedOverflowLLM overflows every turn request until ok is set, and records +// the messages of every turn request. +type pairedOverflowLLM struct { + overflowLLM + ok atomic.Bool + reqs [][]openai.ChatCompletionMessage +} + +func (p *pairedOverflowLLM) CreateChatCompletion(ctx context.Context, req openai.ChatCompletionRequest) (cogito.LLMReply, cogito.LLMUsage, error) { + if n := len(req.Messages); n > 0 && strings.HasPrefix(req.Messages[n-1].Content, compactInstruction) { + return p.overflowLLM.CreateChatCompletion(ctx, req) + } + p.mu.Lock() + p.reqs = append(p.reqs, append([]openai.ChatCompletionMessage(nil), req.Messages...)) + p.mu.Unlock() + if p.ok.Load() { + return replyWith("fine"), cogito.LLMUsage{PromptTokens: 10, TotalTokens: 10}, nil + } + return p.overflowLLM.CreateChatCompletion(ctx, req) +} + +// Recovery compacts, the retry still overflows, the turn fails. The next turn +// starts from the compacted history, not the pre-turn one, with no orphans. +func TestOverflowFailedRetryKeepsTheCompaction(t *testing.T) { + llm := &pairedOverflowLLM{overflowLLM: overflowLLM{failures: 1 << 20}} + s := newOverflowSession(t, llm) + s.fragment = cogito.NewFragment( + openai.ChatCompletionMessage{Role: "user", Content: "u1"}, + openai.ChatCompletionMessage{Role: "assistant", ToolCalls: []openai.ToolCall{basicToolCall("c1", "read", `{"path":"/a"}`)}}, + openai.ChatCompletionMessage{Role: "tool", ToolCallID: "c1", Content: strings.Repeat("x", 4000)}, + openai.ChatCompletionMessage{Role: "assistant", Content: "a1"}, + openai.ChatCompletionMessage{Role: "user", Content: "u2"}, + openai.ChatCompletionMessage{Role: "assistant", ToolCalls: []openai.ToolCall{basicToolCall("c2", "read", `{"path":"/b"}`)}}, + openai.ChatCompletionMessage{Role: "tool", ToolCallID: "c2", Content: strings.Repeat("y", 4000)}, + openai.ChatCompletionMessage{Role: "assistant", Content: "a2"}, + ) + + if _, err := s.SendMessage("what changed?"); err == nil { + t.Fatal("expected the turn to fail") + } + if s.overflowRetries() != 1 { + t.Fatalf("overflow retries = %d, want 1", s.overflowRetries()) + } + isSummary := func(m openai.ChatCompletionMessage) bool { + return strings.HasPrefix(m.Content, "[Earlier conversation compacted.") + } + if !fragmentHas(s, isSummary) { + t.Fatalf("the failed turn was rolled back past its compaction (%d messages)", len(s.fragment.Messages)) + } + if fragmentHas(s, func(m openai.ChatCompletionMessage) bool { return m.Role == "user" && m.Content == "what changed?" }) { + t.Fatal("the failed turn's own message is still in the fragment") + } + if err := validateToolPairing(s.fragment.Messages); err != nil { + t.Fatalf("pairing broken after the rollback: %v", err) + } + + llm.ok.Store(true) + if _, err := s.SendMessage("next"); err != nil { + t.Fatalf("the next turn failed: %v", err) + } + llm.mu.Lock() + next := llm.reqs[len(llm.reqs)-1] + llm.mu.Unlock() + sawSummary := false + for _, m := range next { + if isSummary(m) { + sawSummary = true + } + if m.Content == strings.Repeat("x", 4000) { + t.Fatal("the next turn sent the uncompacted history") + } + } + if !sawSummary { + t.Fatalf("the next turn did not start from the compacted history (%d messages)", len(next)) + } + if err := validateToolPairing(next); err != nil { + t.Fatalf("the next turn has orphaned tool messages: %v", err) + } +} + +// scriptedErrLLM fails its turn requests with errs, in order, then answers. +// Summary requests are counted apart and always succeed. +type scriptedErrLLM struct { + mu sync.Mutex + errs []error + maxTokens []int + summaries int +} + +func (l *scriptedErrLLM) CreateChatCompletion(ctx context.Context, req openai.ChatCompletionRequest) (cogito.LLMReply, cogito.LLMUsage, error) { + l.mu.Lock() + defer l.mu.Unlock() + if n := len(req.Messages); n > 0 && strings.HasPrefix(req.Messages[n-1].Content, compactInstruction) { + l.summaries++ + return replyWith("summary"), cogito.LLMUsage{}, nil + } + l.maxTokens = append(l.maxTokens, req.MaxTokens) + if len(l.errs) > 0 { + err := l.errs[0] + l.errs = l.errs[1:] + return cogito.LLMReply{}, cogito.LLMUsage{}, err + } + return replyWith("ok"), cogito.LLMUsage{PromptTokens: 10, TotalTokens: 10}, nil +} + +func (l *scriptedErrLLM) Ask(ctx context.Context, f cogito.Fragment) (cogito.Fragment, error) { + l.mu.Lock() + l.summaries++ + l.mu.Unlock() + return f.AddMessage("assistant", "summary"), nil +} + +func overflowFixtureErr(t *testing.T, name string) error { + t.Helper() + b, err := os.ReadFile("testdata/overflow/" + name + ".json") + if err != nil { + t.Fatal(err) + } + var fx struct { + Raw string `json:"raw"` + } + if err := json.Unmarshal(b, &fx); err != nil { + t.Fatal(err) + } + return errors.New(fx.Raw) +} + +// The regolo output-cap error lowers the cap and retries, without compacting. +func TestOverflowOutputCapLowersTheCap(t *testing.T) { + llm := &scriptedErrLLM{errs: []error{overflowFixtureErr(t, "regolo-vllm-output-cap-nonstream")}} + s := newOverflowSession(t, llm) + s.outputCap = 250000 + s.compaction.MaxContextTokens = 400000 + + if _, err := s.SendMessage("what changed?"); err != nil { + t.Fatalf("SendMessage: %v", err) + } + if cap, _ := s.requestLimits(); cap != 210000 { + t.Fatalf("output cap = %d, want the stated 210000", cap) + } + if llm.summaries != 0 || s.overflowRetries() != 0 { + t.Fatalf("an output-cap error compacted (summaries %d, retries %d)", llm.summaries, s.overflowRetries()) + } + if len(llm.maxTokens) != 2 || llm.maxTokens[1] > 210000 { + t.Fatalf("max_tokens per request = %v, want a retry at most 210000", llm.maxTokens) + } + + // Never raised: an output-cap error stating more than the cap changes nothing. + s.lowerOutputCap(300000, s.llmModel) + if cap, _ := s.requestLimits(); cap != 210000 { + t.Fatalf("output cap raised to %d", cap) + } +} + +// The regolo budget error retries with a lower max_tokens, without compacting. +func TestOverflowBudgetRetriesWithLowerMaxTokens(t *testing.T) { + // The regolo capture, verbatim but for its input figure: vLLM counted 133 + // input tokens there, less than this session's tool schemas alone, so + // the lowered reservation would not bind. 8000 keeps the shape (input + // well under the window, prompt + reservation over it). + raw := overflowFixtureErr(t, "regolo-vllm-budget-nonstream").Error() + raw = strings.ReplaceAll(raw, "a total of 210123 tokens: 133 tokens", "a total of 217990 tokens: 8000 tokens") + llm := &scriptedErrLLM{errs: []error{errors.New(raw)}} + s := newOverflowSession(t, llm) + s.outputCap = 209990 + s.compaction.MaxContextTokens = 210000 + before := append([]openai.ChatCompletionMessage(nil), s.fragment.Messages...) + + if _, err := s.SendMessage("what changed?"); err != nil { + t.Fatalf("SendMessage: %v", err) + } + if llm.summaries != 0 || s.overflowRetries() != 0 { + t.Fatalf("a budget overflow compacted (summaries %d, retries %d)", llm.summaries, s.overflowRetries()) + } + for i, m := range before { + if s.fragment.Messages[i].Content != m.Content { + t.Fatalf("message %d changed: %+v", i, s.fragment.Messages[i]) + } + } + // window − input − outputSafetyMargin, from the error's own figures. + want := 210000 - 8000 - outputSafetyMargin + if len(llm.maxTokens) != 2 || llm.maxTokens[1] != want || llm.maxTokens[1] >= llm.maxTokens[0] { + t.Fatalf("max_tokens per request = %v, want a retry at %d", llm.maxTokens, want) + } + // One turn only. + if cap, _ := s.requestLimits(); cap != 209990 { + t.Fatalf("output cap after the turn = %d, want the session's 209990", cap) + } +} + +// Regression for the failure diagnosed on regolo (LiteLLM in front of vLLM). +// +// Scaled 1:10 to keep it fast: vLLM's window is 21000 (really 210000), and +// LiteLLM's /model/info reports 9600 input / 9600 output (really 96000). The +// shape is what matters: the LiteLLM window is smaller than vLLM's, and the +// output cap equals the LiteLLM window. The fake vLLM counts a prompt as +// bytes/3, not nib's bytes/4, the way a real tokenizer disagrees with nib's +// estimate, and rejects prompt + max_tokens > 21000 with vLLM's current +// message wrapped the way regolo wraps it. +// +// The session grows past 11400 tokens (window − cap, really 114000). Every +// turn must succeed or recover, and a manual /compact must succeed after an +// overflow. +func TestOverflowRegoloRegressionOverHTTP(t *testing.T) { + xlog.SetLogger(xlog.NewLogger(xlog.LogLevel("error"), "")) + + const ( + vllmWindow = 21000 + litellmCap = 9600 + ) + var rejected, maxInput atomic.Int64 + + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Method != http.MethodPost { + if strings.HasSuffix(r.URL.Path, "/model/info") { + w.Header().Set("Content-Type", "application/json") + _, _ = fmt.Fprintf(w, `{"data":[{"model_name":"m","model_info":{"max_input_tokens":%d,"max_output_tokens":%d}}]}`, litellmCap, litellmCap) + return + } + http.NotFound(w, r) + return + } + body, _ := io.ReadAll(r.Body) + var req struct { + Messages []struct { + Content any `json:"content"` + } `json:"messages"` + Tools json.RawMessage `json:"tools"` + MaxTokens int `json:"max_tokens"` + MaxCompletionTokens int `json:"max_completion_tokens"` + } + _ = json.Unmarshal(body, &req) + bytes, last := 0, "" + for _, m := range req.Messages { + if c, ok := m.Content.(string); ok { + bytes += len(c) + last = c + } + } + input := bytes/3 + len(req.Tools)/4 + output := max(req.MaxTokens, req.MaxCompletionTokens) + if input+output > vllmWindow { + rejected.Add(1) + inner, _ := json.Marshal(map[string]any{ + "object": "error", + "message": fmt.Sprintf("Requested token count exceeds the model's maximum context length of %d tokens. You requested a total of %d tokens: %d tokens from the input messages and %d tokens for the completion. Please reduce the number of tokens in the input messages or the completion to fit within the limit.", + vllmWindow, input+output, input, output), + "type": "BadRequestError", "param": nil, "code": 400, + }) + wrapped := "regolo.BadRequestError: Hosted_vllmException - " + string(inner) + + "No fallback model group found for original model_group=m.. Received Model Group=m\nAvailable Model Group Fallbacks=None\n" + + "regolo.BadRequestError: Hosted_vllmException - " + string(inner) + + "No fallback model group found for original model_group=m. Retried: 2 times" + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusBadRequest) + _ = json.NewEncoder(w).Encode(map[string]any{ + "error": map[string]any{"message": wrapped, "type": nil, "param": nil, "code": "400"}, + }) + return + } + for { + old := maxInput.Load() + if int64(input) <= old || maxInput.CompareAndSwap(old, int64(input)) { + break + } + } + reply := "ok" + if strings.HasPrefix(last, compactInstruction) { + reply = "summary of the earlier turns" + } + w.Header().Set("Content-Type", "application/json") + _ = json.NewEncoder(w).Encode(map[string]any{ + "id": "cmpl-1", "object": "chat.completion", "created": 1, "model": "m", + "choices": []map[string]any{{ + "index": 0, + "message": map[string]any{"role": "assistant", "content": reply}, + "finish_reason": "stop", + }}, + "usage": map[string]any{"prompt_tokens": input, "completion_tokens": 1, "total_tokens": input + 1}, + }) + })) + defer srv.Close() + + cfg := types.Config{ + Model: "m", + APIKey: "k", + BaseURL: srv.URL + "/v1", + LogLevel: "error", + ApprovalMode: "auto", + BuiltinTools: []string{"read"}, + // MaxRetries 1: see TestOverflowRecoveryOverHTTP. + AgentOptions: types.AgentOptions{Iterations: 10, MaxAttempts: 3, MaxRetries: 1}, + // MaxContextTokens 0: the window is discovered from /model/info. + Compaction: types.CompactionConfig{Threshold: 0.8, KeepRecent: 2, ReserveTokens: 1024}, + } + s, err := NewSession(context.Background(), cfg, Callbacks{}) + if err != nil { + t.Fatalf("NewSession: %v", err) + } + defer s.Close() + + send := func(i, size int) error { + _, err := s.SendMessage(fmt.Sprintf("turn %d: ", i) + strings.Repeat("w", size)) + return err + } + // Ordinary turns, budgeted against LiteLLM's window. + turn := 0 + for ; turn < 3; turn++ { + if err := send(turn, 6000); err != nil { + t.Fatalf("turn %d got stuck: %v", turn, err) + } + } + // A paste larger than vLLM's real window. No recovery can make it fit, + // and with the reservation clamped this is the only way a session + // budgeted against LiteLLM's window meets vLLM's. The turn fails with + // the overflow, teaches nib the real window, and must not poison the + // session. + if err := send(turn, 3*(vllmWindow+500)); err == nil || !isContextOverflow(err) { + t.Fatalf("the oversized paste: err = %v, want the overflow", err) + } + turn++ + if got := s.contextWindow(); got != vllmWindow { + t.Fatalf("contextWindow = %d, want the learned %d", got, vllmWindow) + } + // Now the session grows past window − cap against the learned window. + // Every turn must succeed. + for end := turn + 12; turn < end; turn++ { + if err := send(turn, 6000); err != nil { + t.Fatalf("turn %d got stuck: %v", turn, err) + } + } + if rejected.Load() < 2 { + t.Fatalf("the backend rejected %d requests: the regression was not exercised", rejected.Load()) + } + if maxInput.Load() <= vllmWindow-litellmCap { + t.Fatalf("largest accepted prompt = %d tokens; the session never grew past %d", maxInput.Load(), vllmWindow-litellmCap) + } + if _, _, err := s.CompactHistory(); err != nil { + t.Fatalf("manual /compact failed after an overflow: %v", err) + } +} diff --git a/chat/session.go b/chat/session.go index a3eae77..56e7547 100644 --- a/chat/session.go +++ b/chat/session.go @@ -241,6 +241,10 @@ type Session struct { // LLM wrapper clamps each request's reservation against it (see // clampOutputTokens). outputCap int + // turnOutputCap, when above 0, lowers outputCap for the current turn + // only; guarded by modelMu. A budget overflow sets it from the figures + // the backend stated (see budgetRetryOutput), and each turn clears it. + turnOutputCap int // prunedMu guards the tool-output pruning state below. The manipulator reads // it from inside cogito's loop, and nothing here should assume which @@ -1339,6 +1343,86 @@ func (s *Session) mcpToolFilter() func(*mcp.ClientSession, string) bool { } } +// fragmentTokens is the byte/4 estimate of the conversation, read under the +// history lock. +func (s *Session) fragmentTokens() int { + s.historyMu.Lock() + defer s.historyMu.Unlock() + return estimateTokens(s.fragment.Messages) +} + +// noticesText is the message pendingNoticesMessage made for notices, or "" +// when there were none. +func noticesText(notices []string) string { + if len(notices) == 0 { + return "" + } + return pendingNoticesMessage(notices) +} + +// dropFailedTurn removes a failed turn's own messages from msgs, a history +// that overflow recovery compacted during the turn, and keeps the compaction. +// +// The turn's messages start at its user message (and the notices message +// just before it, which the caller hands back to the pending list). When the +// user message is still in msgs, everything from it on goes. When it is not, +// the compaction's summary covers it, and every message after the summary +// (the first non-system message) is the turn's own: the kept tail is a suffix +// of the history, and the user message came before all of it. System messages +// are kept wherever they are, since ensureSystemPrompt may have appended the +// prompt after the tail. +// +// It reports false when the result would break the tool pairing; the caller +// then falls back to the pre-turn history. +func dropFailedTurn(msgs []openai.ChatCompletionMessage, user openai.ChatCompletionMessage, notices string) ([]openai.ChatCompletionMessage, bool) { + cut := -1 + for i := len(msgs) - 1; i >= 0; i-- { + m := msgs[i] + if m.Role == user.Role && m.Content == user.Content && len(m.MultiContent) == len(user.MultiContent) { + cut = i + break + } + } + if cut >= 0 { + if notices != "" && cut > 0 && msgs[cut-1].Role == "user" && msgs[cut-1].Content == notices { + cut-- + } + } else { + cut = len(msgs) + for i, m := range msgs { + if m.Role != "system" { + cut = i + 1 + break + } + } + } + kept := append([]openai.ChatCompletionMessage(nil), msgs[:cut]...) + for _, m := range msgs[cut:] { + if m.Role == "system" { + kept = append(kept, m) + } + } + if len(kept) == 0 || validateToolPairing(kept) != nil { + return nil, false + } + return kept, true +} + +// dropFailedDisplay removes a failed turn's user message, and anything after +// it, from the display copy, as the pre-turn rollback does. +func dropFailedDisplay(msgs []openai.ChatCompletionMessage, text string) []openai.ChatCompletionMessage { + for i := len(msgs) - 1; i >= 0; i-- { + if msgs[i].Role == "user" && msgs[i].Content == text { + return msgs[:i] + } + } + return msgs +} + +// overflowTrim is the first step of overflow recovery. It is a variable so a +// test can make it fail and reach the basicCompact fallback. +var overflowTrim = func(s *Session, ctx context.Context) error { return s.iterativeTrim(ctx) } + // buildUserFragment appends the user turn to the fragment, attaching multimodal // parts. cogito's AddMessage routes image parts into image_url MultiContent and // audio/video into the fragment's transient PendingNativeParts (send-once). @@ -1586,11 +1670,17 @@ func (s *Session) SendMessage(text string, parts ...ContentPart) (string, error) s.overflowMu.Lock() s.overflowRetried = 0 s.overflowMu.Unlock() + // A budget overflow lowers the output reservation for one turn only. + s.setTurnOutputCap(0) + defer s.setTurnOutputCap(0) s.turnRetryMu.Lock() s.turnRetryTotal = 0 s.turnRetryMu.Unlock() // Failed attempts in a row that made no progress; see turnRetryBudget. stalled := 0 + // Output-cap and budget overflows are retried at most once per turn each; + // neither compacts, so neither counts as an overflow recovery. + capRetried, budgetRetried := false, false defer s.endTurn() // Report this turn's own size while it runs; hand authority back to // s.fragment (which compaction may since have shrunk) once it ends. @@ -1644,6 +1734,9 @@ func (s *Session) SendMessage(text string, parts ...ContentPart) (string, error) s.fragment = s.fragment.AddMessage("user", pendingNoticesMessage(notices)) } s.fragment = buildUserFragment(s.fragment, text, parts) + // The turn's own message, so a failed turn can be cut out of a history + // that overflow recovery has since compacted (see dropFailedTurn). + turnUser := s.fragment.Messages[len(s.fragment.Messages)-1] s.messages = append(s.messages, openai.ChatCompletionMessage{ Role: "user", Content: text, @@ -1906,6 +1999,47 @@ func (s *Session) SendMessage(text string, parts ...ContentPart) (string, error) // decides whether a 413 with no token wording is an overflow. turnOverflow := s.classifyTurnOverflow(err) logUnclassifiedRejection(err, turnOverflow) + announce := func(status string) { + if s.callbacks.OnStatus != nil { + s.callbacks.OnStatus(status) + } + } + if turnCtx.Err() == nil { + switch turnOverflow.Kind { + case KindOutputCap: + // The requested output alone is above the model's + // maximum. Compaction cannot fix that; a lower cap can. + // Once per turn, and only when the cap really drops, so + // the retry is not the same request again. + if !capRetried && s.lowerOutputCap(turnOverflow.Window, mainModel) { + capRetried = true + xlog.Warn("output cap above the model's maximum; lowered it and retrying", "cap", turnOverflow.Window) + announce("Output limit too large — lowering it and retrying…") + announce(retryResumeStatus) + continue + } + case KindBudget: + // The prompt fits, but prompt + reserved output does not. + // Never compact for this: lower the reservation for the + // rest of the turn, from the backend's exact figures. + if w, ok := learnedWindowFrom(err); ok { + s.rememberWindow(w, mainModel) + } + if out, ok := budgetRetryOutput(turnOverflow); ok && !budgetRetried { + budgetRetried = true + s.setTurnOutputCap(out) + xlog.Warn("prompt plus output reservation above the window; retrying with a smaller reservation", "max_tokens", out) + announce("Context window nearly full — reserving less output and retrying…") + announce(retryResumeStatus) + continue + } + // The prompt leaves less than minOutputTokens of room, or + // a smaller reservation did not fit either. Only a + // smaller prompt helps now, so this is a context + // overflow. + turnOverflow.Kind = KindContext + } + } if turnCtx.Err() == nil && turnOverflow.Kind == KindContext { if w, ok := learnedWindowFrom(err); ok { s.rememberWindow(w, mainModel) @@ -1915,41 +2049,45 @@ func (s *Session) SendMessage(text string, parts ...ContentPart) (string, error) s.overflowMu.Unlock() if first { - cb, ca, cerr := s.compactHistory(turnCtx) - switch { - case cerr != nil: - // Report the ORIGINAL overflow, not the summariser's + // The escalation chain: iterativeTrim (LLM compaction, + // then a smaller tail, a forced prune and a hard + // truncation), then basicCompact when none of that + // fits. The whole chain is ONE recovery attempt. + cb := s.fragmentTokens() + status := "" + if terr := overflowTrim(s, turnCtx); terr == nil { + status = "Context window exceeded — compacting and retrying…" + } else if turnCtx.Err() == nil { + // Report the ORIGINAL overflow, not the trim's // failure: the first is the one the user can act on. - xlog.Warn("overflow recovery: compaction failed", "error", cerr) - case cb == ca: - // Nothing to summarise, so a retry would send a - // byte-identical request. - default: + xlog.Warn("overflow recovery: trimming failed, trying basic compaction", "error", terr) + if berr := s.basicCompact(turnCtx); berr == nil { + status = "Context window exceeded — using basic fallback compaction and retrying…" + } else { + xlog.Warn("overflow recovery: basic compaction failed", "error", berr) + } + } + if status != "" { + ca := s.fragmentTokens() s.overflowMu.Lock() s.overflowRetried++ s.overflowMu.Unlock() - // compactHistory rebuilt the fragment as [summary] + - // tail, and renderMessages skips system content, so the - // system prompt was dropped without even being - // represented in the summary. SendMessage's own guard - // sits ABOVE this loop, so the `continue` below would - // re-send the turn with no identity, no working + // Compaction rebuilt the fragment as [summary] + tail, + // and renderMessages skips system content, so the + // system prompt may have been dropped without even + // being represented in the summary. SendMessage's own + // guard sits ABOVE this loop, so the `continue` below + // would re-send the turn with no identity, no working // directory, no skills index and none of the tool // guidance. Re-apply the same guard here: it is a // no-op whenever the prompt survived in the kept tail. s.ensureSystemPrompt() - // Announced here, AFTER compaction, and only on the + // Announced here, AFTER the chain, and only on the // branch that reaches the `continue` below. The status - // promises a retry, and the two branches above are the - // cases where no retry happens: the user would be told - // nib was retrying and then handed the bare overflow - // error, with no corrective status to withdraw the - // promise. Compaction is a whole LLM call, so this does - // cost the user a few silent seconds before the line - // appears — the alternative is a line that lies. - if s.callbacks.OnStatus != nil { - s.callbacks.OnStatus("Context window exceeded — compacting and retrying…") - } + // promises a retry; a chain that changed nothing makes + // none, and the user would be told nib was retrying + // and then handed the bare overflow error. + announce(status) if s.callbacks.OnCompactDone != nil { s.callbacks.OnCompactDone(cb, ca) } @@ -1957,9 +2095,7 @@ func (s *Session) SendMessage(text string, parts ...ContentPart) (string, error) // the status back, or a retry that streams a plain // answer leaves "compacting and retrying" on screen // for the rest of the turn (see retryResumeStatus). - if s.callbacks.OnStatus != nil { - s.callbacks.OnStatus(retryResumeStatus) - } + announce(retryResumeStatus) continue } } @@ -2015,9 +2151,23 @@ func (s *Session) SendMessage(text string, parts ...ContentPart) (string, error) // every later turn overflow too. Token usage is kept — the // backend billed those tokens regardless. if overflow { + recovered := s.overflowRetries() > 0 s.historyMu.Lock() - s.fragment = preTurnFragment - s.messages = preTurnMessages + // When recovery compacted, rolling back to preTurnFragment + // would undo that compaction: the next turn would start from + // the same history that overflowed, and overflow again. Keep + // the compaction and cut out only this turn's own messages. + kept, ok := []openai.ChatCompletionMessage(nil), false + if recovered { + kept, ok = dropFailedTurn(s.fragment.Messages, turnUser, noticesText(notices)) + } + if ok { + s.fragment.Messages = kept + s.messages = dropFailedDisplay(s.messages, text) + } else { + s.fragment = preTurnFragment + s.messages = preTurnMessages + } s.historyMu.Unlock() // The rollback dropped the notices too; keep them for the next turn. s.restorePendingNotices(notices) diff --git a/selfdoc/README.md b/selfdoc/README.md index 1c5f36d..53e655e 100644 --- a/selfdoc/README.md +++ b/selfdoc/README.md @@ -580,7 +580,12 @@ log_level: error # compact when the prompt itself does not fit (kind "context"), lower the # requested output when only prompt + output does not fit ("budget"), or # lower the output cap when the output alone is above the model's maximum -# ("output_cap"). It knows the wording of llama.cpp/LocalAI, vLLM, OpenAI, +# ("output_cap"), and then retry the turn once. When the summary alone does +# not make the history fit, nib stubs old tool outputs, then drops older turns +# (saved as an artifact:// the model can read), and as a last resort compacts +# without the model. If the retry still does not fit, the turn fails, but the +# compaction is kept, so the next turn does not overflow again. It knows the +# wording of llama.cpp/LocalAI, vLLM, OpenAI, # Anthropic and Gemini. overflow_patterns teaches it another backend without a # rebuild: the rows are tried before the built-in ones, first match wins. A # regex states its figures with the named groups window, total, input and From be387c92c1dbe6e2f9c1385eaeef5697553d1fc1 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Fri, 25 Sep 2026 22:01:28 +0000 Subject: [PATCH 11/16] chore(deps): bump cogito to return typed status errors from streams cogito 01190bf (mudler/cogito#79) returns a non-200 streaming answer as an *openai.APIError or *openai.RequestError that carries the HTTP status, instead of the text "localai stream: status 400: ". errorStatus now reads the status from the typed error on the streaming path too; the text form stays supported. Co-Authored-By: Claude Opus 5.5 --- go.mod | 2 +- go.sum | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/go.mod b/go.mod index f51c397..c10ca8c 100644 --- a/go.mod +++ b/go.mod @@ -37,7 +37,7 @@ require ( github.com/msuozzo/bonsai/bonsai-tsx v0.4.0 github.com/msuozzo/bonsai/bonsai-typescript v0.4.0 github.com/msuozzo/bonsai/bonsai-yaml v0.4.0 - github.com/mudler/cogito v0.11.1-0.20260923083907-5cbfdb15460b + github.com/mudler/cogito v0.11.1-0.20260925214156-01190bfd0565 github.com/mudler/xlog v0.0.1 github.com/muesli/termenv v0.16.0 github.com/sashabaranov/go-openai v1.41.2 diff --git a/go.sum b/go.sum index aa4b123..129f836 100644 --- a/go.sum +++ b/go.sum @@ -207,8 +207,8 @@ github.com/msuozzo/bonsai/bonsai-typescript v0.4.0 h1:BD6JHfCs3rGNUttL+tMnz/iav1 github.com/msuozzo/bonsai/bonsai-typescript v0.4.0/go.mod h1:DElq3FKqbTkJXcat1o8m98ijL3wUf+jW+DNInzh5Mgo= github.com/msuozzo/bonsai/bonsai-yaml v0.4.0 h1:PJfyfjMQrcWU5e3ib0xGsVuOafNuimf5fK0NhTk8olU= github.com/msuozzo/bonsai/bonsai-yaml v0.4.0/go.mod h1:z8jc0tjXDSOQRKRG0G3uymvVxHdBbrVsWd149nPOyHE= -github.com/mudler/cogito v0.11.1-0.20260923083907-5cbfdb15460b h1:DAyr5sxdltoLcLqG32zHYyfR5YOZow3Vufxoz3jH5SA= -github.com/mudler/cogito v0.11.1-0.20260923083907-5cbfdb15460b/go.mod h1:6sfja3lcu2nWRzEc0wwqGNu/eCG3EWgij+8s7xyUeQ4= +github.com/mudler/cogito v0.11.1-0.20260925214156-01190bfd0565 h1:A6CT8LINsF5thZlgMCC1iicAfOTdlWjSyURJfw3JpSQ= +github.com/mudler/cogito v0.11.1-0.20260925214156-01190bfd0565/go.mod h1:6sfja3lcu2nWRzEc0wwqGNu/eCG3EWgij+8s7xyUeQ4= github.com/mudler/xlog v0.0.1 h1:yR3/wszd3ZM6u1n96YITJZ4yUcDgqHSwvQmzUJa+8vg= github.com/mudler/xlog v0.0.1/go.mod h1:39f5vcd05Qd6GWKM8IjyHNQ7AmOx3ZM0YfhfIGhC18U= github.com/muesli/ansi v0.0.0-20230316100256-276c6243b2f6 h1:ZK8zHtRHOkbHy6Mmr5D264iyp3TiX5OmNcI5cIARiQI= From 43b51c35cf6ae18c700d8b7c317fcca3ca54b672 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Fri, 25 Sep 2026 22:13:52 +0000 Subject: [PATCH 12/16] fix(chat): roll a summary head known to be too large from the start With a known window, a head larger than the summary prompt budget still went out whole first, and fitSummaryInput cut long pieces and dropped middle messages to make it fit. That built a lossy summary input on purpose, and cost the same one call before rolling would start. Now such a head goes directly to rollingCover with the initial target (window less the summary output reservation, capped by ContextBudget). A head that fits, or an unknown window, keeps the single whole-head call and rolls only after an overflow. fitSummaryInput still cuts a single piece too large for a chunk on its own. Co-Authored-By: Claude Opus 5.5 --- chat/compact.go | 38 +++++++++-- chat/compact_fit_test.go | 9 ++- chat/overflow_trim_test.go | 25 +++++-- chat/summary_rolling_test.go | 126 +++++++++++++++++++++++++++++++++++ 4 files changed, 185 insertions(+), 13 deletions(-) diff --git a/chat/compact.go b/chat/compact.go index 27c63f5..cbeabc4 100644 --- a/chat/compact.go +++ b/chat/compact.go @@ -696,10 +696,13 @@ func (s *Session) summaryPromptLimit(output int) int { // (the pruned view of them). An empty head, with a nil error, means there was // nothing to compact. // -// The whole head goes out in one request first. When the backend rejects it -// as too large, the head is summarized in chunks instead (see rollingCover), -// each fitted to the retry target summaryRetryTarget derives from the -// rejection. When the chunks cover the whole head, the result is the summary +// When the window is known and the rendered head is larger than the summary +// prompt budget (summaryPromptLimit), the head is summarized in chunks fitted +// to that budget from the start (see rollingCover): fitting it into one +// request would cut and drop messages on purpose. Otherwise the whole head +// goes out in one request first, and when the backend rejects it as too +// large, the head is summarized in chunks instead, each fitted to the retry +// target summaryRetryTarget derives from the rejection. When the chunks cover the whole head, the result is the summary // and the original tail. When they stop at maxRollingChunks, the head ends // where they stopped and the rest joins the tail, verbatim: nothing is lost, // and the next compaction can summarize it. The loop stops on a non-overflow @@ -712,7 +715,14 @@ func (s *Session) summarizeFitting(ctx context.Context, msgs []openai.ChatComple return "", nil, nil, nil } limit := s.summaryPromptLimit(s.summaryOutputTokens()) - summary, sent, serr := s.summarize(ctx, renderMessages(view(head)), limit) + pieces := renderMessages(view(head)) + if limit > 0 && piecesTokens(pieces) > limit { + // The head is known not to fit: fitSummaryInput would cut it down + // on purpose, so it goes out in chunks from the start instead. The + // target is the whole prompt, instruction included. + return s.rollingHead(ctx, msgs, head, view, limit+tokensOf(summaryPrefix), 0) + } + summary, sent, serr := s.summarize(ctx, pieces, limit) if serr == nil { return summary, head, tail, nil } @@ -723,7 +733,14 @@ func (s *Session) summarizeFitting(ctx context.Context, msgs []openai.ChatComple if target <= 0 { return "", nil, nil, serr } - summary, covered, err := s.rollingCover(ctx, head, view, target, 1) + return s.rollingHead(ctx, msgs, head, view, target, 1) +} + +// rollingHead summarizes head in chunks fitted to target (see rollingCover) +// and splits msgs where the chunks stopped. failures is passed through to +// rollingCover. +func (s *Session) rollingHead(ctx context.Context, msgs, head []openai.ChatCompletionMessage, view func([]openai.ChatCompletionMessage) []openai.ChatCompletionMessage, target, failures int) (summary string, h, tail []openai.ChatCompletionMessage, err error) { + summary, covered, err := s.rollingCover(ctx, head, view, target, failures) if err != nil { return "", nil, nil, err } @@ -732,6 +749,15 @@ func (s *Session) summarizeFitting(ctx context.Context, msgs []openai.ChatComple return summary, msgs[:covered], msgs[covered:], nil } +// piecesTokens is the byte/4 size of pieces joined. +func piecesTokens(pieces []summaryPiece) int { + n := 0 + for _, p := range pieces { + n += len(p.text) + } + return n / 4 +} + // maxRollingChunks bounds how many chunks rollingCover summarizes for one // compaction. What is left after that stays verbatim. const maxRollingChunks = 8 diff --git a/chat/compact_fit_test.go b/chat/compact_fit_test.go index 212a025..dc29d56 100644 --- a/chat/compact_fit_test.go +++ b/chat/compact_fit_test.go @@ -93,13 +93,16 @@ func TestCompactHistoryFitsTheSummaryPromptInTheBudget(t *testing.T) { if _, _, err := s.CompactHistory(); err != nil { t.Fatalf("CompactHistory: %v", err) } - if got := tokensOf(llm.prompts[0]); got > budget { - t.Fatalf("summary prompt is ~%d tokens, budget is %d", got, budget) + // The head is known not to fit, so it goes out in chunks from the start. + for i, p := range llm.prompts { + if got := tokensOf(p); got > budget { + t.Fatalf("summary prompt %d is ~%d tokens, budget is %d", i, got, budget) + } } if !strings.Contains(llm.prompts[0], "goal: fix the parser") { t.Fatal("fitting the prompt dropped the user's goal") } - if !strings.Contains(llm.prompts[0], "f5.go") { + if !strings.Contains(strings.Join(llm.prompts, ""), "f5.go") { t.Fatal("fitting the prompt dropped the record of which tools ran") } } diff --git a/chat/overflow_trim_test.go b/chat/overflow_trim_test.go index 7549f4b..be95169 100644 --- a/chat/overflow_trim_test.go +++ b/chat/overflow_trim_test.go @@ -42,8 +42,11 @@ func TestIterativeTrimStep1CompactionFits(t *testing.T) { if err := s.iterativeTrim(context.Background()); err != nil { t.Fatalf("iterativeTrim: %v", err) } - if llm.calls != 1 { - t.Fatalf("want 1 summary call, got %d", llm.calls) + // The head (u1..a2) is larger than the summary prompt budget, so step 1 + // summarizes it in exactly 2 rolling chunks. The [summary, u3, a3] + // shape below shows it was step 1 (KeepRecent 2), not the shrunk keep. + if llm.calls != 2 { + t.Fatalf("want 2 summary chunk calls from step 1, got %d", llm.calls) } got := trimFragment(s) if len(got) != 3 || !strings.Contains(got[0].Content, "SUMMARY-TEXT") { @@ -67,11 +70,25 @@ func TestIterativeTrimNoOpFallsThroughToShrinkKeep(t *testing.T) { // KeepRecent 10 > len(frag): compactHistory has nothing to compact. s := newCompactTestSession(llm, 10, frag, frag) + // Step 1 on its own: nothing to compact, so no summary call and no change. + if _, _, err := s.compactHistory(context.Background()); err != nil { + t.Fatalf("compactHistory: %v", err) + } + if llm.calls != 0 { + t.Fatalf("want 0 summary calls from step 1, got %d", llm.calls) + } + if got := trimFragment(s); len(got) != len(frag) { + t.Fatalf("step 1 changed the fragment: %d messages", len(got)) + } + if err := s.iterativeTrim(context.Background()); err != nil { t.Fatalf("iterativeTrim: %v", err) } - if llm.calls != 1 { - t.Fatalf("want exactly 1 summary call (from the shrunk keep), got %d", llm.calls) + // Step 1 makes no call, so every call comes from the shrunk keep: its + // head is larger than the summary prompt budget and goes out in exactly + // 2 rolling chunks. + if llm.calls != 2 { + t.Fatalf("want exactly 2 summary chunk calls (from the shrunk keep), got %d", llm.calls) } got := trimFragment(s) if len(got) >= len(frag) || !strings.Contains(got[0].Content, "SUMMARY-TEXT") { diff --git a/chat/summary_rolling_test.go b/chat/summary_rolling_test.go index 20afb69..2ad2740 100644 --- a/chat/summary_rolling_test.go +++ b/chat/summary_rolling_test.go @@ -280,3 +280,129 @@ func TestRollingSummaryCutsAMessageLargerThanTheTarget(t *testing.T) { t.Fatal("the artifact must hold the over-large message in full") } } + +// knownWindowSession is a session that knows the backend's 8192-token window, +// where every summary reserves 4000 of it. +func knownWindowSession(llm cogito.LLM, frag []openai.ChatCompletionMessage) *Session { + s := newCompactTestSession(llm, 2, frag, frag) + s.compaction.MaxContextTokens = 8192 + s.compaction.SummaryMaxTokens = 4000 + s.artifacts = mcp.NewArtifactStore() + return s +} + +const ( + cutMarker = "bytes omitted to fit the summary" + droppedMarker = "earlier messages omitted to fit the summary" +) + +func TestRollingStartsAtOnceForAHeadKnownToBeTooLarge(t *testing.T) { + frag := shortHistory(24) // ~13k tokens against a ~3.7k target + orig := cloneMessages(frag) + llm := &rollingLLM{limit: 8192} + s := knownWindowSession(llm, frag) + + if _, _, err := s.CompactHistory(); err != nil { + t.Fatalf("CompactHistory: %v", err) + } + if len(llm.reqs) < 3 || llm.served != len(llm.reqs) { + t.Fatalf("requests = %d, served = %d: want several chunks and no overflow", len(llm.reqs), llm.served) + } + for i := range llm.reqs { + p := llm.prompt(i) + if strings.Contains(p, "BODY00") && strings.Contains(p, "BODY23") { + t.Fatalf("request %d carries the whole head", i) + } + if strings.Contains(p, cutMarker) || strings.Contains(p, droppedMarker) { + t.Fatalf("request %d was cut by fitSummaryInput", i) + } + } + all := "" + for i := range llm.reqs { + all += llm.prompt(i) + } + for i := range 24 { + if c := strings.Count(all, fmt.Sprintf("BODY%02d", i)); c != 1 { + t.Fatalf("BODY%02d appears in %d chunks, want 1", i, c) + } + } + got := s.fragment.Messages + if len(got) != 3 || !strings.Contains(got[0].Content, fmt.Sprintf("SUMMARY-%d", len(llm.reqs))) { + t.Fatalf("want [final summary] + the original tail, got %d messages", len(got)) + } + if !reflect.DeepEqual(got[1:], orig[len(orig)-2:]) { + t.Fatal("the tail is not the original KeepRecent tail") + } +} + +func TestRollingNotUsedForAHeadThatFits(t *testing.T) { + frag := shortHistory(2) + llm := &rollingLLM{limit: 8192} + s := knownWindowSession(llm, frag) + if _, _, err := s.CompactHistory(); err != nil { + t.Fatalf("CompactHistory: %v", err) + } + if len(llm.reqs) != 1 { + t.Fatalf("summary calls = %d, want 1", len(llm.reqs)) + } + if !strings.HasPrefix(llm.prompt(0), summaryPrefix) || !strings.Contains(llm.prompt(0), "BODY01") { + t.Fatal("want the whole head in one normal summary call") + } +} + +func TestRollingWithAnUnknownWindowStartsOnOverflow(t *testing.T) { + frag := shortHistory(24) + llm := &rollingLLM{limit: 8192} + s := knownWindowSession(llm, frag) + s.compaction.MaxContextTokens = 0 + + if _, _, err := s.CompactHistory(); err != nil { + t.Fatalf("CompactHistory: %v", err) + } + if len(llm.reqs) < 3 || llm.served != len(llm.reqs)-1 { + t.Fatalf("requests = %d, served = %d: want an overflow then chunks", len(llm.reqs), llm.served) + } + first := llm.prompt(0) + if !strings.Contains(first, "BODY00") || !strings.Contains(first, "BODY23") { + t.Fatal("the first request must carry the whole head") + } +} + +func TestRollingProactiveCutsOnlyThePieceLargerThanTheTarget(t *testing.T) { + frag := shortHistory(24) + // One result alone is far larger than the target. + frag[4].Content = "HUGE" + strings.Repeat("y", 40000) + orig := cloneMessages(frag) + llm := &rollingLLM{limit: 8192} + s := knownWindowSession(llm, frag) + + if _, _, err := s.CompactHistory(); err != nil { + t.Fatalf("CompactHistory: %v", err) + } + if llm.served != len(llm.reqs) { + t.Fatalf("requests = %d, served = %d: want no overflow", len(llm.reqs), llm.served) + } + cuts := 0 + for i := range llm.reqs { + p := llm.prompt(i) + if strings.Contains(p, droppedMarker) { + t.Fatalf("request %d dropped whole messages", i) + } + if n := strings.Count(p, cutMarker); n > 0 { + cuts += n + if !strings.Contains(p, "tool: HUGE") { + t.Fatalf("request %d cut a piece other than the huge one", i) + } + } + } + if cuts != 1 { + t.Fatalf("cuts = %d, want exactly the one huge piece cut", cuts) + } + got := s.fragment.Messages + if len(got) != 3 || !reflect.DeepEqual(got[1:], orig[len(orig)-2:]) { + t.Fatalf("want [summary] + the original tail, got %d messages", len(got)) + } + if !strings.Contains(s.artifacts.Get(1).Content, frag[4].Content) { + t.Fatal("the artifact must hold the over-large message in full") + } +} From dd5fcf103431f63c6c67fa71db5acf5e65fc949f Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Fri, 25 Sep 2026 22:21:18 +0000 Subject: [PATCH 13/16] feat(chat): measure the tool-schema floor on the requests the turn sends SchemaBudget counted only the tools nib registers itself. cogito adds every MCP server's tools and its agent tools to each request, so with many MCP servers the real floor was far larger than the estimate and no warning fired. The request wrapper now records the size of the tools and system messages of each request that carries tools, raised to the backend's own overhead when the reported prompt tokens say more. SchemaBudget uses that floor once a request was sent, and the MCP servers' tools in the estimate before. It attributes the schemas per MCP server (cached until the server set changes), and the session reports the floor once through OnStatus when it crosses the warn threshold, and again at the error threshold. Co-Authored-By: Claude Opus 5.5 --- chat/contextlive.go | 113 ++++++++++++++++++--- chat/midturn.go | 4 +- chat/schema_budget.go | 201 ++++++++++++++++++++++++++++++++++-- chat/schema_budget_test.go | 202 +++++++++++++++++++++++++++++++++++++ chat/session.go | 14 +++ 5 files changed, 510 insertions(+), 24 deletions(-) diff --git a/chat/contextlive.go b/chat/contextlive.go index cc4ba76..bf5a61d 100644 --- a/chat/contextlive.go +++ b/chat/contextlive.go @@ -31,9 +31,52 @@ import ( // active distinguishes "no turn is running" from "a turn is running and has // not been billed yet". Without it a zero would be indistinguishable from an // idle session and ContextTokens could not tell which source to trust. +// +// floor is different: it is the fixed per-request overhead (tool schemas and +// system messages) the wrapper measured on the last request that carried +// tools, raised by the backend's own count when that is larger. It is not +// reset between turns, because it describes the session's configuration, not +// the conversation. See SchemaBudget. type liveUsage struct { prompt atomic.Int64 active atomic.Bool + + floor atomic.Int64 + measured atomic.Bool +} + +// recordFloor stores the measured fixed overhead of the latest request. +func (l *liveUsage) recordFloor(n int) { + l.floor.Store(int64(n)) + l.measured.Store(true) +} + +// raiseFloor lifts the latest floor to n when n is larger: the backend's +// count of the same request also includes tokenizer skew. +func (l *liveUsage) raiseFloor(n int) { + for { + cur := l.floor.Load() + if int64(n) <= cur || l.floor.CompareAndSwap(cur, int64(n)) { + return + } + } +} + +// schemaFloor returns the latest measured floor, and whether any request +// was measured yet. +func (l *liveUsage) schemaFloor() (int, bool) { + return int(l.floor.Load()), l.measured.Load() +} + +// promptOverhead is how many more tokens the backend reported for a request +// than the byte/4 estimate of the messages it was measured against: the +// tool schemas, anything else the estimate left out, and tokenizer skew. +// Zero when the report is not larger. +func promptOverhead(reported, estimated int) int { + if reported <= estimated { + return 0 + } + return reported - estimated } // begin arms the tracker for a new turn, discarding the previous turn's @@ -88,17 +131,48 @@ type trackedLLM struct { limits func() (cap, window int) } -// clamp applies clampOutputTokens with the session's current limits. -func (t *trackedLLM) clamp(request openai.ChatCompletionRequest) openai.ChatCompletionRequest { +// prepare measures the request's fixed overhead and applies clampOutputTokens +// with the session's current limits. The tool schemas are marshalled once for +// both. It returns the byte/4 estimate of the request's non-system messages, +// which calibrate compares with the backend's report, or -1 when the request +// was not measured. +// +// Only a request that carries tools is measured: a tool-less request (a +// summary, say) says nothing about the overhead of a turn's requests. +func (t *trackedLLM) prepare(request openai.ChatCompletionRequest) (openai.ChatCompletionRequest, int) { + toolBytes := toolSchemaBytes(request.Tools) + rest := -1 + if len(request.Tools) > 0 { + sysBytes := 0 + var other []openai.ChatCompletionMessage + for _, m := range request.Messages { + if m.Role == openai.ChatMessageRoleSystem { + sysBytes += len(m.Content) + } else { + other = append(other, m) + } + } + t.live.recordFloor((toolBytes + sysBytes) / 4) + rest = estimateTokens(other) + } if t.limits == nil { - return request + return request, rest } cap, window := t.limits() - return clampOutputTokens(request, cap, window) + return clampOutputTokensSized(request, cap, window, toolBytes/4), rest +} + +// calibrate raises the measured floor to the backend's overhead for the +// request prepare measured (rest >= 0) when the backend reported more. +func (t *trackedLLM) calibrate(promptTokens, rest int) { + if rest < 0 || promptTokens <= 0 { + return + } + t.live.raiseFloor(promptOverhead(promptTokens, rest)) } func (t *trackedLLM) CreateChatCompletion(ctx context.Context, request openai.ChatCompletionRequest) (cogito.LLMReply, cogito.LLMUsage, error) { - request = t.clamp(request) + request, rest := t.prepare(request) var usage cogito.LLMUsage reply, err := retryRequest(ctx, func() (cogito.LLMReply, error) { r, u, e := t.LLM.CreateChatCompletion(ctx, request) @@ -106,6 +180,7 @@ func (t *trackedLLM) CreateChatCompletion(ctx context.Context, request openai.Ch return r, e }) t.live.record(usage.PromptTokens) + t.calibrate(usage.PromptTokens, rest) return reply, usage, err } @@ -137,7 +212,7 @@ type trackedStreamingLLM struct { // The relay goroutine is bounded by the source channel: it ends when the // provider closes it, which cogito's own consumer loop already depends on. func (t *trackedStreamingLLM) CreateChatCompletionStream(ctx context.Context, request openai.ChatCompletionRequest) (<-chan cogito.StreamEvent, error) { - request = t.clamp(request) + request, rest := t.prepare(request) src, err := retryRequest(ctx, func() (<-chan cogito.StreamEvent, error) { return t.stream.CreateChatCompletionStream(ctx, request) }) @@ -150,6 +225,7 @@ func (t *trackedStreamingLLM) CreateChatCompletionStream(ctx context.Context, re for ev := range src { if ev.Type == cogito.StreamEventDone { t.live.record(ev.Usage.PromptTokens) + t.calibrate(ev.Usage.PromptTokens, rest) } out <- ev } @@ -193,15 +269,28 @@ const ( // A request with an explicit value, or unknown limits (cap or window 0), is // returned unchanged. func clampOutputTokens(request openai.ChatCompletionRequest, cap, window int) openai.ChatCompletionRequest { + return clampOutputTokensSized(request, cap, window, toolSchemaBytes(request.Tools)/4) +} + +// toolSchemaBytes is the size of the JSON of a request's tools. +func toolSchemaBytes(tools []openai.Tool) int { + if len(tools) == 0 { + return 0 + } + b, err := json.Marshal(tools) + if err != nil { + return 0 + } + return len(b) +} + +// clampOutputTokensSized is clampOutputTokens with the tool schemas already +// measured (toolTokens), so the wrapper does not marshal them twice. +func clampOutputTokensSized(request openai.ChatCompletionRequest, cap, window, toolTokens int) openai.ChatCompletionRequest { if cap <= 0 || window <= 0 || request.MaxTokens != 0 || request.MaxCompletionTokens != 0 { return request } - prompt := estimateTokens(request.Messages) - if len(request.Tools) > 0 { - if b, err := json.Marshal(request.Tools); err == nil { - prompt += len(b) / 4 - } - } + prompt := estimateTokens(request.Messages) + toolTokens v := min(cap, window-prompt-outputSafetyMargin) if v < minOutputTokens { v = min(minOutputTokens, cap) diff --git a/chat/midturn.go b/chat/midturn.go index 80460d6..39e61b7 100644 --- a/chat/midturn.go +++ b/chat/midturn.go @@ -112,10 +112,10 @@ func (c *turnCompactor) manipulate(msgs []openai.ChatCompletionMessage) []openai // report measured both, so its excess over the last estimate is added back. func (c *turnCompactor) overhead() int { reported := c.s.live.promptTokens() - if c.sent <= 0 || reported <= c.sent { + if c.sent <= 0 { return 0 } - return reported - c.sent + return promptOverhead(reported, c.sent) } // compact summarizes the head of out and returns the request with the summary diff --git a/chat/schema_budget.go b/chat/schema_budget.go index 73ad850..3bb7f58 100644 --- a/chat/schema_budget.go +++ b/chat/schema_budget.go @@ -1,12 +1,25 @@ package chat import ( + "cmp" + "context" "encoding/json" + "fmt" + "maps" + "slices" + "strings" + "time" + "github.com/modelcontextprotocol/go-sdk/mcp" "github.com/mudler/cogito" "github.com/mudler/xlog" + openai "github.com/sashabaranov/go-openai" ) +// mcpListTimeout bounds how long mcpSchemaCosts waits for one server to list +// its tools. +const mcpListTimeout = 3 * time.Second + // schemaWarnFraction is the share of the context budget the fixed prompt // overhead (tool schemas + system prompt) may take before SchemaBudget warns. const schemaWarnFraction = 0.60 @@ -17,16 +30,29 @@ const schemaWarnFraction = 0.60 const schemaErrorFraction = 0.90 // SchemaBudget is a read-only observation of how much of the context the fixed -// per-request overhead takes. Floor is the estimated token count of the tool -// schemas plus the system prompt — the floor no compaction can go below. +// per-request overhead takes. Floor is the token count of the tool schemas plus +// the system prompt — the floor no compaction can go below. Measured says Floor +// comes from a request the session actually sent (see trackedLLM.prepare); +// before the first one it is an estimate. Servers is the MCP servers' share of +// it, largest first. type SchemaBudget struct { Floor int + Measured bool + Servers []ServerSchemaCost WarnThreshold float64 ErrorThreshold float64 Warn bool Error bool } +// ServerSchemaCost is what one MCP server's tools add to every request: the +// number of tools it advertises and the byte/4 size of their schemas. +type ServerSchemaCost struct { + Server string + Tools int + Tokens int +} + // estimateSchemaTokens estimates the tokens the tool schemas and the system // prompt take in every request, with the same byte/4 heuristic estimateTokens // uses for messages. @@ -61,24 +87,126 @@ func (s *Session) resetSchemaTools() { s.schemaToolsMu.Unlock() } -// SchemaBudget estimates the fixed prompt overhead and compares it with the +// invalidateSchemaCosts drops the cached mcpSchemaCosts result; called when +// the set of MCP servers changes. +func (s *Session) invalidateSchemaCosts() { + s.schemaCostsMu.Lock() + s.schemaCosts, s.schemaCostsValid = nil, false + s.schemaCostsMu.Unlock() +} + +// namedClients returns the session's MCP clients with a name for each: the +// configured name of a config server, "skills" for the skills server, and the +// server's own reported name for a built-in one. +func (s *Session) namedClients() ([]string, []*mcp.ClientSession) { + var names []string + var clients []*mcp.ClientSession + for i, c := range s.clients { + name := fmt.Sprintf("mcp-%d", i+1) + if r := c.InitializeResult(); r != nil && r.ServerInfo != nil && r.ServerInfo.Name != "" { + name = r.ServerInfo.Name + } + names, clients = append(names, name), append(clients, c) + } + if s.skillsClient != nil { + names, clients = append(names, "skills"), append(clients, s.skillsClient) + } + for _, name := range slices.Sorted(maps.Keys(s.cfgClients)) { + names, clients = append(names, name), append(clients, s.cfgClients[name]) + } + return names, clients +} + +// mcpSchemaCosts lists the tools of every connected MCP client (the ones +// toolOptions hands cogito.WithMCPs) and returns what each server adds to a +// request, largest first. Tools the turn's filter drops are not counted, as +// cogito does not send them. Each listing is bounded by mcpListTimeout; a +// server that fails to list is left out. The result is cached until the set of +// servers changes (ReconcileMCPServers, SetSkills). +// +// Like allClients, it reads the client set without a lock: that set changes +// only at turn start, on the turn goroutine. +func (s *Session) mcpSchemaCosts(ctx context.Context) []ServerSchemaCost { + s.schemaCostsMu.Lock() + if s.schemaCostsValid { + out := slices.Clone(s.schemaCosts) + s.schemaCostsMu.Unlock() + return out + } + s.schemaCostsMu.Unlock() + + names, clients := s.namedClients() + filter := s.mcpToolFilter() + var out []ServerSchemaCost + for i, c := range clients { + lctx, cancel := context.WithTimeout(ctx, mcpListTimeout) + res, err := c.ListTools(lctx, nil) + cancel() + if err != nil { + xlog.Debug("schema budget: cannot list MCP tools", "server", names[i], "error", err) + continue + } + cost := ServerSchemaCost{Server: names[i]} + bytes := 0 + for _, t := range res.Tools { + if !filter(c, t.Name) { + continue + } + ot := openai.Tool{Type: openai.ToolTypeFunction, Function: &openai.FunctionDefinition{ + Name: t.Name, Description: t.Description, Parameters: t.InputSchema}} + if b, err := json.Marshal(ot); err == nil { + bytes += len(b) + } + cost.Tools++ + } + if cost.Tools == 0 { + continue + } + cost.Tokens = bytes / 4 + out = append(out, cost) + } + slices.SortStableFunc(out, func(a, b ServerSchemaCost) int { return cmp.Compare(b.Tokens, a.Tokens) }) + + s.schemaCostsMu.Lock() + s.schemaCosts, s.schemaCostsValid = out, true + s.schemaCostsMu.Unlock() + return slices.Clone(out) +} + +// SchemaBudget reports the fixed prompt overhead and compares it with the // context budget and window. It does not change behavior; it only logs a // warning when the overhead is large, for a caller to show the user. // -// The tool definitions are the ones the last toolOptions call registered -// through withTool. MCP tools and the agent-spawning tools are added inside -// cogito and are not counted. +// Once the session sent a request with tools, Floor is what the request +// wrapper measured on the latest one (all the tools cogito sent, MCP and agent +// tools included, plus the system messages), raised by the backend's own count +// when that was larger. Before that it is an estimate: the tool definitions the +// last toolOptions call registered through withTool, the system prompt, and +// the MCP servers' tools. func (s *Session) SchemaBudget() SchemaBudget { - s.schemaToolsMu.Lock() - tools := append([]cogito.ToolDefinitionInterface(nil), s.schemaTools...) - s.schemaToolsMu.Unlock() + ctx := s.ctx + if ctx == nil { + ctx = context.Background() + } + servers := s.mcpSchemaCosts(ctx) - floor := estimateSchemaTokens(tools, s.systemPrompt) + floor, measured := s.live.schemaFloor() + if !measured { + s.schemaToolsMu.Lock() + tools := append([]cogito.ToolDefinitionInterface(nil), s.schemaTools...) + s.schemaToolsMu.Unlock() + floor = estimateSchemaTokens(tools, s.systemPrompt) + for _, c := range servers { + floor += c.Tokens + } + } window := s.contextWindow() budget := ContextBudget(s.compactionConfig(), window) sb := SchemaBudget{ Floor: floor, + Measured: measured, + Servers: servers, WarnThreshold: schemaWarnFraction, ErrorThreshold: schemaErrorFraction, } @@ -98,3 +226,56 @@ func (s *Session) SchemaBudget() SchemaBudget { } return sb } + +// notifySchemaBudget tells the user, through OnStatus, when the fixed prompt +// overhead first passes the warn threshold, and again only when it passes the +// error threshold after that. +func (s *Session) notifySchemaBudget() { + if s.callbacks.OnStatus == nil { + return + } + sb := s.SchemaBudget() + level := 0 + switch { + case sb.Error: + level = 2 + case sb.Warn: + level = 1 + } + s.schemaCostsMu.Lock() + if level <= s.schemaNoticeLevel { + s.schemaCostsMu.Unlock() + return + } + s.schemaNoticeLevel = level + s.schemaCostsMu.Unlock() + s.callbacks.OnStatus(schemaBudgetNotice(sb, s.contextWindow())) +} + +// schemaBudgetNotice is the one-line notice for sb, naming at most the three +// largest MCP servers. +func schemaBudgetNotice(sb SchemaBudget, window int) string { + var b strings.Builder + fmt.Fprintf(&b, "tool schemas use %s of the %s context window", fmtTokensK(sb.Floor), fmtTokensK(window)) + if len(sb.Servers) > 0 { + parts := make([]string, 0, 4) + for _, c := range sb.Servers[:min(3, len(sb.Servers))] { + parts = append(parts, fmt.Sprintf("%s %s (%d tools)", c.Server, fmtTokensK(c.Tokens), c.Tools)) + } + if len(sb.Servers) > 3 { + parts = append(parts, "…") + } + b.WriteString(" — ") + b.WriteString(strings.Join(parts, ", ")) + } + b.WriteString(": disable MCP servers or filter their tools") + return b.String() +} + +// fmtTokensK renders a token count as "34k", or as is below a thousand. +func fmtTokensK(n int) string { + if n < 1000 { + return fmt.Sprintf("%d", n) + } + return fmt.Sprintf("%dk", (n+500)/1000) +} diff --git a/chat/schema_budget_test.go b/chat/schema_budget_test.go index 47dec47..8bb2d1f 100644 --- a/chat/schema_budget_test.go +++ b/chat/schema_budget_test.go @@ -1,11 +1,15 @@ package chat import ( + "context" "encoding/json" + "fmt" "strings" "testing" + sdkmcp "github.com/modelcontextprotocol/go-sdk/mcp" "github.com/mudler/cogito" + "github.com/mudler/nib/types" openai "github.com/sashabaranov/go-openai" ) @@ -88,3 +92,201 @@ func TestSchemaBudgetCountsRecordedTools(t *testing.T) { t.Fatalf("Floor after reset = %d, want 0", got) } } + +// promptUsageLLM reports a fixed prompt-token count for every request. +type promptUsageLLM struct { + cogito.LLM + prompt int +} + +func (u *promptUsageLLM) CreateChatCompletion(_ context.Context, _ openai.ChatCompletionRequest) (cogito.LLMReply, cogito.LLMUsage, error) { + return cogito.LLMReply{}, cogito.LLMUsage{PromptTokens: u.prompt}, nil +} + +// mcpLikeTools builds n tools the session never registered, each with a +// description of size bytes. +func mcpLikeTools(n, size int) []openai.Tool { + out := make([]openai.Tool, n) + for i := range out { + out[i] = openai.Tool{Type: openai.ToolTypeFunction, Function: &openai.FunctionDefinition{ + Name: fmt.Sprintf("mcp_tool_%d", i), Description: strings.Repeat("m", size)}} + } + return out +} + +func TestSchemaBudgetMeasuredFromRequest(t *testing.T) { + s := newCompactTestSession(&fakeSummaryLLM{}, 2, nil, nil) + s.systemPrompt = "short" + before := s.SchemaBudget() + if before.Measured { + t.Fatal("Measured before any request") + } + + tools := mcpLikeTools(20, 400) + sys := strings.Repeat("s", 800) + req := openai.ChatCompletionRequest{ + Messages: []openai.ChatCompletionMessage{{Role: "system", Content: sys}, {Role: "user", Content: "hi"}}, + Tools: tools, + } + llm := trackUsage(&promptUsageLLM{}, &s.live, s.requestLimits) + if _, _, err := llm.CreateChatCompletion(context.Background(), req); err != nil { + t.Fatal(err) + } + b, _ := json.Marshal(tools) + want := (len(b) + len(sys)) / 4 + + got := s.SchemaBudget() + if !got.Measured { + t.Fatal("Measured = false after a request") + } + if got.Floor != want { + t.Fatalf("Floor = %d, want %d (tools the request carried)", got.Floor, want) + } + if !got.Error { + t.Fatalf("floor %d over a 1000-token window must set Error", got.Floor) + } +} + +func TestSchemaBudgetUsageCalibration(t *testing.T) { + s := newCompactTestSession(&fakeSummaryLLM{}, 2, nil, nil) + tools := mcpLikeTools(1, 40) + user := strings.Repeat("u", 400) // 100 tokens + req := openai.ChatCompletionRequest{ + Messages: []openai.ChatCompletionMessage{{Role: "system", Content: "sys"}, {Role: "user", Content: user}}, + Tools: tools, + } + llm := trackUsage(&promptUsageLLM{prompt: 700}, &s.live, s.requestLimits) + if _, _, err := llm.CreateChatCompletion(context.Background(), req); err != nil { + t.Fatal(err) + } + if got := s.SchemaBudget().Floor; got != 700-100 { + t.Fatalf("Floor = %d, want the backend overhead %d", got, 600) + } + + // A report below the measurement keeps the measurement. + llm = trackUsage(&promptUsageLLM{prompt: 101}, &s.live, s.requestLimits) + if _, _, err := llm.CreateChatCompletion(context.Background(), req); err != nil { + t.Fatal(err) + } + b, _ := json.Marshal(tools) + if got, want := s.SchemaBudget().Floor, (len(b)+len("sys"))/4; got != want { + t.Fatalf("Floor = %d, want the measurement %d", got, want) + } +} + +// connectTestMCP stands up an in-process MCP server with n tools whose +// descriptions are size bytes, and returns a connected client session. +func connectTestMCP(t *testing.T, name string, n, size int) *sdkmcp.ClientSession { + t.Helper() + ctx := context.Background() + server := sdkmcp.NewServer(&sdkmcp.Implementation{Name: name, Version: "v0"}, nil) + for i := 0; i < n; i++ { + sdkmcp.AddTool(server, &sdkmcp.Tool{Name: fmt.Sprintf("%s_%d", name, i), Description: strings.Repeat("d", size)}, + func(context.Context, *sdkmcp.CallToolRequest, struct{}) (*sdkmcp.CallToolResult, any, error) { + return &sdkmcp.CallToolResult{}, nil, nil + }) + } + serverT, clientT := sdkmcp.NewInMemoryTransports() + if _, err := server.Connect(ctx, serverT, nil); err != nil { + t.Fatal(err) + } + client := sdkmcp.NewClient(&sdkmcp.Implementation{Name: "test", Version: "v0"}, nil) + sess, err := client.Connect(ctx, clientT, nil) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _ = sess.Close() }) + return sess +} + +func TestMCPSchemaCostAttributesAndSorts(t *testing.T) { + s := newReloadTestSession(t) + s.cfgClients["small"] = connectTestMCP(t, "small", 2, 40) + s.cfgClients["big"] = connectTestMCP(t, "big", 5, 2000) + s.cfgClients["mid"] = connectTestMCP(t, "mid", 3, 400) + + costs := s.mcpSchemaCosts(context.Background()) + if len(costs) != 3 { + t.Fatalf("costs = %+v, want 3 servers", costs) + } + names := []string{costs[0].Server, costs[1].Server, costs[2].Server} + if names[0] != "big" || names[1] != "mid" || names[2] != "small" { + t.Fatalf("order = %v, want big, mid, small", names) + } + if costs[0].Tools != 5 || costs[1].Tools != 3 || costs[2].Tools != 2 { + t.Fatalf("tool counts = %+v", costs) + } + if costs[0].Tokens < 5*2000/4 { + t.Fatalf("big tokens = %d, want at least %d", costs[0].Tokens, 5*2000/4) + } + + // Cached until the set of servers changes. + s.cfgClients["late"] = connectTestMCP(t, "late", 1, 10) + if got := s.mcpSchemaCosts(context.Background()); len(got) != 3 { + t.Fatalf("cache not used: %d servers", len(got)) + } + s.cfgServers = map[string]types.MCPServer{} + if err := s.ReconcileMCPServers(map[string]types.MCPServer{}); err != nil { + t.Fatal(err) + } + if got := s.mcpSchemaCosts(context.Background()); len(got) != 0 { + t.Fatalf("after reconcile removed every server: %+v", got) + } +} + +func TestSchemaBudgetEstimateIncludesMCPTools(t *testing.T) { + s := newReloadTestSession(t) + s.compaction = types.CompactionConfig{MaxContextTokens: 100000} + s.cfgClients["big"] = connectTestMCP(t, "big", 4, 1000) + sb := s.SchemaBudget() + if sb.Measured { + t.Fatal("Measured before any request") + } + if len(sb.Servers) != 1 || sb.Floor < sb.Servers[0].Tokens || sb.Servers[0].Tokens < 1000 { + t.Fatalf("estimate blind to MCP tools: %+v", sb) + } +} + +func TestSchemaBudgetNoticeOnce(t *testing.T) { + s := newReloadTestSession(t) + s.compaction = types.CompactionConfig{MaxContextTokens: 100000} + for i, n := range []int{8, 6, 4, 2} { + name := fmt.Sprintf("srv%d", i) + s.cfgClients[name] = connectTestMCP(t, name, n, 4000) + } + var notices []string + s.callbacks.OnStatus = func(m string) { notices = append(notices, m) } + + // Window 100000, budget 100000-min(4096, 50000) = 95904: + // warn above 57542, error above 90000. + s.live.recordFloor(1000) + s.notifySchemaBudget() + if len(notices) != 0 { + t.Fatalf("notice below the threshold: %q", notices) + } + s.live.recordFloor(60000) + s.notifySchemaBudget() + s.notifySchemaBudget() + if len(notices) != 1 { + t.Fatalf("warn notices = %d, want 1: %q", len(notices), notices) + } + n := notices[0] + for _, want := range []string{"60k", "100k", "srv0", "srv1", "srv2", "8 tools", "…", "disable MCP servers"} { + if !strings.Contains(n, want) { + t.Fatalf("notice %q lacks %q", n, want) + } + } + if strings.Contains(n, "srv3") { + t.Fatalf("notice names more than three servers: %q", n) + } + if strings.Index(n, "srv0") > strings.Index(n, "srv1") { + t.Fatalf("largest server not first: %q", n) + } + + s.live.recordFloor(95000) + s.notifySchemaBudget() + s.notifySchemaBudget() + if len(notices) != 2 { + t.Fatalf("notices after error threshold = %d, want 2", len(notices)) + } +} diff --git a/chat/session.go b/chat/session.go index 56e7547..b43e07d 100644 --- a/chat/session.go +++ b/chat/session.go @@ -129,6 +129,14 @@ type Session struct { schemaTools []cogito.ToolDefinitionInterface schemaToolsMu sync.Mutex + // schemaCosts caches mcpSchemaCosts until the set of MCP servers changes; + // schemaNoticeLevel is the highest SchemaBudget level (1 warn, 2 error) + // already reported to the user. Both guarded by schemaCostsMu. + schemaCosts []ServerSchemaCost + schemaCostsValid bool + schemaNoticeLevel int + schemaCostsMu sync.Mutex + agentManager *cogito.AgentManager agentDefs []cogito.AgentDefinition agentModels map[string]bool // models configured per agent type (for the LLM-model guard) @@ -1957,6 +1965,9 @@ func (s *Session) SendMessage(text string, parts ...ContentPart) (string, error) var newFragment cogito.Fragment midTurn.reset() newFragment, err = cogito.ExecuteTools(llm, runFragment, cogitoOpts...) + // The run's requests measured the tool-schema floor; tell the user + // once when it takes a large share of the window. + s.notifySchemaBudget() if err != nil && !errors.Is(err, cogito.ErrNoToolSelected) { // Interrupt (turnCtx cancelled) surfaces here as a context error; // pause the goal so the user's stop sticks and it doesn't re-arm, @@ -2403,12 +2414,14 @@ func (s *Session) ReconcileMCPServers(desired map[string]types.MCPServer) error _ = sess.Close() delete(s.cfgClients, name) delete(s.cfgServers, name) + s.invalidateSchemaCosts() } } for name, srv := range desired { if _, ok := s.cfgClients[name]; ok { continue } + s.invalidateSchemaCosts() transport := wizmcp.TransportForServer(srv) // Deliberately pass s.ctx (the session's long-lived context) directly, // with no per-connect timeout. The go-sdk's mcp.Client.Connect stores the @@ -2438,6 +2451,7 @@ func (s *Session) ReconcileMCPServers(desired map[string]types.MCPServer) error // Called from Reload at turn start (deferred while background sub-agents run), // so closing the old skills client cannot race with a detached agent. func (s *Session) SetSkills(skills []types.Skill) error { + s.invalidateSchemaCosts() if s.skillsClient != nil { _ = s.skillsClient.Close() s.skillsClient = nil From 7de31af352a597597d05e92015ebe07334c9d324 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Fri, 25 Sep 2026 22:38:37 +0000 Subject: [PATCH 14/16] feat(chat): keep the conversation when the tool-schema floor fills the window When the tool schemas and system prompt alone leave no room for the user's message and a minimal answer, no compaction makes a request fit. Overflow recovery compacted, trimmed and basic-compacted the history anyway, and then failed. Overflow recovery now checks the floor before the chain. When the floor plus the user's message and minOutputTokens is over the window, the turn ends with errSchemaFloor: the history is restored to the pre-turn state, the user's message stays on screen, and the message names the largest MCP servers, in the same words as the schema-budget notice. End-of-turn and mid-turn auto-compaction are skipped when the floor plus the last user message is over the context budget. A floor that is only an estimate blocks only when it alone is over the limit. TestSessionAutoCompacts reported 1000 prompt tokens against a 100-token window, which calibrates the floor far over the budget. It now reports 14000 against a 20000-token window: over the trigger, under the budget. Co-Authored-By: Claude Opus 5.5 --- README.md | 6 +- chat/errors.go | 13 ++ chat/midturn.go | 2 +- chat/overflow_floor_test.go | 205 +++++++++++++++++++++++++++++++ chat/session.go | 71 ++++++++++- chat/session_autocompact_test.go | 7 +- selfdoc/README.md | 6 +- 7 files changed, 304 insertions(+), 6 deletions(-) create mode 100644 chat/overflow_floor_test.go diff --git a/README.md b/README.md index 53e655e..7dc2ffa 100644 --- a/README.md +++ b/README.md @@ -584,7 +584,11 @@ log_level: error # not make the history fit, nib stubs old tool outputs, then drops older turns # (saved as an artifact:// the model can read), and as a last resort compacts # without the model. If the retry still does not fit, the turn fails, but the -# compaction is kept, so the next turn does not overflow again. It knows the +# compaction is kept, so the next turn does not overflow again. When the tool +# schemas and system prompt alone fill the window, compaction cannot help: +# nib keeps the conversation as it was, does not compact (neither on +# overflow nor automatically), and names the MCP servers that take the most +# room, so you can disable them or filter their tools. It knows the # wording of llama.cpp/LocalAI, vLLM, OpenAI, # Anthropic and Gemini. overflow_patterns teaches it another backend without a # rebuild: the rows are tried before the built-in ones, first match wins. A diff --git a/chat/errors.go b/chat/errors.go index 0ed57a4..c2c45ca 100644 --- a/chat/errors.go +++ b/chat/errors.go @@ -174,3 +174,16 @@ func contextOverflowRetriedMessage(raw string) string { return "the request is still larger than the model's context window" + overflowDetail(raw) + " after compacting the conversation and retrying. Compacting again will not help: increase the backend's context size, or reduce the enabled tools/MCP servers." } + +// errSchemaFloor marks a turn that failed because the tool schemas and the +// system prompt alone leave no room for a request: no compaction can help, so +// nib kept the conversation instead of compacting it away. +var errSchemaFloor = errors.New("tool schemas and system prompt do not fit the context window") + +// schemaFloorMessage is the user-facing text for errSchemaFloor. It reads +// like the schema-budget notice (sizes, the largest servers, what to do), and +// says the conversation was kept. +func schemaFloorMessage(sb SchemaBudget, window int) string { + return schemaBudgetNotice(sb, window) + + ", or increase the backend's context size. The conversation was kept: compacting it cannot make the request fit." +} diff --git a/chat/midturn.go b/chat/midturn.go index 39e61b7..6beabe5 100644 --- a/chat/midturn.go +++ b/chat/midturn.go @@ -94,7 +94,7 @@ func (c *turnCompactor) manipulate(msgs []openai.ChatCompletionMessage) []openai out := c.s.progressivePrune(view) overhead := c.overhead() - if !c.failed && c.s.shouldCompactNow(estimateTokens(out)+overhead) { + if !c.failed && c.s.shouldCompactNow(estimateTokens(out)+overhead) && !c.s.autoCompactBlocked(lastUserTokens(out)) { out = c.compact(msgs, out, base, repl, overhead) } c.sent = estimateTokens(out) diff --git a/chat/overflow_floor_test.go b/chat/overflow_floor_test.go new file mode 100644 index 0000000..8b5a4e0 --- /dev/null +++ b/chat/overflow_floor_test.go @@ -0,0 +1,205 @@ +package chat + +import ( + "context" + "errors" + "strings" + "sync" + "testing" + + sdkmcp "github.com/modelcontextprotocol/go-sdk/mcp" + "github.com/mudler/cogito" + openai "github.com/sashabaranov/go-openai" +) + +// A floor above the window: compaction cannot make any request fit, so an +// overflow must not summarize, must not touch the history, and must say which +// MCP servers take the room. +func TestOverflowFloorKeepsConversation(t *testing.T) { + llm := &overflowLLM{failures: 99} + s := newOverflowSession(t, llm) + // The overflow teaches a 262144-token window; github alone is about 300k. + s.clients = []*sdkmcp.ClientSession{ + connectTestMCP(t, "github", 60, 20000), + connectTestMCP(t, "jira", 10, 4000), + connectTestMCP(t, "linear", 4, 2000), + connectTestMCP(t, "slack", 1, 100), + } + s.ensureSystemPrompt() + before := append([]openai.ChatCompletionMessage(nil), s.fragment.Messages...) + + _, err := s.SendMessage("what changed?") + if !errors.Is(err, errSchemaFloor) { + t.Fatalf("err = %v, want the schema-floor error", err) + } + msg := err.Error() + for _, want := range []string{"github", "jira", "linear", "kept"} { + if !strings.Contains(msg, want) { + t.Fatalf("message %q does not mention %q", msg, want) + } + } + if strings.Contains(msg, "slack") { + t.Fatalf("message %q names more than the three largest servers", msg) + } + llm.mu.Lock() + asks := llm.asks + llm.mu.Unlock() + if asks != 0 { + t.Fatalf("the summarizer ran %d times; the floor is the problem, not the history", asks) + } + if s.overflowRetries() != 0 { + t.Fatalf("overflow retries = %d, want 0", s.overflowRetries()) + } + got := s.fragment.Messages + if len(got) != len(before) { + t.Fatalf("fragment has %d messages, want the %d from before the turn", len(got), len(before)) + } + for i := range got { + if got[i].Role != before[i].Role || got[i].Content != before[i].Content { + t.Fatalf("message %d changed: %+v", i, got[i]) + } + } + if n := len(s.messages); n == 0 || s.messages[n-1].Role != "user" || s.messages[n-1].Content != "what changed?" { + t.Fatalf("the user's message is not the last in the display copy: %+v", s.messages) + } +} + +// A floor well below the window: recovery runs as before. +func TestOverflowFloorBelowWindowRecovers(t *testing.T) { + llm := &overflowLLM{failures: 1} + s := newOverflowSession(t, llm) + if _, err := s.SendMessage("what changed?"); err != nil { + t.Fatalf("SendMessage: %v", err) + } + llm.mu.Lock() + asks := llm.asks + llm.mu.Unlock() + if asks == 0 || s.overflowRetries() != 1 { + t.Fatalf("asks = %d, retries = %d: recovery did not run", asks, s.overflowRetries()) + } +} + +// An estimate may undercount, so it blocks only when it alone is over the +// limit; a measured floor blocks once it leaves no room. +func TestOverflowFloorEstimateOnly(t *testing.T) { + s := newCompactTestSession(&fakeSummaryLLM{}, 2, nil, nil) // window 1000 + s.systemPrompt = strings.Repeat("s", 3200) // ~800 tokens + + sb, blocked := s.schemaFloorBlocks(1000, 300) + if sb.Measured || blocked { + t.Fatalf("estimate %d below the window blocked (measured %v)", sb.Floor, sb.Measured) + } + s.live.recordFloor(sb.Floor) + if _, blocked := s.schemaFloorBlocks(1000, 300); !blocked { + t.Fatal("a measured floor that leaves no room did not block") + } + + s2 := newCompactTestSession(&fakeSummaryLLM{}, 2, nil, nil) + s2.systemPrompt = strings.Repeat("s", 4400) // ~1100 tokens + if sb, blocked := s2.schemaFloorBlocks(1000, 0); !blocked { + t.Fatalf("estimate %d over the window did not block", sb.Floor) + } +} + +// floorUsageLLM answers every turn request and reports its byte/4 size plus +// extra as prompt tokens, standing in for a large tool-schema floor. +type floorUsageLLM struct { + mu sync.Mutex + extra int + asks int +} + +func (f *floorUsageLLM) summary() string { + f.mu.Lock() + f.asks++ + f.mu.Unlock() + return "summary of earlier turns" +} + +func (f *floorUsageLLM) CreateChatCompletion(ctx context.Context, req openai.ChatCompletionRequest) (cogito.LLMReply, cogito.LLMUsage, error) { + if n := len(req.Messages); n > 0 && strings.HasPrefix(req.Messages[n-1].Content, compactInstruction) { + return replyWith(f.summary()), cogito.LLMUsage{PromptTokens: 20, TotalTokens: 22}, nil + } + p := estimateTokens(req.Messages) + f.extra + return replyWith("ok"), cogito.LLMUsage{PromptTokens: p, CompletionTokens: 1, TotalTokens: p + 1}, nil +} + +func (f *floorUsageLLM) Ask(ctx context.Context, fr cogito.Fragment) (cogito.Fragment, error) { + out := fr.AddMessage("assistant", f.summary()) + if out.Status != nil { + out.Status.LastUsage = cogito.LLMUsage{PromptTokens: 20, TotalTokens: 22} + } + return out, nil +} + +// End-of-turn auto-compaction is skipped when the floor alone is over the +// budget, and runs otherwise. +func TestOverflowFloorEndOfTurnCompaction(t *testing.T) { + run := func(t *testing.T, extra int, history string) (*floorUsageLLM, *Session) { + llm := &floorUsageLLM{extra: extra} + s := newOverflowSession(t, llm) + // Budget 40000 - 4096 = 35904; trigger 28723. + s.compaction.MaxContextTokens = 40000 + s.fragment = cogito.NewFragment( + openai.ChatCompletionMessage{Role: "user", Content: "u1"}, + openai.ChatCompletionMessage{Role: "assistant", Content: "a1 " + history}, + openai.ChatCompletionMessage{Role: "user", Content: "u2"}, + openai.ChatCompletionMessage{Role: "assistant", Content: "a2 " + history}, + ) + if _, err := s.SendMessage("next"); err != nil { + t.Fatalf("SendMessage: %v", err) + } + return llm, s + } + + t.Run("floor over budget", func(t *testing.T) { + llm, s := run(t, 37000, "short") + if llm.asks != 0 { + t.Fatalf("auto-compaction summarized %d times with the floor (%d) over the budget", llm.asks, s.SchemaBudget().Floor) + } + if !fragmentHas(s, func(m openai.ChatCompletionMessage) bool { return m.Content == "u1" }) { + t.Fatal("the history was compacted away") + } + }) + t.Run("floor under budget", func(t *testing.T) { + llm, _ := run(t, 5000, strings.Repeat("earlier detail ", 3500)) // ~13k tokens each + if llm.asks == 0 { + t.Fatal("auto-compaction did not run") + } + }) +} + +// Mid-turn compaction is skipped when the floor alone is over the budget, and +// runs otherwise. +func TestOverflowFloorMidTurnCompaction(t *testing.T) { + msgs := func() []openai.ChatCompletionMessage { + long := strings.Repeat("x", 400) // 100 tokens + return []openai.ChatCompletionMessage{ + {Role: "user", Content: "u1 " + long}, {Role: "assistant", Content: "a1 " + long}, + {Role: "user", Content: "u2 " + long}, {Role: "assistant", Content: "a2 " + long}, + {Role: "user", Content: "u3 " + long}, {Role: "assistant", Content: "a3 " + long}, + } + } + for _, tc := range []struct { + name string + floor int + compact bool + }{ + {"floor over budget", 600, false}, // window 1000, budget 500 + {"floor under budget", 10, true}, + } { + t.Run(tc.name, func(t *testing.T) { + llm := &fakeSummaryLLM{reply: "S"} + s := newCompactTestSession(llm, 2, nil, nil) + s.live.recordFloor(tc.floor) + in := msgs() + out := s.newTurnCompactor(context.Background()).manipulate(in) + if got := llm.calls > 0; got != tc.compact { + t.Fatalf("summarized = %v, want %v", got, tc.compact) + } + if !tc.compact && len(out) != len(in) { + t.Fatalf("request changed: %d messages, want %d", len(out), len(in)) + } + }) + } +} diff --git a/chat/session.go b/chat/session.go index b43e07d..fe8c1df 100644 --- a/chat/session.go +++ b/chat/session.go @@ -1359,6 +1359,49 @@ func (s *Session) fragmentTokens() int { return estimateTokens(s.fragment.Messages) } +// schemaFloorBlocks reports whether the fixed prompt overhead (tool schemas +// and system prompt, see SchemaBudget) leaves no room for the smallest request +// that could succeed: room tokens on top of the floor within limit. When it +// does not, compacting the conversation cannot help, and destroys history +// that was never the problem. +// +// A measured floor is trusted as is. An estimate (no request sent yet) may +// undercount, so it blocks only when it alone is over limit. An unknown limit +// never blocks. +func (s *Session) schemaFloorBlocks(limit, room int) (SchemaBudget, bool) { + sb := s.SchemaBudget() + if limit <= 0 { + return sb, false + } + if !sb.Measured { + return sb, sb.Floor > limit + } + return sb, sb.Floor+room > limit +} + +// autoCompactBlocked reports whether auto-compaction (end of turn or mid-turn) +// must be skipped because the floor plus lastUser tokens is over the context +// budget: the request would stay over the trigger whatever is summarized. +func (s *Session) autoCompactBlocked(lastUser int) bool { + budget := ContextBudget(s.compactionConfig(), s.contextWindow()) + sb, blocked := s.schemaFloorBlocks(budget, lastUser) + if blocked { + xlog.Warn("auto-compaction skipped: the tool schemas and system prompt alone exceed the context budget", + "schema_tokens", sb.Floor, "budget", budget) + } + return blocked +} + +// lastUserTokens is the byte/4 size of the last user message in msgs. +func lastUserTokens(msgs []openai.ChatCompletionMessage) int { + for i := len(msgs) - 1; i >= 0; i-- { + if msgs[i].Role == "user" { + return estimateTokens(msgs[i : i+1]) + } + } + return 0 +} + // noticesText is the message pendingNoticesMessage made for notices, or "" // when there were none. func noticesText(notices []string) string { @@ -2059,6 +2102,30 @@ func (s *Session) SendMessage(text string, parts ...ContentPart) (string, error) first := s.overflowRetried == 0 s.overflowMu.Unlock() + // When the tool schemas and system prompt alone leave no + // room for even the user's message and a minimal answer, no + // compaction can make the request fit. Do not compact away + // the conversation: keep it as it was before the turn, with + // the user's message still on screen, and say what takes + // the room. + if first { + window := s.contextWindow() + if sb, blocked := s.schemaFloorBlocks(window, estimateTokens([]openai.ChatCompletionMessage{turnUser})+minOutputTokens); blocked { + xlog.Warn("overflow recovery skipped: the tool schemas and system prompt alone fill the context window", + "schema_tokens", sb.Floor, "window", window) + err = &FriendlyError{err: fmt.Errorf("%w: %w", errSchemaFloor, err), msg: schemaFloorMessage(sb, window)} + if s.callbacks.OnError != nil { + s.callbacks.OnError(err) + } + s.historyMu.Lock() + s.fragment = preTurnFragment + s.messages = append(slices.Clone(preTurnMessages), openai.ChatCompletionMessage{Role: "user", Content: text}) + s.historyMu.Unlock() + s.restorePendingNotices(notices) + return "", err + } + } + if first { // The escalation chain: iterativeTrim (LLM compaction, // then a smaller tail, a forced prune and a hard @@ -2273,7 +2340,9 @@ func (s *Session) SendMessage(text string, parts ...ContentPart) (string, error) if s.fragment.Status != nil { promptTokens = s.fragment.Status.LastUsage.PromptTokens } - if s.shouldCompactNow(promptTokens) { + // Not when the tool schemas alone are over the budget: the next request + // would cross the trigger again whatever the summary kept. + if s.shouldCompactNow(promptTokens) && !s.autoCompactBlocked(estimateTokens([]openai.ChatCompletionMessage{turnUser})) { if s.callbacks.OnStatus != nil { s.callbacks.OnStatus("Compacting conversation…") } diff --git a/chat/session_autocompact_test.go b/chat/session_autocompact_test.go index 4ebbc65..3e5796b 100644 --- a/chat/session_autocompact_test.go +++ b/chat/session_autocompact_test.go @@ -76,7 +76,10 @@ func autoCompactOpenAI(promptTokens int) http.HandlerFunc { func TestSessionAutoCompacts(t *testing.T) { xlog.SetLogger(xlog.NewLogger(xlog.LogLevel("error"), "")) - srv := httptest.NewServer(autoCompactOpenAI(1000)) + // 14000 is over the 12723 trigger (0.8 of a 20000 window less the 4096 + // reserve) and under the 15904 budget. The backend's count calibrates the + // tool-schema floor, and a floor over the budget skips auto-compaction. + srv := httptest.NewServer(autoCompactOpenAI(14000)) defer srv.Close() var mu sync.Mutex @@ -91,7 +94,7 @@ func TestSessionAutoCompacts(t *testing.T) { ApprovalMode: "auto", AgentOptions: types.AgentOptions{Iterations: 10, MaxAttempts: 3, MaxRetries: 3}, Compaction: types.CompactionConfig{ - MaxContextTokens: 100, // limit = 80 with default 0.8 threshold; 1000 >> 80 + MaxContextTokens: 20000, KeepRecent: 0, }, } diff --git a/selfdoc/README.md b/selfdoc/README.md index 53e655e..7dc2ffa 100644 --- a/selfdoc/README.md +++ b/selfdoc/README.md @@ -584,7 +584,11 @@ log_level: error # not make the history fit, nib stubs old tool outputs, then drops older turns # (saved as an artifact:// the model can read), and as a last resort compacts # without the model. If the retry still does not fit, the turn fails, but the -# compaction is kept, so the next turn does not overflow again. It knows the +# compaction is kept, so the next turn does not overflow again. When the tool +# schemas and system prompt alone fill the window, compaction cannot help: +# nib keeps the conversation as it was, does not compact (neither on +# overflow nor automatically), and names the MCP servers that take the most +# room, so you can disable them or filter their tools. It knows the # wording of llama.cpp/LocalAI, vLLM, OpenAI, # Anthropic and Gemini. overflow_patterns teaches it another backend without a # rebuild: the rows are tried before the built-in ones, first match wins. A From 422c5a39ca6dff281af3b06cd4f03305d50388f1 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Fri, 25 Sep 2026 23:20:46 +0000 Subject: [PATCH 15/16] fix(chat): scale the schema floor by the tokenizer ratio The floor was raised to prompt_tokens minus the estimate of the non-system messages. That residual holds the tool schemas and also the tokenizer skew on the whole conversation, so the floor grew with the history: 5k of schemas under 60k of history on a backend counting 1.5x read as a 35k floor. With a large reserve, auto-compaction was skipped although the history was the problem, and the schema notice reported the skew as tool schemas. The wrapper now learns a ratio, prompt_tokens over the byte/4 estimate of the whole request, clamped to [1, maxTokenizerRatio], and the floor is the measured schemas and system prompt times that ratio. A request with an image part does not calibrate, as the estimate does not count images. The mid-turn trigger still uses the full residual. TestSessionAutoCompacts now runs against a backend that counts 1.5x the estimate, with a large history and half the window reserved. Co-Authored-By: Claude Opus 5.5 --- chat/contextlive.go | 101 ++++++++++++++-------- chat/overflow_floor_test.go | 53 +++++++++++- chat/schema_budget_test.go | 144 ++++++++++++++++++++++++++++--- chat/session_autocompact_test.go | 38 +++++--- 4 files changed, 274 insertions(+), 62 deletions(-) diff --git a/chat/contextlive.go b/chat/contextlive.go index bf5a61d..d570ca1 100644 --- a/chat/contextlive.go +++ b/chat/contextlive.go @@ -3,6 +3,7 @@ package chat import ( "context" "encoding/json" + "math" "sync/atomic" "github.com/mudler/cogito" @@ -33,39 +34,55 @@ import ( // idle session and ContextTokens could not tell which source to trust. // // floor is different: it is the fixed per-request overhead (tool schemas and -// system messages) the wrapper measured on the last request that carried -// tools, raised by the backend's own count when that is larger. It is not -// reset between turns, because it describes the session's configuration, not -// the conversation. See SchemaBudget. +// system messages) the wrapper measured, byte/4, on the last request that +// carried tools. ratio is how many tokens the backend counts per byte/4 +// estimated token, learned from the reported usage; the floor is reported +// scaled by it (see schemaFloor). Neither is reset between turns, because they +// describe the session's configuration and backend, not the conversation. See +// SchemaBudget. type liveUsage struct { prompt atomic.Int64 active atomic.Bool floor atomic.Int64 measured atomic.Bool + ratio atomic.Uint64 // math.Float64bits; 0 means not calibrated yet } -// recordFloor stores the measured fixed overhead of the latest request. +// maxTokenizerRatio caps how many tokens per byte/4 estimated token the floor +// calibration believes a backend counts. +const maxTokenizerRatio = 4.0 + +// recordFloor stores the measured fixed overhead (byte/4) of the latest +// request. func (l *liveUsage) recordFloor(n int) { l.floor.Store(int64(n)) l.measured.Store(true) } -// raiseFloor lifts the latest floor to n when n is larger: the backend's -// count of the same request also includes tokenizer skew. -func (l *liveUsage) raiseFloor(n int) { - for { - cur := l.floor.Load() - if int64(n) <= cur || l.floor.CompareAndSwap(cur, int64(n)) { - return - } +// setRatio stores the tokenizer ratio, clamped to [1, maxTokenizerRatio]. +func (l *liveUsage) setRatio(r float64) { + switch { + case r < 1: + r = 1 + case r > maxTokenizerRatio: + r = maxTokenizerRatio } + l.ratio.Store(math.Float64bits(r)) } -// schemaFloor returns the latest measured floor, and whether any request -// was measured yet. +// tokenizerRatio returns the calibrated ratio, or 1 before any calibration. +func (l *liveUsage) tokenizerRatio() float64 { + if b := l.ratio.Load(); b != 0 { + return math.Float64frombits(b) + } + return 1 +} + +// schemaFloor returns the latest measured floor scaled by the tokenizer +// ratio, and whether any request was measured yet. func (l *liveUsage) schemaFloor() (int, bool) { - return int(l.floor.Load()), l.measured.Load() + return int(float64(l.floor.Load()) * l.tokenizerRatio()), l.measured.Load() } // promptOverhead is how many more tokens the backend reported for a request @@ -133,17 +150,21 @@ type trackedLLM struct { // prepare measures the request's fixed overhead and applies clampOutputTokens // with the session's current limits. The tool schemas are marshalled once for -// both. It returns the byte/4 estimate of the request's non-system messages, -// which calibrate compares with the backend's report, or -1 when the request -// was not measured. +// both. It returns the byte/4 estimate of the whole request (messages, tools +// and system messages), which calibrate compares with the backend's report, +// or -1 when the request must not calibrate. // // Only a request that carries tools is measured: a tool-less request (a -// summary, say) says nothing about the overhead of a turn's requests. +// summary, say) says nothing about the overhead of a turn's requests. A +// request with an image part is measured but does not calibrate: the backend +// counts the image, the byte/4 estimate does not, and the ratio would absorb +// it. func (t *trackedLLM) prepare(request openai.ChatCompletionRequest) (openai.ChatCompletionRequest, int) { toolBytes := toolSchemaBytes(request.Tools) - rest := -1 + whole := -1 if len(request.Tools) > 0 { sysBytes := 0 + image := false var other []openai.ChatCompletionMessage for _, m := range request.Messages { if m.Role == openai.ChatMessageRoleSystem { @@ -151,28 +172,40 @@ func (t *trackedLLM) prepare(request openai.ChatCompletionRequest) (openai.ChatC } else { other = append(other, m) } + for _, p := range m.MultiContent { + if p.Type == openai.ChatMessagePartTypeImageURL { + image = true + } + } + } + fixed := (toolBytes + sysBytes) / 4 + t.live.recordFloor(fixed) + if !image { + whole = estimateTokens(other) + fixed } - t.live.recordFloor((toolBytes + sysBytes) / 4) - rest = estimateTokens(other) } if t.limits == nil { - return request, rest + return request, whole } cap, window := t.limits() - return clampOutputTokensSized(request, cap, window, toolBytes/4), rest + return clampOutputTokensSized(request, cap, window, toolBytes/4), whole } -// calibrate raises the measured floor to the backend's overhead for the -// request prepare measured (rest >= 0) when the backend reported more. -func (t *trackedLLM) calibrate(promptTokens, rest int) { - if rest < 0 || promptTokens <= 0 { +// calibrate sets the tokenizer ratio from the backend's report for a request +// prepare measured (whole >= 0): promptTokens / whole, clamped to +// [1, maxTokenizerRatio]. The floor is the measured overhead times this ratio, +// never the raw residual promptTokens - estimate: that residual also holds the +// skew on the whole conversation, and would make the floor grow with the +// history. +func (t *trackedLLM) calibrate(promptTokens, whole int) { + if whole <= 0 || promptTokens <= 0 { return } - t.live.raiseFloor(promptOverhead(promptTokens, rest)) + t.live.setRatio(float64(promptTokens) / float64(whole)) } func (t *trackedLLM) CreateChatCompletion(ctx context.Context, request openai.ChatCompletionRequest) (cogito.LLMReply, cogito.LLMUsage, error) { - request, rest := t.prepare(request) + request, whole := t.prepare(request) var usage cogito.LLMUsage reply, err := retryRequest(ctx, func() (cogito.LLMReply, error) { r, u, e := t.LLM.CreateChatCompletion(ctx, request) @@ -180,7 +213,7 @@ func (t *trackedLLM) CreateChatCompletion(ctx context.Context, request openai.Ch return r, e }) t.live.record(usage.PromptTokens) - t.calibrate(usage.PromptTokens, rest) + t.calibrate(usage.PromptTokens, whole) return reply, usage, err } @@ -212,7 +245,7 @@ type trackedStreamingLLM struct { // The relay goroutine is bounded by the source channel: it ends when the // provider closes it, which cogito's own consumer loop already depends on. func (t *trackedStreamingLLM) CreateChatCompletionStream(ctx context.Context, request openai.ChatCompletionRequest) (<-chan cogito.StreamEvent, error) { - request, rest := t.prepare(request) + request, whole := t.prepare(request) src, err := retryRequest(ctx, func() (<-chan cogito.StreamEvent, error) { return t.stream.CreateChatCompletionStream(ctx, request) }) @@ -225,7 +258,7 @@ func (t *trackedStreamingLLM) CreateChatCompletionStream(ctx context.Context, re for ev := range src { if ev.Type == cogito.StreamEventDone { t.live.record(ev.Usage.PromptTokens) - t.calibrate(ev.Usage.PromptTokens, rest) + t.calibrate(ev.Usage.PromptTokens, whole) } out <- ev } diff --git a/chat/overflow_floor_test.go b/chat/overflow_floor_test.go index 8b5a4e0..ba6ffac 100644 --- a/chat/overflow_floor_test.go +++ b/chat/overflow_floor_test.go @@ -9,6 +9,7 @@ import ( sdkmcp "github.com/modelcontextprotocol/go-sdk/mcp" "github.com/mudler/cogito" + "github.com/mudler/nib/types" openai "github.com/sashabaranov/go-openai" ) @@ -135,9 +136,12 @@ func (f *floorUsageLLM) Ask(ctx context.Context, fr cogito.Fragment) (cogito.Fra // End-of-turn auto-compaction is skipped when the floor alone is over the // budget, and runs otherwise. func TestOverflowFloorEndOfTurnCompaction(t *testing.T) { - run := func(t *testing.T, extra int, history string) (*floorUsageLLM, *Session) { + run := func(t *testing.T, extra int, sys, history string) (*floorUsageLLM, *Session) { llm := &floorUsageLLM{extra: extra} s := newOverflowSession(t, llm) + if sys != "" { + s.systemPrompt = sys + } // Budget 40000 - 4096 = 35904; trigger 28723. s.compaction.MaxContextTokens = 40000 s.fragment = cogito.NewFragment( @@ -153,7 +157,9 @@ func TestOverflowFloorEndOfTurnCompaction(t *testing.T) { } t.Run("floor over budget", func(t *testing.T) { - llm, s := run(t, 37000, "short") + // A system prompt of 37000 tokens: the floor is the fixed overhead + // itself, not a skew the backend reports on top of the history. + llm, s := run(t, 0, strings.Repeat("s", 37000*4), "short") if llm.asks != 0 { t.Fatalf("auto-compaction summarized %d times with the floor (%d) over the budget", llm.asks, s.SchemaBudget().Floor) } @@ -162,7 +168,7 @@ func TestOverflowFloorEndOfTurnCompaction(t *testing.T) { } }) t.Run("floor under budget", func(t *testing.T) { - llm, _ := run(t, 5000, strings.Repeat("earlier detail ", 3500)) // ~13k tokens each + llm, _ := run(t, 5000, "", strings.Repeat("earlier detail ", 3500)) // ~13k tokens each if llm.asks == 0 { t.Fatal("auto-compaction did not run") } @@ -203,3 +209,44 @@ func TestOverflowFloorMidTurnCompaction(t *testing.T) { }) } } + +// With a large reserve (budget window/2) and a backend counting 1.5x, a long +// history must not read as a schema floor: the history is the problem, so +// end-of-turn auto-compaction runs. The old calibration took the whole +// residual (schemas plus the skew on the history) as the floor and skipped it. +func TestOverflowFloorSkewDoesNotBlockAutoCompaction(t *testing.T) { + s := newCompactTestSession(&fakeSummaryLLM{}, 2, nil, nil) + s.compaction = types.CompactionConfig{MaxContextTokens: 96000, ReserveTokens: 48000, Threshold: 0.8, KeepRecent: 2} + llm := trackUsage(&skewLLM{ratio: 1.5}, &s.live, s.requestLimits) + req, raw := skewRequest(5000, 90000) + _, usage, err := llm.CreateChatCompletion(context.Background(), req) + if err != nil { + t.Fatal(err) + } + if !s.shouldCompactNow(usage.PromptTokens) { + t.Fatalf("%d prompt tokens did not cross the trigger", usage.PromptTokens) + } + if s.autoCompactBlocked(10) { + t.Fatalf("auto-compaction blocked by a floor of %d (schemas about %d x 1.5) with budget 48000", + s.SchemaBudget().Floor, raw) + } +} + +// A tools-heavy request, schemas 80k of a 96k window and no skew, still +// blocks: there the schemas are the problem, whatever is compacted. +func TestOverflowFloorToolsHeavyBlocks(t *testing.T) { + s := newCompactTestSession(&fakeSummaryLLM{}, 2, nil, nil) + s.compaction = types.CompactionConfig{MaxContextTokens: 96000, ReserveTokens: 48000, Threshold: 0.8, KeepRecent: 2} + llm := trackUsage(&skewLLM{ratio: 1}, &s.live, s.requestLimits) + req, raw := skewRequest(80000, 1000) + if _, _, err := llm.CreateChatCompletion(context.Background(), req); err != nil { + t.Fatal(err) + } + sb, blocked := s.schemaFloorBlocks(ContextBudget(s.compactionConfig(), s.contextWindow()), 10) + if !blocked || !sb.Measured || sb.Floor != raw { + t.Fatalf("floor %d (measured %v, want %d) did not block", sb.Floor, sb.Measured, raw) + } + if !s.autoCompactBlocked(10) { + t.Fatal("auto-compaction not blocked by an 80k floor") + } +} diff --git a/chat/schema_budget_test.go b/chat/schema_budget_test.go index 8bb2d1f..fea35ec 100644 --- a/chat/schema_budget_test.go +++ b/chat/schema_budget_test.go @@ -4,6 +4,7 @@ import ( "context" "encoding/json" "fmt" + "slices" "strings" "testing" @@ -147,30 +148,145 @@ func TestSchemaBudgetMeasuredFromRequest(t *testing.T) { } } -func TestSchemaBudgetUsageCalibration(t *testing.T) { - s := newCompactTestSession(&fakeSummaryLLM{}, 2, nil, nil) - tools := mcpLikeTools(1, 40) - user := strings.Repeat("u", 400) // 100 tokens +// skewLLM reports the whole request (messages and tool schemas) at byte/4 +// times ratio, standing in for a backend whose tokenizer counts more than the +// byte/4 estimate. +type skewLLM struct { + cogito.LLM + ratio float64 +} + +func (u *skewLLM) CreateChatCompletion(_ context.Context, req openai.ChatCompletionRequest) (cogito.LLMReply, cogito.LLMUsage, error) { + n := estimateTokens(req.Messages) + toolSchemaBytes(req.Tools)/4 + return cogito.LLMReply{}, cogito.LLMUsage{PromptTokens: int(float64(n) * u.ratio)}, nil +} + +// skewRequest is a request with about schemaTokens of tools and system +// prompt and historyTokens of byte/4 conversation. It returns the byte/4 size +// of the tools and system messages too. +func skewRequest(schemaTokens, historyTokens int) (openai.ChatCompletionRequest, int) { + tools := mcpLikeTools(schemaTokens/1000, 3900) + sys := "you are nib" req := openai.ChatCompletionRequest{ - Messages: []openai.ChatCompletionMessage{{Role: "system", Content: "sys"}, {Role: "user", Content: user}}, - Tools: tools, + Messages: []openai.ChatCompletionMessage{ + {Role: "system", Content: sys}, + {Role: "user", Content: strings.Repeat("u", historyTokens*4)}, + }, + Tools: tools, } - llm := trackUsage(&promptUsageLLM{prompt: 700}, &s.live, s.requestLimits) + return req, (toolSchemaBytes(tools) + len(sys)) / 4 +} + +func near(got, want int) bool { + d := got - want + return d >= -want/100-2 && d <= want/100+2 +} + +// The worked example: 5k of schemas, 60k of history, a backend counting 1.5x. +// The floor is the schemas at 1.5x, not the whole residual, and it does not +// grow with the history. +func TestSchemaBudgetScalesByTokenizerRatio(t *testing.T) { + s := newCompactTestSession(&fakeSummaryLLM{}, 2, nil, nil) + llm := trackUsage(&skewLLM{ratio: 1.5}, &s.live, s.requestLimits) + + req, raw := skewRequest(5000, 60000) if _, _, err := llm.CreateChatCompletion(context.Background(), req); err != nil { t.Fatal(err) } - if got := s.SchemaBudget().Floor; got != 700-100 { - t.Fatalf("Floor = %d, want the backend overhead %d", got, 600) + want := raw * 3 / 2 + if got := s.SchemaBudget().Floor; !near(got, want) { + t.Fatalf("Floor = %d, want about %d (schemas x 1.5), not the residual", got, want) } - // A report below the measurement keeps the measurement. - llm = trackUsage(&promptUsageLLM{prompt: 101}, &s.live, s.requestLimits) + req, _ = skewRequest(5000, 120000) if _, _, err := llm.CreateChatCompletion(context.Background(), req); err != nil { t.Fatal(err) } - b, _ := json.Marshal(tools) - if got, want := s.SchemaBudget().Floor, (len(b)+len("sys"))/4; got != want { - t.Fatalf("Floor = %d, want the measurement %d", got, want) + if got := s.SchemaBudget().Floor; !near(got, want) { + t.Fatalf("Floor = %d after the history doubled, want about %d", got, want) + } +} + +// The ratio is clamped to [1, maxTokenizerRatio]: a report below the estimate +// keeps the measurement, and a wild one cannot multiply the floor without end. +func TestSchemaBudgetRatioClamped(t *testing.T) { + s := newCompactTestSession(&fakeSummaryLLM{}, 2, nil, nil) + req, raw := skewRequest(2000, 1000) + + llm := trackUsage(&skewLLM{ratio: 0.5}, &s.live, s.requestLimits) + if _, _, err := llm.CreateChatCompletion(context.Background(), req); err != nil { + t.Fatal(err) + } + if got := s.SchemaBudget().Floor; got != raw { + t.Fatalf("Floor = %d with a report below the estimate, want the measurement %d", got, raw) + } + + llm = trackUsage(&skewLLM{ratio: 10}, &s.live, s.requestLimits) + if _, _, err := llm.CreateChatCompletion(context.Background(), req); err != nil { + t.Fatal(err) + } + if got, want := s.SchemaBudget().Floor, int(float64(raw)*maxTokenizerRatio); got != want { + t.Fatalf("Floor = %d with a 10x report, want %d (clamped to %vx)", got, want, maxTokenizerRatio) + } +} + +// Image parts are not in the byte/4 estimate, so a request that carries one +// says nothing about the tokenizer: it does not change the ratio. +func TestSchemaBudgetImageKeepsRatio(t *testing.T) { + s := newCompactTestSession(&fakeSummaryLLM{}, 2, nil, nil) + llm := trackUsage(&skewLLM{ratio: 1.5}, &s.live, s.requestLimits) + req, raw := skewRequest(2000, 1000) + if _, _, err := llm.CreateChatCompletion(context.Background(), req); err != nil { + t.Fatal(err) + } + want := s.SchemaBudget().Floor + if !near(want, raw*3/2) { + t.Fatalf("Floor = %d, want about %d", want, raw*3/2) + } + + img := req + img.Messages = append(slices.Clone(req.Messages), openai.ChatCompletionMessage{Role: "user", MultiContent: []openai.ChatMessagePart{ + {Type: openai.ChatMessagePartTypeText, Text: "look"}, + {Type: openai.ChatMessagePartTypeImageURL, ImageURL: &openai.ChatMessageImageURL{URL: "data:image/png;base64,AAAA"}}, + }}) + llm = trackUsage(&skewLLM{ratio: 4}, &s.live, s.requestLimits) + if _, _, err := llm.CreateChatCompletion(context.Background(), img); err != nil { + t.Fatal(err) + } + if got := s.SchemaBudget().Floor; got != want { + t.Fatalf("Floor = %d after a request with an image, want the previous %d", got, want) + } + + // With no earlier calibration, an image request leaves the ratio at 1. + s2 := newCompactTestSession(&fakeSummaryLLM{}, 2, nil, nil) + llm = trackUsage(&skewLLM{ratio: 4}, &s2.live, s2.requestLimits) + if _, _, err := llm.CreateChatCompletion(context.Background(), img); err != nil { + t.Fatal(err) + } + if got := s2.SchemaBudget().Floor; got != raw { + t.Fatalf("Floor = %d, want the unscaled measurement %d", got, raw) + } +} + +// The schema-budget notice reports the corrected size, not the tokenizer skew +// on the whole conversation. +func TestSchemaBudgetNoticeCorrectedSize(t *testing.T) { + s := newCompactTestSession(&fakeSummaryLLM{}, 2, nil, nil) + s.compaction.MaxContextTokens = 20000 // budget 15904, warn above 9542 + var notices []string + s.callbacks.OnStatus = func(m string) { notices = append(notices, m) } + + llm := trackUsage(&skewLLM{ratio: 1.5}, &s.live, s.requestLimits) + req, raw := skewRequest(8000, 60000) + if _, _, err := llm.CreateChatCompletion(context.Background(), req); err != nil { + t.Fatal(err) + } + s.notifySchemaBudget() + if len(notices) != 1 { + t.Fatalf("notices = %q, want one", notices) + } + if want := fmtTokensK(raw * 3 / 2); !strings.Contains(notices[0], "use "+want+" ") { + t.Fatalf("notice %q does not report the corrected %s", notices[0], want) } } diff --git a/chat/session_autocompact_test.go b/chat/session_autocompact_test.go index 3e5796b..26ba24d 100644 --- a/chat/session_autocompact_test.go +++ b/chat/session_autocompact_test.go @@ -5,6 +5,7 @@ import ( "encoding/json" "net/http" "net/http/httptest" + "strings" "sync" "testing" @@ -14,15 +15,25 @@ import ( ) // autoCompactOpenAI is a minimal OpenAI-compatible endpoint that always replies -// with a plain stop message and reports a fixed, high prompt-token usage so the -// auto-compaction threshold is deterministically crossed after a single turn. -func autoCompactOpenAI(promptTokens int) http.HandlerFunc { +// with a plain stop message and reports the request's prompt tokens as ratio +// times its byte/4 size (messages and tool schemas): a backend whose tokenizer +// counts more than the estimate, as real ones do. +func autoCompactOpenAI(ratio float64) http.HandlerFunc { return func(w http.ResponseWriter, r *http.Request) { var req struct { - Stream bool `json:"stream"` + Stream bool `json:"stream"` + Messages []struct { + Content string `json:"content"` + } `json:"messages"` + Tools json.RawMessage `json:"tools"` } body, _ := readAll(r) _ = json.Unmarshal(body, &req) + n := len(req.Tools) + for _, m := range req.Messages { + n += len(m.Content) + } + promptTokens := int(float64(n/4) * ratio) usage := map[string]any{ "prompt_tokens": promptTokens, @@ -71,15 +82,13 @@ func autoCompactOpenAI(promptTokens int) http.HandlerFunc { } // TestSessionAutoCompacts drives a real chat.Session against a fake LLM that -// reports a prompt-token usage above the configured threshold, and proves the -// OnCompactDone callback fires after the turn. +// counts 1.5x the byte/4 estimate, with a large history and a large reserve, +// and proves the OnCompactDone callback fires. The tokenizer skew on the +// history must not be taken for a tool-schema floor that blocks compaction. func TestSessionAutoCompacts(t *testing.T) { xlog.SetLogger(xlog.NewLogger(xlog.LogLevel("error"), "")) - // 14000 is over the 12723 trigger (0.8 of a 20000 window less the 4096 - // reserve) and under the 15904 budget. The backend's count calibrates the - // tool-schema floor, and a floor over the budget skips auto-compaction. - srv := httptest.NewServer(autoCompactOpenAI(14000)) + srv := httptest.NewServer(autoCompactOpenAI(1.5)) defer srv.Close() var mu sync.Mutex @@ -94,7 +103,11 @@ func TestSessionAutoCompacts(t *testing.T) { ApprovalMode: "auto", AgentOptions: types.AgentOptions{Iterations: 10, MaxAttempts: 3, MaxRetries: 3}, Compaction: types.CompactionConfig{ - MaxContextTokens: 20000, + // Budget 20000 (window less half of it), trigger 16000. The + // history below is 40000 byte/4 tokens, 60000 as counted: the + // skew alone on it is 20000, the whole budget. + MaxContextTokens: 40000, + ReserveTokens: 20000, KeepRecent: 0, }, } @@ -115,6 +128,9 @@ func TestSessionAutoCompacts(t *testing.T) { } defer session.Close() + if _, err := session.SendMessage(strings.Repeat("earlier detail ", 40000*4/15)); err != nil { + t.Fatalf("SendMessage: %v", err) + } if _, err := session.SendMessage("hi"); err != nil { t.Fatalf("SendMessage: %v", err) } From 72ea606c8004b3c80042b5b6dc68cb4152cb25fd Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Sat, 26 Sep 2026 07:51:48 +0000 Subject: [PATCH 16/16] feat(chat): classify interrupted streams and truncated tool calls cogito now names why a decision failed. An interrupted stream is retried at most 3 times and invalid tool arguments once, instead of the 10 turn retries every unknown error got. A tool call cut by finish_reason=length is a context problem when the window ran out: prompt plus completion reached 95% of it, or the output clamp lowered max_tokens for that request. nib then retries the turn once with a one-turn note that asks the model to split the call. It compacts first when the prompt is at least a quarter of the window, and lowers the reasoning effort for the turn when the reasoning filled it. The output cap is never lowered. When the cap was reached with room left, the turn fails with a message that names max_tokens. cogito's text for a truncated call says "context window", so the overflow catalog no longer reads it as a rejected request. Co-Authored-By: Claude Opus 5.5 --- README.md | 9 +- chat/contextlive.go | 52 +++++- chat/errors.go | 75 ++++++++ chat/overflow_patterns.go | 7 + chat/retry.go | 42 +++++ chat/retry_toolargs_test.go | 363 ++++++++++++++++++++++++++++++++++++ chat/session.go | 160 +++++++++++++++- go.mod | 2 +- go.sum | 4 +- selfdoc/README.md | 9 +- 10 files changed, 716 insertions(+), 7 deletions(-) create mode 100644 chat/retry_toolargs_test.go diff --git a/README.md b/README.md index 7dc2ffa..1dd043f 100644 --- a/README.md +++ b/README.md @@ -588,7 +588,14 @@ log_level: error # schemas and system prompt alone fill the window, compaction cannot help: # nib keeps the conversation as it was, does not compact (neither on # overflow nor automatically), and names the MCP servers that take the most -# room, so you can disable them or filter their tools. It knows the +# room, so you can disable them or filter their tools. When a tool call is +# cut off because the reply ran out of room in the window (for example a +# whole file in one write), nib retries the turn once with a note that asks +# the model to split the call, after compacting when the prompt is large. +# When the reasoning filled the window, the retry also uses one lower +# reasoning effort. nib never caps the output for this. A stream that ends +# before the reply is finished is retried at most 3 times, and a tool call +# with invalid JSON arguments once. It knows the # wording of llama.cpp/LocalAI, vLLM, OpenAI, # Anthropic and Gemini. overflow_patterns teaches it another backend without a # rebuild: the rows are tried before the built-in ones, first match wins. A diff --git a/chat/contextlive.go b/chat/contextlive.go index d570ca1..d8be410 100644 --- a/chat/contextlive.go +++ b/chat/contextlive.go @@ -4,6 +4,7 @@ import ( "context" "encoding/json" "math" + "slices" "sync/atomic" "github.com/mudler/cogito" @@ -47,6 +48,38 @@ type liveUsage struct { floor atomic.Int64 measured atomic.Bool ratio atomic.Uint64 // math.Float64bits; 0 means not calibrated yet + + // clamped records whether clampOutputTokens lowered the last request's + // max_tokens below the cap. A tool call truncated on such a request ran + // out of window, not of cap (see truncationCause). + clamped atomic.Bool + // note is a one-turn user-role message the next request carries on the + // wire only, then drops: it never reaches the fragment (see + // trackedLLM.prepare). + note atomic.Pointer[string] +} + +// setNote arms the note the next request carries; "" disarms it. +func (l *liveUsage) setNote(n string) { + if n == "" { + l.note.Store(nil) + return + } + l.note.Store(&n) +} + +// takeNote returns the armed note and disarms it. +func (l *liveUsage) takeNote() string { + if p := l.note.Swap(nil); p != nil { + return *p + } + return "" +} + +// lastClamped reports whether the last request's output reservation was +// lowered below the cap by clampOutputTokens. +func (l *liveUsage) lastClamped() bool { + return l.clamped.Load() } // maxTokenizerRatio caps how many tokens per byte/4 estimated token the floor @@ -101,6 +134,8 @@ func promptOverhead(reported, estimated int) int { func (l *liveUsage) begin() { l.prompt.Store(0) l.active.Store(true) + l.clamped.Store(false) + l.note.Store(nil) } // end disarms it, handing authority back to s.fragment — which by then holds @@ -108,6 +143,7 @@ func (l *liveUsage) begin() { func (l *liveUsage) end() { l.active.Store(false) l.prompt.Store(0) + l.note.Store(nil) } // reset drops the in-flight figure without disarming: compaction replaced the @@ -160,6 +196,12 @@ type trackedLLM struct { // counts the image, the byte/4 estimate does not, and the ratio would absorb // it. func (t *trackedLLM) prepare(request openai.ChatCompletionRequest) (openai.ChatCompletionRequest, int) { + // A one-turn note goes on this request only. The fragment is cogito's, + // built from its own messages, so nothing here is stored. + if note := t.live.takeNote(); note != "" { + request.Messages = append(slices.Clip(request.Messages), + openai.ChatCompletionMessage{Role: openai.ChatMessageRoleUser, Content: note}) + } toolBytes := toolSchemaBytes(request.Tools) whole := -1 if len(request.Tools) > 0 { @@ -184,11 +226,19 @@ func (t *trackedLLM) prepare(request openai.ChatCompletionRequest) (openai.ChatC whole = estimateTokens(other) + fixed } } + explicit := request.MaxTokens != 0 || request.MaxCompletionTokens != 0 if t.limits == nil { + t.live.clamped.Store(false) return request, whole } cap, window := t.limits() - return clampOutputTokensSized(request, cap, window, toolBytes/4), whole + out := clampOutputTokensSized(request, cap, window, toolBytes/4) + // An explicit value is cogito's own retry of the previous request with + // a raised cap: it keeps that request's record. + if !explicit { + t.live.clamped.Store(cap > 0 && window > 0 && out.MaxTokens < cap) + } + return out, whole } // calibrate sets the tokenizer ratio from the backend's report for a request diff --git a/chat/errors.go b/chat/errors.go index c2c45ca..9af1b9f 100644 --- a/chat/errors.go +++ b/chat/errors.go @@ -4,6 +4,8 @@ import ( "errors" "fmt" "strings" + + "github.com/mudler/cogito" ) // FriendlyError wraps a noisy backend error with a short, actionable message @@ -32,9 +34,82 @@ func humanizeError(err error) error { if isEmptyReply(err) { return &FriendlyError{err: err, msg: emptyReplyMessage} } + // cogito's sentinels. The turn builds these messages with the figures + // only it knows; this is the fallback for an error it did not rewrite. + if _, done := err.(*FriendlyError); !done { + var te *cogito.ToolArgumentsTruncatedError + switch { + case errors.Is(err, cogito.ErrStreamInterrupted): + return &FriendlyError{err: err, msg: streamInterruptedMessage(0)} + case errors.Is(err, cogito.ErrToolArgumentsInvalid): + return &FriendlyError{err: err, msg: toolArgsInvalidMessage} + case errors.As(err, &te): + return &FriendlyError{err: err, msg: toolArgsTruncatedMessage(te, truncCap)} + } + } return err } +// streamInterruptedMessage is the text for a turn whose stream kept ending +// before the backend finished it. tries is how many times the turn was sent, +// or 0 when unknown. +func streamInterruptedMessage(tries int) string { + n := "" + if tries > 0 { + n = fmt.Sprintf(" (tried %d times)", tries) + } + return "the connection to the model ended before the reply finished" + n + + ". The backend or a proxy in front of it may have timed out on a long reply." +} + +const toolArgsInvalidMessage = "the model produced a tool call with invalid arguments several times. " + + "Try rephrasing, or ask it to split the change into smaller steps." + +// truncationCause is why a tool call was cut by finish_reason=length. +type truncationCause int + +const ( + // truncCap: the output cap was reached with room left in the window. + truncCap truncationCause = iota + // truncToolCall: the window ran out while the model wrote a long call. + truncToolCall + // truncReasoning: the window ran out while the model was reasoning. + truncReasoning +) + +// toolArgsTruncatedMessage is the text for a tool call cut by the output +// limit, by cause. +func toolArgsTruncatedMessage(te *cogito.ToolArgumentsTruncatedError, cause truncationCause) string { + switch cause { + case truncToolCall: + return fmt.Sprintf("the model's call to %s did not fit in the context window (about %d tokens of arguments on a prompt of %d). "+ + "Ask it to make the change in smaller steps.", te.ToolName, te.ArgumentsBytes/4, te.PromptTokens) + case truncReasoning: + return fmt.Sprintf("the model's reasoning filled the context window (about %d tokens on a prompt of %d). "+ + "Ask it to split the task, or lower the reasoning effort.", te.ReasoningBytes/4, te.PromptTokens) + } + limit := "the output limit" + if te.MaxTokens > 0 { + limit = fmt.Sprintf("the output limit (max_tokens %d)", te.MaxTokens) + } + return "the model's tool call was longer than " + limit + + ". Raise the model's max_tokens, or ask for smaller edits." +} + +// truncationNote is the one-turn user-role note sent with the retry after a +// tool call was cut because the window ran out. It is never stored in the +// history. +func truncationNote(te *cogito.ToolArgumentsTruncatedError, cause truncationCause) string { + if cause == truncReasoning { + return fmt.Sprintf("Your previous reply was cut off: its reasoning (about %d tokens) ran out of room in the context window before the call to %s was complete. "+ + "Keep the plan shorter and work in smaller steps (for example write a file in parts, or use edit for targeted changes).", + te.ReasoningBytes/4, te.ToolName) + } + return fmt.Sprintf("Your previous call to %s was cut off after about %d tokens because the reply ran out of room in the context window. "+ + "Split it into smaller calls (for example write the file in parts, or use edit for targeted changes).", + te.ToolName, te.ArgumentsBytes/4) +} + // emptyReplyMarkers are the texts cogito uses when every decision attempt came // back with no text and no tool call: "streaming decision produced no content // (finish_reason=...)" on the streaming path, "no choices: 0" on the other. diff --git a/chat/overflow_patterns.go b/chat/overflow_patterns.go index eafe242..10fbe9e 100644 --- a/chat/overflow_patterns.go +++ b/chat/overflow_patterns.go @@ -7,6 +7,7 @@ import ( "strings" "sync/atomic" + "github.com/mudler/cogito" "github.com/mudler/nib/types" "github.com/mudler/xlog" openai "github.com/sashabaranov/go-openai" @@ -208,6 +209,12 @@ func classifyOverflowWith(err error, lastRequestTokens, window int) overflowInfo if err == nil { return overflowInfo{} } + // A tool call cut by finish_reason=length is not a rejected request, even + // though cogito's text for it says "the context window ran out". The + // turn decides what it means (see truncationCause). + if errors.Is(err, cogito.ErrToolArgumentsTruncated) { + return overflowInfo{} + } status := errorStatus(err) weakOK := status == 0 || status == 400 || status == 413 m := matchOverflow(err, weakOK) diff --git a/chat/retry.go b/chat/retry.go index 397e48e..36f4e8d 100644 --- a/chat/retry.go +++ b/chat/retry.go @@ -2,6 +2,7 @@ package chat import ( "context" + "errors" "fmt" "math/rand/v2" "regexp" @@ -52,6 +53,24 @@ const ( errFatal backendErrorClass = iota errRateLimited errTransient + // errStreamInterrupted: the stream ended before the backend finished it + // (cogito.ErrStreamInterrupted). Retried with backoff, at most + // streamInterruptedBudget times in a row. + errStreamInterrupted + // errToolArgs: the model's tool-call arguments were not valid JSON after + // every cogito attempt (cogito.ErrToolArgumentsInvalid). cogito already + // showed the model its parse error, so the turn is retried at most + // toolArgsBudget time. + errToolArgs +) + +const ( + // streamInterruptedBudget caps the turn retries in a row after an + // interrupted stream. Each retry already made cogito's own attempts, so + // turnRetryBudget (10) would re-send the same long reply dozens of times. + streamInterruptedBudget = 3 + // toolArgsBudget caps the turn retries after invalid tool arguments. + toolArgsBudget = 1 ) const ( @@ -141,6 +160,17 @@ func classifyBackendError(err error) backendErrorClass { if err == nil || classifyOverflow(err).Kind != KindNone { return errFatal } + // cogito names these causes with sentinels; match them before the + // strings. A truncated tool call is handled by the turn itself (see + // truncationCause): repeating the same request truncates it the same way. + switch { + case errors.Is(err, cogito.ErrToolArgumentsTruncated): + return errFatal + case errors.Is(err, cogito.ErrStreamInterrupted): + return errStreamInterrupted + case errors.Is(err, cogito.ErrToolArgumentsInvalid): + return errToolArgs + } low := strings.ToLower(err.Error()) if m := statusRe.FindStringSubmatch(low); m != nil { switch code := m[1]; { @@ -314,6 +344,18 @@ func retryRequest[T any](ctx context.Context, fn func() (T, error)) (T, error) { } } +// classRetryBudget is how many turn retries in a row an error class gets: +// turnRetryBudget, or the class's own budget when that is smaller. +func classRetryBudget(c backendErrorClass) int { + switch c { + case errStreamInterrupted: + return min(streamInterruptedBudget, turnRetryBudget) + case errToolArgs: + return min(toolArgsBudget, turnRetryBudget) + } + return turnRetryBudget +} + // canRetryTurn reports whether the session should run the turn again after // err. It requires the turn context to still be alive: a cancelled turnCtx // means the user pressed Ctrl+C, and re-sending is the opposite of what they diff --git a/chat/retry_toolargs_test.go b/chat/retry_toolargs_test.go new file mode 100644 index 0000000..daf4236 --- /dev/null +++ b/chat/retry_toolargs_test.go @@ -0,0 +1,363 @@ +package chat + +import ( + "context" + "errors" + "fmt" + "strings" + "sync" + "testing" + + "github.com/mudler/cogito" + openai "github.com/sashabaranov/go-openai" +) + +// toolArgsStep answers one turn request of toolArgsLLM. +type toolArgsStep func(req openai.ChatCompletionRequest) (cogito.LLMReply, cogito.LLMUsage, error) + +// toolArgsLLM answers turn requests from script, in order, then with "ok". +// Summary requests are counted apart and always succeed. It records every +// turn request and every reasoning effort it is given. +type toolArgsLLM struct { + mu sync.Mutex + script []toolArgsStep + repeat toolArgsStep // answers every turn request when set + reqs []openai.ChatCompletionRequest + summaries int + efforts []string +} + +func (l *toolArgsLLM) CreateChatCompletion(ctx context.Context, req openai.ChatCompletionRequest) (cogito.LLMReply, cogito.LLMUsage, error) { + if err := ctx.Err(); err != nil { + return cogito.LLMReply{}, cogito.LLMUsage{}, err + } + l.mu.Lock() + if n := len(req.Messages); n > 0 && strings.HasPrefix(req.Messages[n-1].Content, compactInstruction) { + l.summaries++ + l.mu.Unlock() + return replyWith(fmt.Sprintf("summary %d", l.summaries)), cogito.LLMUsage{PromptTokens: 20, TotalTokens: 20}, nil + } + l.reqs = append(l.reqs, req) + var step toolArgsStep + switch { + case l.repeat != nil: + step = l.repeat + case len(l.script) > 0: + step = l.script[0] + l.script = l.script[1:] + } + l.mu.Unlock() + if step != nil { + return step(req) + } + return replyWith("ok"), cogito.LLMUsage{PromptTokens: 10, CompletionTokens: 1, TotalTokens: 11}, nil +} + +func (l *toolArgsLLM) Ask(ctx context.Context, f cogito.Fragment) (cogito.Fragment, error) { + l.mu.Lock() + l.summaries++ + n := l.summaries + l.mu.Unlock() + return f.AddMessage("assistant", fmt.Sprintf("summary %d", n)), nil +} + +func (l *toolArgsLLM) SetReasoningEffort(effort string) { + l.mu.Lock() + l.efforts = append(l.efforts, effort) + l.mu.Unlock() +} + +func (l *toolArgsLLM) requests() []openai.ChatCompletionRequest { + l.mu.Lock() + defer l.mu.Unlock() + return append([]openai.ChatCompletionRequest(nil), l.reqs...) +} + +// truncatedCall is a reply whose tool call was cut by finish_reason=length. +// Like cogito's clients, it reports the output cap the request carried. +func truncatedCall(tool, args, reasoning string, prompt, completion int) toolArgsStep { + return func(req openai.ChatCompletionRequest) (cogito.LLMReply, cogito.LLMUsage, error) { + return cogito.LLMReply{ + ChatCompletionResponse: openai.ChatCompletionResponse{Choices: []openai.ChatCompletionChoice{{ + Message: openai.ChatCompletionMessage{Role: "assistant", ToolCalls: []openai.ToolCall{{ + ID: "call_1", Type: openai.ToolTypeFunction, + Function: openai.FunctionCall{Name: tool, Arguments: args}, + }}}, + FinishReason: openai.FinishReasonLength, + }}}, + ReasoningContent: reasoning, + MaxTokens: max(req.MaxTokens, req.MaxCompletionTokens), + }, cogito.LLMUsage{PromptTokens: prompt, CompletionTokens: completion, TotalTokens: prompt + completion}, + nil + } +} + +// toolArgsSession is newOverflowSession with a 100k window. +func toolArgsSession(t *testing.T, llm cogito.LLM, outputCap int) *Session { + t.Helper() + stubRetrySleep(t) + s := newOverflowSession(t, llm) + s.compaction.MaxContextTokens = 100000 + s.outputCap = outputCap + return s +} + +// lastUserNote returns the content of the request's last message when it is +// the one-turn truncation note, or "". +func lastUserNote(req openai.ChatCompletionRequest) string { + if n := len(req.Messages); n > 0 { + m := req.Messages[n-1] + if m.Role == "user" && strings.Contains(m.Content, "Your previous") { + return m.Content + } + } + return "" +} + +func historyHasNote(s *Session) bool { + s.historyMu.Lock() + defer s.historyMu.Unlock() + for _, m := range s.fragment.Messages { + if strings.Contains(m.Content, "Your previous") { + return true + } + } + for _, m := range s.messages { + if strings.Contains(m.Content, "Your previous") { + return true + } + } + return false +} + +func TestRetryClassifiesCogitoSentinels(t *testing.T) { + interrupted := fmt.Errorf("failed to make a streaming decision after 3 attempts: %w", + fmt.Errorf("localai stream: body ended without [DONE] or a finish_reason: %w", cogito.ErrStreamInterrupted)) + invalid := fmt.Errorf("failed to make a decision after 3 attempts: %w", + fmt.Errorf("%w: tool %q: %v", cogito.ErrToolArgumentsInvalid, "write", "unexpected end of JSON input")) + truncated := fmt.Errorf("failed to make a decision after 1 attempts: %w", &cogito.ToolArgumentsTruncatedError{ + PromptTokens: 73000, CompletionTokens: 27000, MaxTokens: 99000, ToolName: "write", ArgumentsBytes: 108000, + }) + cases := []struct { + name string + err error + want backendErrorClass + }{ + {"stream interrupted", interrupted, errStreamInterrupted}, + {"invalid arguments", invalid, errToolArgs}, + {"truncated arguments are not retried as transient", truncated, errFatal}, + {"old cogito string form", errors.New("failed to make a streaming decision after 3 attempts: unexpected end of JSON input"), errTransient}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + if got := classifyBackendError(tc.err); got != tc.want { + t.Fatalf("classifyBackendError(%v) = %d, want %d", tc.err, got, tc.want) + } + }) + } + if k := classifyOverflow(truncated).Kind; k != KindNone { + t.Fatalf("a truncated tool call classified as overflow kind %v", k) + } +} + +func TestRetryStreamInterruptedBudget(t *testing.T) { + llm := &toolArgsLLM{repeat: func(openai.ChatCompletionRequest) (cogito.LLMReply, cogito.LLMUsage, error) { + return cogito.LLMReply{}, cogito.LLMUsage{}, + fmt.Errorf("localai stream: body ended without [DONE] or a finish_reason: %w", cogito.ErrStreamInterrupted) + }} + s := toolArgsSession(t, llm, 0) + + _, err := s.SendMessage("what changed?") + if err == nil { + t.Fatal("want an error") + } + // Each turn attempt makes requestAttempts calls (one cogito attempt). + if got, want := len(llm.requests()), requestAttempts*(streamInterruptedBudget+1); got != want { + t.Fatalf("requests = %d, want %d (%d turn retries)", got, want, streamInterruptedBudget) + } + if !strings.Contains(err.Error(), "the connection to the model ended before the reply finished") || + !strings.Contains(err.Error(), fmt.Sprintf("tried %d times", streamInterruptedBudget+1)) { + t.Fatalf("err = %q, want the interrupted message", err) + } + if !errors.Is(err, cogito.ErrStreamInterrupted) { + t.Fatal("the friendly error lost the sentinel") + } +} + +func TestToolArgsInvalidRetriedOnce(t *testing.T) { + llm := &toolArgsLLM{repeat: func(openai.ChatCompletionRequest) (cogito.LLMReply, cogito.LLMUsage, error) { + return cogito.LLMReply{ChatCompletionResponse: openai.ChatCompletionResponse{Choices: []openai.ChatCompletionChoice{{ + Message: openai.ChatCompletionMessage{Role: "assistant", ToolCalls: []openai.ToolCall{{ + ID: "call_1", Type: openai.ToolTypeFunction, + Function: openai.FunctionCall{Name: "write", Arguments: `{"path": "a.txt", "content": oops}`}, + }}}, + FinishReason: openai.FinishReasonToolCalls, + }}}}, cogito.LLMUsage{PromptTokens: 100, CompletionTokens: 10}, nil + }} + s := toolArgsSession(t, llm, 0) + + _, err := s.SendMessage("write a file") + if err == nil { + t.Fatal("want an error") + } + if got, want := len(llm.requests()), toolArgsBudget+1; got != want { + t.Fatalf("requests = %d, want %d", got, want) + } + if !strings.Contains(err.Error(), "the model produced a tool call with invalid arguments several times") { + t.Fatalf("err = %q, want the invalid-arguments message", err) + } +} + +// The clamp lowered max_tokens and a tool call was truncated: the window ran +// out, and the prompt is large enough that compaction can make room. +func TestToolArgsTruncatedAfterClampCompacts(t *testing.T) { + llm := &toolArgsLLM{script: []toolArgsStep{ + truncatedCall("write", `{"path":"a.txt","content":"`+strings.Repeat("x", 40000), "", 60000, 10000), + }} + s := toolArgsSession(t, llm, 100000) + + if _, err := s.SendMessage("write the file"); err != nil { + t.Fatalf("SendMessage: %v", err) + } + if llm.summaries == 0 { + t.Fatal("the recovery chain did not run (no summary requested)") + } + if reqs := llm.requests(); len(reqs) != 2 { + t.Fatalf("turn requests = %d, want 2 (one retry)", len(reqs)) + } +} + +// The reported incident: a 73k prompt of a 100k window, below the compaction +// trigger, and a write call truncated at 27k completion tokens. +func TestToolArgsTruncatedIncidentNotesTheRetry(t *testing.T) { + args := `{"path":"big.go","content":"` + strings.Repeat("y", 108000) + llm := &toolArgsLLM{script: []toolArgsStep{truncatedCall("write", args, "", 73000, 27000)}} + s := toolArgsSession(t, llm, 100000) + + if _, err := s.SendMessage("rewrite big.go"); err != nil { + t.Fatalf("SendMessage: %v", err) + } + if llm.summaries == 0 { + t.Fatal("the recovery chain did not run") + } + reqs := llm.requests() + if len(reqs) != 2 { + t.Fatalf("turn requests = %d, want 2", len(reqs)) + } + if lastUserNote(reqs[0]) != "" { + t.Fatal("the first request carried the note") + } + note := lastUserNote(reqs[1]) + if !strings.Contains(note, "Your previous call to write was cut off after about") || + !strings.Contains(note, "Split it into smaller calls") { + t.Fatalf("retry note = %q", note) + } + if historyHasNote(s) { + t.Fatal("the one-turn note was persisted in the history") + } +} + +// Runaway reasoning: a small prompt, the output filled the window. No +// compaction, the output cap unchanged, the reasoning effort lowered for the +// turn only; a second truncation surfaces the reasoning message. +func TestToolArgsRunawayReasoning(t *testing.T) { + reasoning := strings.Repeat("hmm ", 75000) + step := truncatedCall("write", `{"path":"a`, reasoning, 20000, 79000) + llm := &toolArgsLLM{script: []toolArgsStep{step, step}} + s := toolArgsSession(t, llm, 100000) + s.reasoningEffort = "high" + + _, err := s.SendMessage("do the big task") + if err == nil { + t.Fatal("want an error after the second truncation") + } + if llm.summaries != 0 || s.overflowRetries() != 0 { + t.Fatalf("runaway reasoning compacted (summaries %d, retries %d)", llm.summaries, s.overflowRetries()) + } + reqs := llm.requests() + if len(reqs) != 2 { + t.Fatalf("turn requests = %d, want 2", len(reqs)) + } + note := lastUserNote(reqs[1]) + // The cap is unchanged: the retry reserves what the window leaves, less + // only the note's own size. + if cap, _ := s.requestLimits(); cap != 100000 { + t.Fatalf("output cap after the turn = %d, want 100000", cap) + } + if d := reqs[0].MaxTokens - reqs[1].MaxTokens; d > len(note)/4+8 { + t.Fatalf("the retry lowered max_tokens: %d then %d", reqs[0].MaxTokens, reqs[1].MaxTokens) + } + if !strings.Contains(note, "reasoning") || !strings.Contains(note, "smaller steps") { + t.Fatalf("retry note = %q, want the reasoning note", note) + } + llm.mu.Lock() + efforts := append([]string(nil), llm.efforts...) + llm.mu.Unlock() + if len(efforts) != 2 || efforts[0] != "medium" || efforts[1] != "high" { + t.Fatalf("reasoning efforts set = %v, want [medium high] (lowered for the turn, then restored)", efforts) + } + if !strings.Contains(err.Error(), "the model's reasoning filled the context window") { + t.Fatalf("err = %q, want the reasoning message", err) + } + if historyHasNote(s) { + t.Fatal("the one-turn note was persisted in the history") + } +} + +// A long legitimate reply with a small prompt is never capped: no output +// ceiling exists, the request carries the whole cap. +func TestToolArgsNoOutputCeiling(t *testing.T) { + llm := &toolArgsLLM{} + s := toolArgsSession(t, llm, 60000) + + if _, err := s.SendMessage("write a long file"); err != nil { + t.Fatalf("SendMessage: %v", err) + } + if reqs := llm.requests(); len(reqs) != 1 || reqs[0].MaxTokens != 60000 { + t.Fatalf("max_tokens = %v, want the full cap 60000", reqs[0].MaxTokens) + } +} + +// The cap was reached with room left in the window: not exhausted, so no +// compaction and the max_tokens message. +func TestToolArgsTruncatedNotExhausted(t *testing.T) { + llm := &toolArgsLLM{repeat: func(req openai.ChatCompletionRequest) (cogito.LLMReply, cogito.LLMUsage, error) { + return truncatedCall("write", `{"path":"a","content":"zz`, "", 20000, req.MaxTokens)(req) + }} + s := toolArgsSession(t, llm, 20000) + + _, err := s.SendMessage("write it") + if err == nil { + t.Fatal("want an error") + } + if llm.summaries != 0 { + t.Fatalf("compacted %d times, want none", llm.summaries) + } + if !strings.Contains(err.Error(), "longer than the output limit (max_tokens") { + t.Fatalf("err = %q, want the max_tokens message", err) + } + if lastUserNote(llm.requests()[len(llm.requests())-1]) != "" { + t.Fatal("a note was sent although the turn was not retried") + } +} + +// No clamp (unknown limits) and no usage figures: no compaction, the +// max_tokens message. +func TestToolArgsTruncatedWithoutClamp(t *testing.T) { + llm := &toolArgsLLM{repeat: truncatedCall("write", `{"path":"a","content":"zz`, "", 0, 0)} + s := toolArgsSession(t, llm, 0) + + _, err := s.SendMessage("write it") + if err == nil { + t.Fatal("want an error") + } + if llm.summaries != 0 { + t.Fatalf("compacted %d times, want none", llm.summaries) + } + if len(llm.requests()) != 1 { + t.Fatalf("turn requests = %d, want 1 (no retry)", len(llm.requests())) + } + if !strings.Contains(err.Error(), "max_tokens") { + t.Fatalf("err = %q, want the max_tokens message", err) + } +} diff --git a/chat/session.go b/chat/session.go index fe8c1df..0c215f3 100644 --- a/chat/session.go +++ b/chat/session.go @@ -20,6 +20,7 @@ import ( "github.com/mudler/nib/hooks" "github.com/mudler/nib/internal" "github.com/mudler/nib/llmprovider" + "github.com/mudler/nib/llmprovider/catalog" "github.com/mudler/nib/llmprovider/copilot" "github.com/mudler/nib/lsp" "github.com/mudler/nib/manage" @@ -1470,6 +1471,101 @@ func dropFailedDisplay(msgs []openai.ChatCompletionMessage, text string) []opena return msgs } +const ( + // windowExhaustedShare: a truncated tool call whose prompt plus completion + // reached this share of the window stopped at the window, not the cap + // (llama.cpp and LocalAI stop at n_ctx with finish_reason=length). + windowExhaustedShare = 0.95 + // compactableShare: a prompt at least this share of the window is large + // enough that compaction can make room for the call. + compactableShare = 0.25 +) + +// truncationCause decides what a tool call cut by finish_reason=length means. +// The window was exhausted when prompt + completion reached +// windowExhaustedShare of it, or when clampOutputTokens lowered the request's +// max_tokens below the cap (clamped). Otherwise the cap was reached with room +// left: truncCap, which no retry fixes. When exhausted, the cause is the +// reasoning if it took most of the reply, else the call; compact reports +// whether the prompt is large enough for compaction to make room. +func (s *Session) truncationCause(te *cogito.ToolArgumentsTruncatedError, clamped bool) (cause truncationCause, compact bool) { + window := s.contextWindow() + used := te.PromptTokens + te.CompletionTokens + exhausted := clamped || (window > 0 && used > 0 && float64(used) >= windowExhaustedShare*float64(window)) + if !exhausted { + return truncCap, false + } + cause = truncToolCall + if te.ReasoningBytes > 2*te.ArgumentsBytes { + cause = truncReasoning + } + prompt := te.PromptTokens + if prompt <= 0 { + prompt = s.fragmentTokens() + } + return cause, window > 0 && float64(prompt) >= compactableShare*float64(window) +} + +// reasoningEffortSetter is a client whose reasoning effort can be changed. +type reasoningEffortSetter interface { + SetReasoningEffort(effort string) +} + +// minUnknownEffort is the lowest effort lowerEffort picks when the catalog +// does not list the efforts a model accepts: most reasoning models accept +// "low", fewer accept "minimal" or "none". +const minUnknownEffort = "low" + +// lowerEffort returns the effort one level below effort that the model +// accepts (supported, from the catalog; empty when unknown), or "" when there +// is none. +func lowerEffort(effort string, supported []string) string { + idx := slices.Index(types.EffortOrder, strings.ToLower(effort)) + if idx <= 0 { + return "" + } + if len(supported) == 0 { + floor := slices.Index(types.EffortOrder, minUnknownEffort) + if idx-1 < floor { + return "" + } + return types.EffortOrder[idx-1] + } + for i := idx - 1; i >= 0; i-- { + if slices.ContainsFunc(supported, func(s string) bool { return strings.EqualFold(strings.TrimSpace(s), types.EffortOrder[i]) }) { + return types.EffortOrder[i] + } + } + return "" +} + +// lowerReasoningForTurn lowers llm's reasoning effort by one level the model +// accepts and returns the func that restores it. It does nothing (and returns +// a no-op) when the session sets no effort or the client cannot change it. +func (s *Session) lowerReasoningForTurn(llm cogito.LLM) (restore func()) { + noop := func() {} + setter, ok := llm.(reasoningEffortSetter) + if !ok { + return noop + } + provider := s.resolvedSessionProvider() + effort := provider.ReasoningEffort + if effort == "" { + return noop + } + var supported []string + if m, ok := catalog.Lookup(provider.Provider, provider.Model, provider.BaseURL); ok { + supported = m.Compat.SupportedReasoningEfforts + } + lower := lowerEffort(effort, supported) + if lower == "" { + return noop + } + xlog.Warn("reasoning filled the context window; lowering the reasoning effort for this turn", "from", effort, "to", lower) + setter.SetReasoningEffort(lower) + return func() { setter.SetReasoningEffort(effort) } +} + // overflowTrim is the first step of overflow recovery. It is a variable so a // test can make it fail and reach the basicCompact fallback. var overflowTrim = func(s *Session, ctx context.Context) error { return s.iterativeTrim(ctx) } @@ -1732,6 +1828,13 @@ func (s *Session) SendMessage(text string, parts ...ContentPart) (string, error) // Output-cap and budget overflows are retried at most once per turn each; // neither compacts, so neither counts as an overflow recovery. capRetried, budgetRetried := false, false + // Turn retries in a row per error class with its own budget (see + // classRetryBudget), and whether a truncated tool call was retried. + classRetries := map[backendErrorClass]int{} + truncRetried := false + // Restores a reasoning effort lowered for this turn. + restoreEffort := func() {} + defer func() { restoreEffort() }() defer s.endTurn() // Report this turn's own size while it runs; hand authority back to // s.fragment (which compaction may since have shrunk) once it ends. @@ -1817,6 +1920,7 @@ func (s *Session) SendMessage(text string, parts ...ContentPart) (string, error) // retrying client instead (WithAgentLLM below; see agentretry.go): // without it cogito would hand them this tracked one. agentLLM := retryForAgent(llm, &s.agentBackoff) + baseLLM := llm llm = trackUsage(llm, &s.live, s.requestLimits) // Build cogito options from config @@ -2094,6 +2198,41 @@ func (s *Session) SendMessage(text string, parts ...ContentPart) (string, error) turnOverflow.Kind = KindContext } } + // A tool call cut by finish_reason=length. When the window, not + // the cap, ran out, retry once with a note that asks for smaller + // calls: after compacting when the prompt is large enough for + // that to make room, else as it is (never with a lower output + // cap: a long write is legitimate). See truncationCause. + var trunc *cogito.ToolArgumentsTruncatedError + truncNote := "" + var truncErr error + if turnCtx.Err() == nil && turnOverflow.Kind == KindNone && errors.As(err, &trunc) { + cause, compact := s.truncationCause(trunc, s.live.lastClamped()) + friendly := &FriendlyError{err: err, msg: toolArgsTruncatedMessage(trunc, cause)} + switch { + case cause == truncCap || truncRetried: + err = friendly + case compact: + truncRetried = true + turnOverflow.Kind = KindContext + truncNote = truncationNote(trunc, cause) + truncErr = friendly + xlog.Warn("tool call truncated by the context window; compacting and retrying", "tool", trunc.ToolName, "prompt_tokens", trunc.PromptTokens) + default: + truncRetried = true + if cause == truncReasoning { + restoreEffort() + restoreEffort = s.lowerReasoningForTurn(baseLLM) + } + xlog.Warn("reply filled the context window; retrying with a note", "tool", trunc.ToolName, "prompt_tokens", trunc.PromptTokens, "completion_tokens", trunc.CompletionTokens) + s.live.setNote(truncationNote(trunc, cause)) + announce("Reply ran out of room in the context window — retrying in smaller steps…") + announce(retryResumeStatus) + resume, _ := resumableFragment(runFragment, newFragment) + s.commitRun(midTurn, resume) + continue + } + } if turnCtx.Err() == nil && turnOverflow.Kind == KindContext { if w, ok := learnedWindowFrom(err); ok { s.rememberWindow(w, mainModel) @@ -2174,10 +2313,15 @@ func (s *Session) SendMessage(text string, parts ...ContentPart) (string, error) // answer leaves "compacting and retrying" on screen // for the rest of the turn (see retryResumeStatus). announce(retryResumeStatus) + s.live.setNote(truncNote) continue } } } + // The recovery a truncated call asked for did not retry. + if truncErr != nil { + err = truncErr + } // Backend failure recovery: a rate limit or a transient error // (5xx, dropped connection, timeout) must not end a turn that may @@ -2189,11 +2333,14 @@ func (s *Session) SendMessage(text string, parts ...ContentPart) (string, error) // and never after an interrupt: Ctrl+C cancels turnCtx, which // also ends the wait. if canRetryTurn(turnCtx, err) { + class := classifyBackendError(err) resume, progressed := resumableFragment(runFragment, newFragment) if progressed { stalled = 0 + clear(classRetries) } - if stalled < turnRetryBudget { + if stalled < turnRetryBudget && classRetries[class] < classRetryBudget(class) { + classRetries[class]++ wait := turnWait(err, stalled) announce := func(status string) { if s.callbacks.OnStatus != nil { @@ -2213,6 +2360,17 @@ func (s *Session) SendMessage(text string, parts ...ContentPart) (string, error) } } + // cogito's sentinels get a message that names the cause, not the + // raw chain. The stream count is only known here. + if turnCtx.Err() == nil { + switch classifyBackendError(err) { + case errStreamInterrupted: + err = &FriendlyError{err: err, msg: streamInterruptedMessage(classRetries[errStreamInterrupted] + 1)} + case errToolArgs: + err = &FriendlyError{err: err, msg: toolArgsInvalidMessage} + } + } + // Reached only when no retry is happening: either this turn never // qualified for one, or it already spent its single recovery and // overflowed again. In that second case the message must not advise diff --git a/go.mod b/go.mod index c10ca8c..f3ad08d 100644 --- a/go.mod +++ b/go.mod @@ -37,7 +37,7 @@ require ( github.com/msuozzo/bonsai/bonsai-tsx v0.4.0 github.com/msuozzo/bonsai/bonsai-typescript v0.4.0 github.com/msuozzo/bonsai/bonsai-yaml v0.4.0 - github.com/mudler/cogito v0.11.1-0.20260925214156-01190bfd0565 + github.com/mudler/cogito v0.11.1-0.20260926074150-5aa9d6dc8f65 github.com/mudler/xlog v0.0.1 github.com/muesli/termenv v0.16.0 github.com/sashabaranov/go-openai v1.41.2 diff --git a/go.sum b/go.sum index 129f836..632c904 100644 --- a/go.sum +++ b/go.sum @@ -207,8 +207,8 @@ github.com/msuozzo/bonsai/bonsai-typescript v0.4.0 h1:BD6JHfCs3rGNUttL+tMnz/iav1 github.com/msuozzo/bonsai/bonsai-typescript v0.4.0/go.mod h1:DElq3FKqbTkJXcat1o8m98ijL3wUf+jW+DNInzh5Mgo= github.com/msuozzo/bonsai/bonsai-yaml v0.4.0 h1:PJfyfjMQrcWU5e3ib0xGsVuOafNuimf5fK0NhTk8olU= github.com/msuozzo/bonsai/bonsai-yaml v0.4.0/go.mod h1:z8jc0tjXDSOQRKRG0G3uymvVxHdBbrVsWd149nPOyHE= -github.com/mudler/cogito v0.11.1-0.20260925214156-01190bfd0565 h1:A6CT8LINsF5thZlgMCC1iicAfOTdlWjSyURJfw3JpSQ= -github.com/mudler/cogito v0.11.1-0.20260925214156-01190bfd0565/go.mod h1:6sfja3lcu2nWRzEc0wwqGNu/eCG3EWgij+8s7xyUeQ4= +github.com/mudler/cogito v0.11.1-0.20260926074150-5aa9d6dc8f65 h1:y/4ggWOV3Z3OAZhFZwdWR37mzHEEIC/cE+2wigZwdl0= +github.com/mudler/cogito v0.11.1-0.20260926074150-5aa9d6dc8f65/go.mod h1:UxGNMBRakV0A2uVHv2JQIHjpnVPkEWFaS7SDhyynZqQ= github.com/mudler/xlog v0.0.1 h1:yR3/wszd3ZM6u1n96YITJZ4yUcDgqHSwvQmzUJa+8vg= github.com/mudler/xlog v0.0.1/go.mod h1:39f5vcd05Qd6GWKM8IjyHNQ7AmOx3ZM0YfhfIGhC18U= github.com/muesli/ansi v0.0.0-20230316100256-276c6243b2f6 h1:ZK8zHtRHOkbHy6Mmr5D264iyp3TiX5OmNcI5cIARiQI= diff --git a/selfdoc/README.md b/selfdoc/README.md index 7dc2ffa..1dd043f 100644 --- a/selfdoc/README.md +++ b/selfdoc/README.md @@ -588,7 +588,14 @@ log_level: error # schemas and system prompt alone fill the window, compaction cannot help: # nib keeps the conversation as it was, does not compact (neither on # overflow nor automatically), and names the MCP servers that take the most -# room, so you can disable them or filter their tools. It knows the +# room, so you can disable them or filter their tools. When a tool call is +# cut off because the reply ran out of room in the window (for example a +# whole file in one write), nib retries the turn once with a note that asks +# the model to split the call, after compacting when the prompt is large. +# When the reasoning filled the window, the retry also uses one lower +# reasoning effort. nib never caps the output for this. A stream that ends +# before the reply is finished is retried at most 3 times, and a tool call +# with invalid JSON arguments once. It knows the # wording of llama.cpp/LocalAI, vLLM, OpenAI, # Anthropic and Gemini. overflow_patterns teaches it another backend without a # rebuild: the rows are tried before the built-in ones, first match wins. A