mirror of
https://github.com/coder/coder.git
synced 2026-09-24 15:04:27 +08:00
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.
<details>
<summary>Implementation notes</summary>
- 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.
</details>
---
Resolves CODAGT-678
Generated by Coder Agents on behalf of @kylecarbs.
This commit is contained in:
@@ -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"),
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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() {
|
||||
|
||||
@@ -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
|
||||
}
|
||||
@@ -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))
|
||||
})
|
||||
}
|
||||
@@ -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,
|
||||
|
||||
Reference in New Issue
Block a user