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