From 32217259b73886b182eb06d79eec2d69ae820e90 Mon Sep 17 00:00:00 2001 From: Kyle Carberry Date: Wed, 24 Jun 2026 09:16:38 -0600 Subject: [PATCH] feat: cap tool output to fit the model context window (#26637) ## Problem Local tool results were persisted and replayed to the model verbatim, with no size cap. A single oversized result, most often a multi-megabyte response from an MCP tool, overflows the prompt on the next request. Every retry rebuilds the same history and fails the same way, leaving the chat wedged in `error`. Auto-compaction is reactive (token usage is only known after a response), so it can't catch a single result that blows the very next request. ## Fix Cap every locally-executed tool result at its single choke point, `executeSingleTool` in `chatloop`, so the cap covers built-in tools, **global (deployment-pinned) MCP**, **workspace MCP**, and provider runners uniformly. Because this runs before the result is published to the live stream and before it is committed, the SSE preview, the persisted message, and the model replay all see the same bounded output. The budget is token-aware: a single tool result may use at most half the model's context window (`~4 bytes/token`), with a `16KB` floor and a `64KB` default when the window is unknown. Truncation keeps the head and tail of the output and replaces the middle with a marker telling the model how much was removed and to narrow its query; it is UTF-8 safe and never exceeds the budget. Binary media `Data` is passed through untouched (only the text payload is bounded). A `coderd_chatd_tool_result_truncated_total{provider,model,tool_name}` counter and a warning log record each truncation. ## Out of scope - Provider-executed results (e.g. web search) arrive via the stream, not `executeSingleTool`. - Dynamic/external tool results submitted through the `/tool-results` API are validated as JSON elsewhere. - Cumulative growth across many results is still handled by context compaction; this change only bounds any single result.
Implementation notes - New `coderd/x/chatd/chatloop/tooltruncate.go`: `toolResultByteBudget(contextLimitTokens)` and `truncateToolResultText(text, maxBytes)` (pure, unit-tested). - `chatloop.go`: added `ContextLimit` to `ExecuteLocalToolsOptions`; threaded a computed byte budget through `executeTools` into `executeSingleTool`, where `resp.Content` is capped for the text, media-text, and error branches. - `generation.go`: passes `ContextLimit: prepared.ContextLimitFallback` (the model's configured context limit). - `metrics.go`: new `ToolResultTruncatedTotal` counter + `RecordToolResultTruncated`. - Tunable knobs live as constants in `tooltruncate.go` (`toolResultContextDivisor = 2`, `bytesPerTokenEstimate`, `minToolResultBytes`, `defaultToolResultBytes`). Verified: `go build ./coderd/x/chatd/...`, `go test ./coderd/x/chatd/chatloop/...`, and the `chatd` test binary compiles.
--- Resolves CODAGT-678 Generated by Coder Agents on behalf of @kylecarbs. --- coderd/x/chatd/chatloop/chatloop.go | 36 +++++- .../chatloop/chatloop_run_internal_test.go | 3 + coderd/x/chatd/chatloop/metrics.go | 20 ++++ coderd/x/chatd/chatloop/tooltruncate.go | 101 +++++++++++++++++ .../chatloop/tooltruncate_internal_test.go | 105 ++++++++++++++++++ coderd/x/chatd/generation.go | 1 + docs/admin/integrations/prometheus.md | 1 + scripts/metricsdocgen/generated_metrics | 3 + 8 files changed, 266 insertions(+), 4 deletions(-) create mode 100644 coderd/x/chatd/chatloop/tooltruncate.go create mode 100644 coderd/x/chatd/chatloop/tooltruncate_internal_test.go diff --git a/coderd/x/chatd/chatloop/chatloop.go b/coderd/x/chatd/chatloop/chatloop.go index 64c010ab10..90f6be61af 100644 --- a/coderd/x/chatd/chatloop/chatloop.go +++ b/coderd/x/chatd/chatloop/chatloop.go @@ -254,6 +254,12 @@ type ExecuteLocalToolsOptions struct { ModelProvider string ModelName string + // ContextLimit is the model's context window in tokens. It is used + // to derive a per-result byte budget so a single oversized tool + // result cannot overflow the prompt. Zero means unknown, in which + // case a default budget applies. + ContextLimit int64 + PublishMessagePart func(codersdk.ChatMessageRole, codersdk.ChatMessagePart) Logger slog.Logger Metrics *Metrics @@ -520,6 +526,7 @@ func ExecuteLocalTools(ctx context.Context, opts ExecuteLocalToolsOptions) (Tool }}, nil } + maxResultBytes := toolResultByteBudget(opts.ContextLimit) toolResults := executeTools( ctx, opts.Clock, @@ -532,6 +539,7 @@ func ExecuteLocalTools(ctx context.Context, opts ExecuteLocalToolsOptions) (Tool provider, modelName, opts.BuiltinToolNames, + maxResultBytes, func(tr fantasy.ToolResultContent, completedAt time.Time) { recordToolResultTimestamp(&result, tr.ToolCallID, completedAt) publishToolAttachments(ctx, opts.Logger, tr, completedAt, publishMessagePart) @@ -997,6 +1005,7 @@ func executeTools( logger slog.Logger, provider, model string, builtinToolNames map[string]bool, + maxResultBytes int, onResult func(fantasy.ToolResultContent, time.Time), ) []fantasy.ToolResultContent { if len(toolCalls) == 0 { @@ -1075,6 +1084,7 @@ func executeTools( activeTools, providerRunnerNames, resultProviderMetadata, + maxResultBytes, ) }() } @@ -1194,6 +1204,7 @@ func executeSingleTool( activeTools []string, providerRunnerNames map[string]struct{}, resultProviderMetadata map[string]func(fantasy.ToolResponse) fantasy.ProviderMetadata, + maxResultBytes int, ) fantasy.ToolResultContent { result := fantasy.ToolResultContent{ ToolCallID: tc.ToolCallID, @@ -1254,25 +1265,42 @@ func executeSingleTool( } result.ClientMetadata = resp.Metadata + + // Cap tool output so a single oversized result (most often a large + // MCP response) cannot overflow the model's context window on the + // next request. Only the text payload is bounded; binary media data + // is passed through untouched. + content := resp.Content + if truncated, didTruncate := truncateToolResultText(content, maxResultBytes); didTruncate { + metrics.RecordToolResultTruncated(provider, model, tc.ToolName) + logger.Warn(ctx, "tool result truncated to fit model context", + slog.F("tool_name", tc.ToolName), + slog.F("tool_call_id", tc.ToolCallID), + slog.F("original_bytes", len(content)), + slog.F("max_bytes", maxResultBytes), + ) + content = truncated + } + switch { case resp.IsError: result.Result = fantasy.ToolResultOutputContentError{ - Error: xerrors.New(resp.Content), + Error: xerrors.New(content), } logger.Info(ctx, "tool returned error result", slog.F("tool_name", tc.ToolName), slog.F("tool_call_id", tc.ToolCallID), - slog.F("tool_error", resp.Content), + slog.F("tool_error", content), ) case resp.Type == "image" || resp.Type == "media": result.Result = fantasy.ToolResultOutputContentMedia{ Data: base64.StdEncoding.EncodeToString(resp.Data), MediaType: resp.MediaType, - Text: strings.ToValidUTF8(resp.Content, "\uFFFD"), + Text: strings.ToValidUTF8(content, "\uFFFD"), } default: result.Result = fantasy.ToolResultOutputContentText{ - Text: strings.ToValidUTF8(resp.Content, "\uFFFD"), + Text: strings.ToValidUTF8(content, "\uFFFD"), } } diff --git a/coderd/x/chatd/chatloop/chatloop_run_internal_test.go b/coderd/x/chatd/chatloop/chatloop_run_internal_test.go index 80e77b0aa8..02e3c1f25d 100644 --- a/coderd/x/chatd/chatloop/chatloop_run_internal_test.go +++ b/coderd/x/chatd/chatloop/chatloop_run_internal_test.go @@ -914,6 +914,7 @@ func TestExecuteSingleTool_MediaBase64Encoding(t *testing.T) { []string{"screenshot"}, map[string]struct{}{}, nil, + defaultToolResultBytes, ) media, ok := result.Result.(fantasy.ToolResultOutputContentMedia) @@ -961,6 +962,7 @@ func TestExecuteSingleTool_MediaBase64Encoding(t *testing.T) { []string{"screenshot"}, map[string]struct{}{}, nil, + defaultToolResultBytes, ) media, ok := result.Result.(fantasy.ToolResultOutputContentMedia) @@ -1003,6 +1005,7 @@ func TestExecuteSingleTool_MediaBase64Encoding(t *testing.T) { []string{"echo"}, map[string]struct{}{}, nil, + defaultToolResultBytes, ) textOutput, ok := result.Result.(fantasy.ToolResultOutputContentText) diff --git a/coderd/x/chatd/chatloop/metrics.go b/coderd/x/chatd/chatloop/metrics.go index 6f13663017..6beddf9a54 100644 --- a/coderd/x/chatd/chatloop/metrics.go +++ b/coderd/x/chatd/chatloop/metrics.go @@ -32,6 +32,7 @@ type Metrics struct { MessageCount *prometheus.HistogramVec PromptSizeBytes *prometheus.HistogramVec ToolResultSizeBytes *prometheus.HistogramVec + ToolResultTruncatedTotal *prometheus.CounterVec ToolErrorsTotal *prometheus.CounterVec TTFTSeconds *prometheus.HistogramVec CompactionTotal *prometheus.CounterVec @@ -72,6 +73,12 @@ func NewMetrics(reg prometheus.Registerer) *Metrics { Help: "Size in bytes of each tool execution result.", Buckets: prometheus.ExponentialBuckets(64, 4, 9), // 64B .. 4MB }, []string{"provider", "model", "tool_name"}), + ToolResultTruncatedTotal: factory.NewCounterVec(prometheus.CounterOpts{ + Namespace: metricsNamespace, + Subsystem: metricsSubsystem, + Name: "tool_result_truncated_total", + Help: "Total tool results truncated to fit the model context window.", + }, []string{"provider", "model", "tool_name"}), ToolErrorsTotal: factory.NewCounterVec(prometheus.CounterOpts{ Namespace: metricsNamespace, Subsystem: metricsSubsystem, @@ -165,6 +172,19 @@ func (m *Metrics) RecordToolError(provider, model, toolLabel string) { m.ToolErrorsTotal.WithLabelValues(provider, model, toolLabel).Inc() } +// RecordToolResultTruncated increments tool_result_truncated_total for +// the given tool. No-op when m is nil. An empty tool label is +// normalized to "unknown" to match the other tool metrics. +func (m *Metrics) RecordToolResultTruncated(provider, model, toolLabel string) { + if m == nil { + return + } + if toolLabel == "" { + toolLabel = "unknown" + } + m.ToolResultTruncatedTotal.WithLabelValues(provider, model, toolLabel).Inc() +} + // RecordStreamBufferDropped increments stream_buffer_dropped_total // once per dropped event. No-op when m is nil. func (m *Metrics) RecordStreamBufferDropped() { diff --git a/coderd/x/chatd/chatloop/tooltruncate.go b/coderd/x/chatd/chatloop/tooltruncate.go new file mode 100644 index 0000000000..610ae07ce3 --- /dev/null +++ b/coderd/x/chatd/chatloop/tooltruncate.go @@ -0,0 +1,101 @@ +package chatloop + +import ( + "fmt" + "strings" +) + +const ( + // toolResultContextDivisor bounds how much of a model's context + // window a single tool result may occupy: at most 1/N of the + // window. This caps a single oversized result (most often a large + // MCP response) so it cannot overflow the prompt on its own, while + // still letting a generous amount of output through. Cumulative + // growth across many results is handled separately by context + // compaction. + toolResultContextDivisor = 2 + + // bytesPerTokenEstimate converts a token budget into a byte budget. + // Tool output is capped before tokenization, so this is a coarse, + // provider-agnostic estimate. Roughly 4 bytes per token for typical + // text. + bytesPerTokenEstimate = 4 + + // minToolResultBytes is the floor for the per-result byte budget so + // small or unknown context windows still let useful output through. + minToolResultBytes = 16 << 10 // 16KB + + // defaultToolResultBytes is the budget used when the model's context + // window is unknown (context limit <= 0). + defaultToolResultBytes = 64 << 10 // 64KB + + // truncationMarkerReserve is the number of bytes reserved for the + // truncation marker when splitting output into a head and tail. It + // is comfortably larger than the longest marker + // truncateToolResultText can produce, so the assembled result never + // exceeds the budget. + truncationMarkerReserve = 256 +) + +// maxIntValue is the largest value of the platform int type. +const maxIntValue = int(^uint(0) >> 1) + +// toolResultByteBudget converts a model context-window size (in +// tokens) into the maximum number of bytes a single tool result may +// contribute to the prompt. A context limit <= 0 means the window is +// unknown and the default budget is used. The result is never below +// minToolResultBytes. +func toolResultByteBudget(contextLimitTokens int64) int { + if contextLimitTokens <= 0 { + return defaultToolResultBytes + } + budgetBytes := contextLimitTokens / toolResultContextDivisor * bytesPerTokenEstimate + if budgetBytes < minToolResultBytes { + return minToolResultBytes + } + // Clamp to the int range for 32-bit safety on very large windows. + if budgetBytes > int64(maxIntValue) { + return maxIntValue + } + return int(budgetBytes) +} + +// truncateToolResultText caps text to at most maxBytes using a +// head-and-tail strategy: it keeps the start and end of the output and +// replaces the middle with a marker noting how many bytes were +// removed. This preserves the most useful context (a tool's leading +// summary and trailing status) while bounding size. The returned +// string is always valid UTF-8 and never exceeds maxBytes. It returns +// (text, false) unchanged when maxBytes <= 0 or the text already fits. +func truncateToolResultText(text string, maxBytes int) (string, bool) { + if maxBytes <= 0 || len(text) <= maxBytes { + return text, false + } + + // When the budget is too small to fit the marker plus a meaningful + // head and tail, hard-cut the head and drop any partial trailing + // rune. + if maxBytes <= truncationMarkerReserve*2 { + return strings.ToValidUTF8(text[:maxBytes], ""), true + } + + avail := maxBytes - truncationMarkerReserve + head := avail * 2 / 3 + tail := avail - head + + // ToValidUTF8 with an empty replacement drops invalid byte runs, so + // a cut that lands in the middle of a multi-byte rune is discarded + // rather than corrupting the output. This can only shrink the + // slices, so the assembled result stays within budget. + headStr := strings.ToValidUTF8(text[:head], "") + tailStr := strings.ToValidUTF8(text[len(text)-tail:], "") + removed := len(text) - len(headStr) - len(tailStr) + + marker := fmt.Sprintf( + "\n\n[... Coder truncated %d bytes to fit the model context; "+ + "narrow the query or read a specific file or range for the "+ + "full output ...]\n\n", + removed, + ) + return headStr + marker + tailStr, true +} diff --git a/coderd/x/chatd/chatloop/tooltruncate_internal_test.go b/coderd/x/chatd/chatloop/tooltruncate_internal_test.go new file mode 100644 index 0000000000..f050baa98f --- /dev/null +++ b/coderd/x/chatd/chatloop/tooltruncate_internal_test.go @@ -0,0 +1,105 @@ +package chatloop + +import ( + "strings" + "testing" + "unicode/utf8" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestToolResultByteBudget(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + contextLimit int64 + want int + }{ + {name: "Unknown", contextLimit: 0, want: defaultToolResultBytes}, + {name: "Negative", contextLimit: -1, want: defaultToolResultBytes}, + {name: "BelowFloor", contextLimit: 1000, want: minToolResultBytes}, + { + name: "LargeWindow", + contextLimit: 200_000, + want: 200_000 / toolResultContextDivisor * bytesPerTokenEstimate, + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + assert.Equal(t, tt.want, toolResultByteBudget(tt.contextLimit)) + }) + } + + t.Run("NeverBelowFloor", func(t *testing.T) { + t.Parallel() + for limit := int64(1); limit <= 200_000; limit += 137 { + assert.GreaterOrEqual(t, toolResultByteBudget(limit), minToolResultBytes) + } + }) +} + +func TestTruncateToolResultText(t *testing.T) { + t.Parallel() + + t.Run("UnderLimitUnchanged", func(t *testing.T) { + t.Parallel() + in := "small output" + out, truncated := truncateToolResultText(in, 1024) + assert.False(t, truncated) + assert.Equal(t, in, out) + }) + + t.Run("ExactlyAtLimitUnchanged", func(t *testing.T) { + t.Parallel() + in := strings.Repeat("a", 1024) + out, truncated := truncateToolResultText(in, 1024) + assert.False(t, truncated) + assert.Equal(t, in, out) + }) + + t.Run("ZeroBudgetUnchanged", func(t *testing.T) { + t.Parallel() + in := strings.Repeat("a", 1024) + out, truncated := truncateToolResultText(in, 0) + assert.False(t, truncated) + assert.Equal(t, in, out) + }) + + t.Run("PreservesHeadAndTail", func(t *testing.T) { + t.Parallel() + in := strings.Repeat("A", 1000) + "MIDDLE" + strings.Repeat("B", 1000) + const maxBytes = 600 + out, truncated := truncateToolResultText(in, maxBytes) + require.True(t, truncated) + assert.LessOrEqual(t, len(out), maxBytes) + assert.True(t, utf8.ValidString(out)) + assert.True(t, strings.HasPrefix(out, strings.Repeat("A", 100))) + assert.True(t, strings.HasSuffix(out, strings.Repeat("B", 100))) + assert.NotContains(t, out, "MIDDLE") + assert.Contains(t, out, "truncated") + }) + + t.Run("MultibyteStaysValid", func(t *testing.T) { + t.Parallel() + // Each rune is 3 bytes, so cuts routinely land mid-rune. + in := strings.Repeat("界", 1000) + const maxBytes = 600 + out, truncated := truncateToolResultText(in, maxBytes) + require.True(t, truncated) + assert.LessOrEqual(t, len(out), maxBytes) + assert.True(t, utf8.ValidString(out), "truncated output must be valid UTF-8") + }) + + t.Run("TinyBudgetHardCut", func(t *testing.T) { + t.Parallel() + in := strings.Repeat("界", 10) // 30 bytes + const maxBytes = 10 + out, truncated := truncateToolResultText(in, maxBytes) + require.True(t, truncated) + assert.LessOrEqual(t, len(out), maxBytes) + assert.True(t, utf8.ValidString(out)) + }) +} diff --git a/coderd/x/chatd/generation.go b/coderd/x/chatd/generation.go index ef3a7e7a70..e22c797895 100644 --- a/coderd/x/chatd/generation.go +++ b/coderd/x/chatd/generation.go @@ -669,6 +669,7 @@ func (s *taskStarter) executeLocalTools( BuiltinToolNames: prepared.BuiltinToolNames, ModelProvider: provider, ModelName: modelName, + ContextLimit: prepared.ContextLimitFallback, PublishMessagePart: publish, Logger: s.opts.Logger, Metrics: s.server.metrics, diff --git a/docs/admin/integrations/prometheus.md b/docs/admin/integrations/prometheus.md index 0af63166c5..d4cab1ac1c 100644 --- a/docs/admin/integrations/prometheus.md +++ b/docs/admin/integrations/prometheus.md @@ -232,6 +232,7 @@ deployment. They will always be available from the agent. | `coderd_chatd_stream_retries_total` | counter | Total LLM stream retries. | `chain_broken` `kind` `model` `provider` | | `coderd_chatd_tool_errors_total` | counter | Total tool calls that returned an error result. | `model` `provider` `tool_name` | | `coderd_chatd_tool_result_size_bytes` | histogram | Size in bytes of each tool execution result. | `model` `provider` `tool_name` | +| `coderd_chatd_tool_result_truncated_total` | counter | Total tool results truncated to fit the model context window. | `model` `provider` `tool_name` | | `coderd_chatd_ttft_seconds` | histogram | Time-to-first-token: wall time from LLM request to first streamed chunk. | `model` `provider` | | `coderd_db_query_counts_total` | counter | Total number of queries labelled by HTTP route, method, and query name. | `method` `query` `route` | | `coderd_db_query_latencies_seconds` | histogram | Latency distribution of queries in seconds. | `query` | diff --git a/scripts/metricsdocgen/generated_metrics b/scripts/metricsdocgen/generated_metrics index 90bf2ef862..9f3b9f629b 100644 --- a/scripts/metricsdocgen/generated_metrics +++ b/scripts/metricsdocgen/generated_metrics @@ -298,6 +298,9 @@ coderd_chatd_tool_errors_total{provider="",model="",tool_name=""} 0 # HELP coderd_chatd_tool_result_size_bytes Size in bytes of each tool execution result. # TYPE coderd_chatd_tool_result_size_bytes histogram coderd_chatd_tool_result_size_bytes{provider="",model="",tool_name=""} 0 +# HELP coderd_chatd_tool_result_truncated_total Total tool results truncated to fit the model context window. +# TYPE coderd_chatd_tool_result_truncated_total counter +coderd_chatd_tool_result_truncated_total{provider="",model="",tool_name=""} 0 # HELP coderd_chatd_ttft_seconds Time-to-first-token: wall time from LLM request to first streamed chunk. # TYPE coderd_chatd_ttft_seconds histogram coderd_chatd_ttft_seconds{provider="",model=""} 0