From 04834212331e9b3e45e3159266bbb62b9c3dace3 Mon Sep 17 00:00:00 2001 From: Chenyme <118253778+chenyme@users.noreply.github.com> Date: Wed, 15 Jul 2026 16:03:42 +0800 Subject: [PATCH] refactor: enhance token estimation and tool result handling in conversation service --- .../conversation/service_message_context.go | 19 +- .../conversation/service_message_send.go | 91 +++- .../conversation/service_message_usage.go | 182 ++++++++ .../service_message_usage_test.go | 246 +++++++++++ .../conversation/service_tool_execution.go | 412 ++++++++++-------- .../conversation/service_trace_tool_test.go | 135 +++--- 6 files changed, 820 insertions(+), 265 deletions(-) diff --git a/backend/internal/application/conversation/service_message_context.go b/backend/internal/application/conversation/service_message_context.go index faf8c6e9..9b254da0 100644 --- a/backend/internal/application/conversation/service_message_context.go +++ b/backend/internal/application/conversation/service_message_context.go @@ -74,9 +74,24 @@ func estimateMessageTokens(message llm.Message) int64 { for _, part := range message.Parts { tokens += estimateContentPartTokens(part) } - return tokens + estimateTokens(message.ReasoningContent) + } else { + tokens += estimateTokens(message.Content) + } + tokens += estimateTokens(message.ReasoningContent) + for _, call := range message.ToolCalls { + tokens += estimateTokens(call.ToolCallID) + tokens += estimateTokens(call.ToolName) + tokens += estimateTokens(call.ArgumentsJSON) + tokens += 8 + } + for _, result := range message.ToolResults { + tokens += estimateTokens(result.ToolCallID) + tokens += estimateTokens(result.ToolName) + tokens += estimateTokens(result.OutputJSON) + tokens += estimateTokens(result.Error) + tokens += 8 } - return tokens + estimateTokens(message.Content) + estimateTokens(message.ReasoningContent) + return tokens } func estimatePromptTokens(messages []llm.Message) int64 { diff --git a/backend/internal/application/conversation/service_message_send.go b/backend/internal/application/conversation/service_message_send.go index d77eea48..0963246b 100644 --- a/backend/internal/application/conversation/service_message_send.go +++ b/backend/internal/application/conversation/service_message_send.go @@ -1121,29 +1121,53 @@ func (s *Service) sendMessageInternal( } llmCallCount := 1 toolLedger := newToolExecutionLedger() + toolHistoryTrimmedForRun := false for len(upstreamOutput.ToolCalls) > 0 && llmCallCount < maxLLMCalls && remainingToolCalls > 0 { + pendingToolCalls := upstreamOutput.ToolCalls + if len(pendingToolCalls) > remainingToolCalls { + pendingToolCalls = pendingToolCalls[:remainingToolCalls] + } + reasoningContent := "" + if reasoningContentPassback { + reasoningContent = outputReasoningContent(upstreamOutput) + } + assistantToolMessage := llm.Message{ + Role: "assistant", + Content: assistantText, + ReasoningContent: reasoningContent, + ToolCalls: pendingToolCalls, + } + toolResultTokenBudget := resolveToolResultTokenBudget( + generateInput, + llmMessages, + assistantToolMessage, + route.UpstreamModel, + route.ModelCapabilitiesJSON, + ) toolCtx, toolSpan := platformtracing.Start(ctx, "conversation.tool.execute", trace.WithAttributes( attribute.Int64("conversation.id", int64(input.ConversationID)), attribute.Int64("user.id", int64(input.UserID)), attribute.Int("conversation.tool.request_count", len(upstreamOutput.ToolCalls)), attribute.Int("conversation.tool.remaining_count", remainingToolCalls), + attribute.Int64("conversation.tool.result_token_budget", toolResultTokenBudget), ), ) toolResult := s.executeAssistantToolCalls(toolCtx, executeAssistantToolCallsInput{ - UserID: input.UserID, - ConversationID: input.ConversationID, - MessageID: assistantMessage.ID, - RequestID: input.RequestID, - RunID: runID, - ToolCalls: upstreamOutput.ToolCalls, - ToolCallLimit: remainingToolCalls, - TraceRecorder: traceRecorder, - ToolNameMap: toolRuntime.nameMap, - MCPConfigs: toolRuntime.mcpConfigs, - ToolSchemas: toolRuntime.schemas, - Ledger: toolLedger, + UserID: input.UserID, + ConversationID: input.ConversationID, + MessageID: assistantMessage.ID, + RequestID: input.RequestID, + RunID: runID, + ToolCalls: pendingToolCalls, + ToolCallLimit: remainingToolCalls, + TraceRecorder: traceRecorder, + ToolNameMap: toolRuntime.nameMap, + MCPConfigs: toolRuntime.mcpConfigs, + ToolSchemas: toolRuntime.schemas, + Ledger: toolLedger, + ResultTokenBudget: toolResultTokenBudget, }) toolSpan.SetAttributes( attribute.Int("conversation.tool.executed_count", len(toolResult.Rows)), @@ -1163,22 +1187,35 @@ func (s *Service) sendMessageInternal( if len(toolResult.ToolResults) == 0 { break } - reasoningContent := "" - if reasoningContentPassback { - reasoningContent = outputReasoningContent(upstreamOutput) - } + assistantToolMessage.ToolCalls = toolResult.ExecutedToolCalls llmMessages = append(llmMessages, - llm.Message{ - Role: "assistant", - Content: assistantText, - ReasoningContent: reasoningContent, - ToolCalls: toolResult.ExecutedToolCalls, - }, + assistantToolMessage, llm.Message{ Role: "tool", ToolResults: toolResult.ToolResults, }, ) + var toolHistoryTrimmed bool + llmMessages, toolHistoryTrimmed = trimToolFollowUpHistory( + generateInput, + llmMessages, + route.UpstreamModel, + route.ModelCapabilitiesJSON, + ) + if toolHistoryTrimmed { + toolHistoryTrimmedForRun = true + sendSpan.SetAttributes(attribute.Bool("conversation.tool.history_trimmed", true)) + } + var toolResultsRebalanced bool + llmMessages, toolResultsRebalanced = rebalanceToolFollowUpResults( + generateInput, + llmMessages, + route.UpstreamModel, + route.ModelCapabilitiesJSON, + ) + if toolResultsRebalanced { + sendSpan.SetAttributes(attribute.Bool("conversation.tool.results_rebalanced", true)) + } followUpInput := generateInput if llmCallCount+1 >= maxLLMCalls { @@ -1187,7 +1224,7 @@ func (s *Service) sendMessageInternal( followUpInput.DisableTools = true followUpInput.PreviousResponseID = "" applyOpenAIResponsesInstructions(route, routeConfig.Endpoint, &followUpInput) - } else if routeConfig.Endpoint == llm.EndpointResponses && supportsPreviousResponseIDRoute(route) && strings.TrimSpace(upstreamOutput.ResponseID) != "" { + } else if !toolHistoryTrimmed && !toolResultsRebalanced && routeConfig.Endpoint == llm.EndpointResponses && supportsPreviousResponseIDRoute(route) && strings.TrimSpace(upstreamOutput.ResponseID) != "" { followUpInput.PreviousResponseID = strings.TrimSpace(upstreamOutput.ResponseID) followUpInput.Messages = []llm.Message{{Role: "tool", ToolResults: toolResult.ToolResults}} } else { @@ -1284,6 +1321,12 @@ func (s *Service) sendMessageInternal( Tools: toolRuntime.definitions, Options: filteredOptions, }) + responseIDForPersistence := upstreamOutput.ResponseID + // 历史裁剪后的上游 response 不再代表数据库可重建的完整历史，禁止跨轮复用。 + if toolHistoryTrimmedForRun { + responseIDForPersistence = "" + statefulPromptFingerprint = "" + } run.InputTokens = effectiveInputTokens run.OutputTokens = effectiveOutputTokens @@ -1341,7 +1384,7 @@ func (s *Service) sendMessageInternal( OutputTokens: effectiveOutputTokens, ReasoningTokens: totalUsage.ReasoningTokens, AssistantLatency: assistantLatencyMS, - ResponseID: upstreamOutput.ResponseID, + ResponseID: responseIDForPersistence, StatefulPromptFingerprint: statefulPromptFingerprint, ToolCallRows: toolCallRows, PersistedToolCallKeys: persistedToolCallKeys, diff --git a/backend/internal/application/conversation/service_message_usage.go b/backend/internal/application/conversation/service_message_usage.go index d0837fd1..a69e4bf6 100644 --- a/backend/internal/application/conversation/service_message_usage.go +++ b/backend/internal/application/conversation/service_message_usage.go @@ -117,3 +117,185 @@ func estimateToolDefinitionTokens(tools []llm.ToolDefinition) int64 { } return tokens } + +// resolveToolResultTokenBudget 计算当前用户轮次的全部工具结果可使用的模型输入预算。 +// 新批次先使用该上限，回灌前再对同轮全部结果统一分配，不额外透支有效上下文。 +func resolveToolResultTokenBudget( + generateInput llm.GenerateInput, + messages []llm.Message, + pendingAssistant llm.Message, + modelName string, + capabilitiesJSON string, +) int64 { + budgetMessages := toolResultPayloadPlaceholders(prioritizeCurrentToolMessages(messages)) + placeholderResults := make([]llm.ToolResult, 0, len(pendingAssistant.ToolCalls)) + for _, call := range pendingAssistant.ToolCalls { + placeholderResults = append(placeholderResults, llm.ToolResult{ + ToolCallID: call.ToolCallID, + ToolName: call.ToolName, + OutputJSON: "{}", + }) + } + budgetMessages = append( + budgetMessages, + pendingAssistant, + llm.Message{Role: "tool", ToolResults: placeholderResults}, + ) + available := int64(llm.EffectiveContextBudgetFromCapabilities(modelName, capabilitiesJSON)) - + estimateToolFollowUpInputTokens(generateInput, budgetMessages) + if available < 0 { + return 0 + } + return available +} + +// rebalanceToolFollowUpResults 在完整工具回灌请求超预算时，统一压缩当前轮的全部工具结果。 +func rebalanceToolFollowUpResults( + generateInput llm.GenerateInput, + messages []llm.Message, + modelName string, + capabilitiesJSON string, +) ([]llm.Message, bool) { + effectiveBudget := int64(llm.EffectiveContextBudgetFromCapabilities(modelName, capabilitiesJSON)) + if estimateToolFollowUpInputTokens(generateInput, messages) <= effectiveBudget { + return messages, false + } + + _, currentUserIndex := toolHistoryBounds(messages) + if currentUserIndex < 0 { + return messages, false + } + fixedMessages := toolResultPayloadPlaceholders(messages) + resultBudget := effectiveBudget - estimateToolFollowUpInputTokens(generateInput, fixedMessages) + if resultBudget < 0 { + resultBudget = 0 + } + + type resultRef struct { + messageIndex int + resultIndex int + } + result := append([]llm.Message(nil), messages...) + refs := make([]resultRef, 0) + slots := make([]toolExecutionSlot, 0) + for messageIndex := currentUserIndex + 1; messageIndex < len(result); messageIndex++ { + if len(result[messageIndex].ToolResults) == 0 { + continue + } + result[messageIndex].ToolResults = append([]llm.ToolResult(nil), result[messageIndex].ToolResults...) + for resultIndex, toolResult := range result[messageIndex].ToolResults { + refs = append(refs, resultRef{messageIndex: messageIndex, resultIndex: resultIndex}) + slots = append(slots, toolExecutionSlot{result: toolResult}) + } + } + if len(slots) == 0 { + return messages, false + } + + enforceToolResultAggregateBudget(slots, resultBudget) + changed := false + for index, ref := range refs { + if result[ref.messageIndex].ToolResults[ref.resultIndex] != slots[index].result { + changed = true + result[ref.messageIndex].ToolResults[ref.resultIndex] = slots[index].result + } + } + if !changed { + return messages, false + } + return result, true +} + +// toolResultPayloadPlaceholders 保留工具结果的协议结构，但移除可变正文以计算固定上下文开销。 +func toolResultPayloadPlaceholders(messages []llm.Message) []llm.Message { + result := append([]llm.Message(nil), messages...) + for messageIndex := range result { + if len(result[messageIndex].ToolResults) == 0 { + continue + } + placeholders := make([]llm.ToolResult, 0, len(result[messageIndex].ToolResults)) + for _, toolResult := range result[messageIndex].ToolResults { + placeholders = append(placeholders, llm.ToolResult{ + ToolCallID: toolResult.ToolCallID, + ToolName: toolResult.ToolName, + OutputJSON: "{}", + Status: toolResult.Status, + }) + } + result[messageIndex].ToolResults = placeholders + } + return result +} + +// trimToolFollowUpHistory 仅在工具回灌请求超预算时删除最老的完整历史轮次。 +func trimToolFollowUpHistory( + generateInput llm.GenerateInput, + messages []llm.Message, + modelName string, + capabilitiesJSON string, +) ([]llm.Message, bool) { + effectiveBudget := int64(llm.EffectiveContextBudgetFromCapabilities(modelName, capabilitiesJSON)) + estimatedTokens := estimateToolFollowUpInputTokens(generateInput, messages) + if estimatedTokens <= effectiveBudget { + return messages, false + } + + systemEnd, currentUserIndex := toolHistoryBounds(messages) + if currentUserIndex <= systemEnd { + return messages, false + } + for cutFrom := systemEnd; cutFrom < currentUserIndex; cutFrom++ { + estimatedTokens -= estimateMessageTokens(messages[cutFrom]) + nextIndex := cutFrom + 1 + if nextIndex < currentUserIndex && messages[nextIndex].Role != "user" { + continue + } + if estimatedTokens <= effectiveBudget || nextIndex == currentUserIndex { + trimmed := make([]llm.Message, 0, systemEnd+len(messages)-nextIndex) + trimmed = append(trimmed, messages[:systemEnd]...) + trimmed = append(trimmed, messages[nextIndex:]...) + return trimmed, true + } + } + return messages, false +} + +// prioritizeCurrentToolMessages 返回系统指令和当前用户轮次，供工具结果计算最大可用预算。 +func prioritizeCurrentToolMessages(messages []llm.Message) []llm.Message { + systemEnd, currentUserIndex := toolHistoryBounds(messages) + if currentUserIndex <= systemEnd { + return append([]llm.Message(nil), messages...) + } + result := make([]llm.Message, 0, systemEnd+len(messages)-currentUserIndex) + result = append(result, messages[:systemEnd]...) + result = append(result, messages[currentUserIndex:]...) + return result +} + +// toolHistoryBounds 定位系统前缀结束位置和当前轮用户消息。 +func toolHistoryBounds(messages []llm.Message) (int, int) { + systemEnd := 0 + for systemEnd < len(messages) && messages[systemEnd].Role == "system" { + systemEnd++ + } + currentUserIndex := -1 + for index := len(messages) - 1; index >= systemEnd; index-- { + if messages[index].Role == "user" { + currentUserIndex = index + break + } + } + return systemEnd, currentUserIndex +} + +// estimateToolFollowUpInputTokens 按全量请求形状估算工具回灌输入。 +func estimateToolFollowUpInputTokens(generateInput llm.GenerateInput, messages []llm.Message) int64 { + budgetMessages := messages + if strings.TrimSpace(generateInput.Instructions) != "" { + _, budgetMessages = extractOpenAIResponsesInstructions(messages) + } + budgetInput := generateInput + budgetInput.Messages = budgetMessages + budgetInput.PreviousResponseID = "" + return estimateGenerateInputTokens(budgetInput) +} diff --git a/backend/internal/application/conversation/service_message_usage_test.go b/backend/internal/application/conversation/service_message_usage_test.go index c8097a05..44977dcf 100644 --- a/backend/internal/application/conversation/service_message_usage_test.go +++ b/backend/internal/application/conversation/service_message_usage_test.go @@ -1,6 +1,7 @@ package conversation import ( + "strings" "testing" "github.com/DEEIX-AI/DEEIX-Chat/backend/internal/infra/llm" @@ -74,6 +75,251 @@ func TestEstimateGenerateInputTokensIncludesInstructionsAndTools(t *testing.T) { } } +func TestEstimateGenerateInputTokensIncludesToolCallsAndResults(t *testing.T) { + withoutTools := estimatePromptTokens([]llm.Message{{Role: "assistant"}}) + withTools := estimatePromptTokens([]llm.Message{ + { + Role: "assistant", + ToolCalls: []llm.ToolCall{{ + ToolCallID: "call_1", + ToolName: "fetch_transcript", + ArgumentsJSON: `{"url":"https://example.com/video"}`, + }}, + }, + { + Role: "tool", + ToolResults: []llm.ToolResult{{ + ToolCallID: "call_1", + ToolName: "fetch_transcript", + OutputJSON: strings.Repeat("subtitle ", 100), + }}, + }, + }) + if withTools <= withoutTools { + t.Fatalf("expected tool calls and results to increase prompt estimate, got %d <= %d", withTools, withoutTools) + } +} + +func TestResolveToolResultTokenBudgetUsesRemainingModelContext(t *testing.T) { + generateInput := llm.GenerateInput{ + Tools: []llm.ToolDefinition{{ + Name: "fetch_transcript", + Description: "Fetch a video transcript", + InputSchema: []byte(`{"type":"object"}`), + }}, + } + pendingAssistant := llm.Message{ + Role: "assistant", + ToolCalls: []llm.ToolCall{{ + ToolCallID: "call_1", + ToolName: "fetch_transcript", + ArgumentsJSON: "{}", + }}, + } + capabilities := `{"contextWindow":64000,"maxOutputTokens":12000}` + shortBudget := resolveToolResultTokenBudget( + generateInput, + []llm.Message{{Role: "user", Content: "summarize this video"}}, + pendingAssistant, + "custom-model", + capabilities, + ) + longBudget := resolveToolResultTokenBudget( + generateInput, + []llm.Message{{Role: "user", Content: strings.Repeat("history ", 20_000)}}, + pendingAssistant, + "custom-model", + capabilities, + ) + effectiveBudget := int64(llm.EffectiveContextBudgetFromCapabilities("custom-model", capabilities)) + if shortBudget <= 0 || shortBudget >= effectiveBudget { + t.Fatalf("expected short prompt to leave a bounded tool budget, got %d of %d", shortBudget, effectiveBudget) + } + if longBudget >= shortBudget { + t.Fatalf("expected longer prompt to reduce tool result budget, got %d >= %d", longBudget, shortBudget) + } + withOldHistoryBudget := resolveToolResultTokenBudget( + generateInput, + []llm.Message{ + {Role: "user", Content: strings.Repeat("old history ", 20_000)}, + {Role: "assistant", Content: "old answer"}, + {Role: "user", Content: "summarize this video"}, + }, + pendingAssistant, + "custom-model", + capabilities, + ) + if withOldHistoryBudget != shortBudget { + t.Fatalf("expected old history to be reclaimable for the current tool result, got %d want %d", withOldHistoryBudget, shortBudget) + } +} + +func TestResolveToolResultTokenBudgetDoesNotExceedRemainingContext(t *testing.T) { + pendingAssistant := llm.Message{ + Role: "assistant", + ToolCalls: []llm.ToolCall{{ + ToolCallID: "call_1", + ToolName: "fetch_transcript", + ArgumentsJSON: "{}", + }}, + } + capabilities := `{"contextWindow":32000,"maxOutputTokens":4000}` + effectiveBudget := int64(llm.EffectiveContextBudgetFromCapabilities("custom-model", capabilities)) + baseBudget := resolveToolResultTokenBudget( + llm.GenerateInput{}, + []llm.Message{{Role: "user", Content: "summarize this video"}}, + pendingAssistant, + "custom-model", + capabilities, + ) + fillerTokens := baseBudget - 512 + if fillerTokens <= 0 { + t.Fatalf("expected enough baseline budget for boundary test, got %d", baseBudget) + } + messages := []llm.Message{{ + Role: "user", + Content: "summarize this video " + strings.Repeat("x", int(fillerTokens*4)), + }} + remaining := resolveToolResultTokenBudget( + llm.GenerateInput{}, + messages, + pendingAssistant, + "custom-model", + capabilities, + ) + if remaining < 0 || remaining >= 1024 { + t.Fatalf("expected strict sub-1024 remaining budget, got %d", remaining) + } + + placeholder := llm.Message{Role: "tool", ToolResults: []llm.ToolResult{{ + ToolCallID: "call_1", + ToolName: "fetch_transcript", + OutputJSON: "{}", + }}} + fixedTokens := estimateToolFollowUpInputTokens( + llm.GenerateInput{}, + append(messages, pendingAssistant, placeholder), + ) + if fixedTokens+remaining > effectiveBudget { + t.Fatalf("expected tool budget to stay within effective context, got %d > %d", fixedTokens+remaining, effectiveBudget) + } +} + +func TestResolveToolResultTokenBudgetReclaimsPreviousToolPayloads(t *testing.T) { + pendingAssistant := llm.Message{ + Role: "assistant", + ToolCalls: []llm.ToolCall{{ + ToolCallID: "call_2", + ToolName: "fetch_details", + ArgumentsJSON: "{}", + }}, + } + capabilities := `{"contextWindow":32000,"maxOutputTokens":4000}` + withSmallPreviousResult := []llm.Message{ + {Role: "user", Content: "summarize this video"}, + {Role: "assistant", ToolCalls: []llm.ToolCall{{ToolCallID: "call_1", ToolName: "fetch_transcript"}}}, + {Role: "tool", ToolResults: []llm.ToolResult{{ + ToolCallID: "call_1", + ToolName: "fetch_transcript", + OutputJSON: "{}", + }}}, + } + withLargePreviousResult := append([]llm.Message(nil), withSmallPreviousResult...) + withLargePreviousResult[2].ToolResults = []llm.ToolResult{{ + ToolCallID: "call_1", + ToolName: "fetch_transcript", + OutputJSON: strings.Repeat("subtitle ", 20_000), + }} + + smallBudget := resolveToolResultTokenBudget( + llm.GenerateInput{}, + withSmallPreviousResult, + pendingAssistant, + "custom-model", + capabilities, + ) + largeBudget := resolveToolResultTokenBudget( + llm.GenerateInput{}, + withLargePreviousResult, + pendingAssistant, + "custom-model", + capabilities, + ) + if largeBudget != smallBudget { + t.Fatalf("expected previous payloads to be reclaimable, got %d want %d", largeBudget, smallBudget) + } +} + +func TestRebalanceToolFollowUpResultsFitsCurrentRoundWithinContext(t *testing.T) { + capabilities := `{"contextWindow":32000,"maxOutputTokens":4000}` + firstOutput := "FIRST_HEAD " + strings.Repeat("first result ", 10_000) + " FIRST_TAIL" + secondOutput := "SECOND_HEAD " + strings.Repeat("second result ", 10_000) + " SECOND_TAIL" + messages := []llm.Message{ + {Role: "system", Content: "policy"}, + {Role: "user", Content: "compare both tool results"}, + {Role: "assistant", ToolCalls: []llm.ToolCall{{ToolCallID: "call_1", ToolName: "first"}}}, + {Role: "tool", ToolResults: []llm.ToolResult{{ToolCallID: "call_1", ToolName: "first", OutputJSON: firstOutput}}}, + {Role: "assistant", ToolCalls: []llm.ToolCall{{ToolCallID: "call_2", ToolName: "second"}}}, + {Role: "tool", ToolResults: []llm.ToolResult{{ToolCallID: "call_2", ToolName: "second", OutputJSON: secondOutput}}}, + } + + rebalanced, changed := rebalanceToolFollowUpResults( + llm.GenerateInput{}, + messages, + "custom-model", + capabilities, + ) + if !changed { + t.Fatal("expected oversized current tool round to be rebalanced") + } + effectiveBudget := int64(llm.EffectiveContextBudgetFromCapabilities("custom-model", capabilities)) + if tokens := estimateToolFollowUpInputTokens(llm.GenerateInput{}, rebalanced); tokens > effectiveBudget { + t.Fatalf("expected rebalanced follow-up within context, got %d > %d", tokens, effectiveBudget) + } + for _, item := range []struct { + result llm.ToolResult + head string + tail string + }{ + {result: rebalanced[3].ToolResults[0], head: "FIRST_HEAD", tail: "FIRST_TAIL"}, + {result: rebalanced[5].ToolResults[0], head: "SECOND_HEAD", tail: "SECOND_TAIL"}, + } { + if !strings.Contains(item.result.OutputJSON, item.head) || !strings.Contains(item.result.OutputJSON, item.tail) { + t.Fatalf("expected rebalanced result to preserve head and tail, got %q", item.result.OutputJSON) + } + } + if messages[3].ToolResults[0].OutputJSON != firstOutput || messages[5].ToolResults[0].OutputJSON != secondOutput { + t.Fatal("expected rebalancing to leave the source messages unchanged") + } +} + +func TestTrimToolFollowUpHistoryRemovesOldCompleteTurns(t *testing.T) { + capabilities := `{"contextWindow":32000,"maxOutputTokens":4000}` + messages := []llm.Message{ + {Role: "system", Content: "policy"}, + {Role: "user", Content: strings.Repeat("old history ", 20_000)}, + {Role: "assistant", Content: "old answer"}, + {Role: "user", Content: "summarize this video"}, + {Role: "assistant", ToolCalls: []llm.ToolCall{{ToolCallID: "call_1", ToolName: "fetch_transcript"}}}, + {Role: "tool", ToolResults: []llm.ToolResult{{ + ToolCallID: "call_1", + ToolName: "fetch_transcript", + OutputJSON: strings.Repeat("subtitle ", 2000), + }}}, + } + trimmed, changed := trimToolFollowUpHistory(llm.GenerateInput{}, messages, "custom-model", capabilities) + if !changed { + t.Fatal("expected oversized follow-up context to trim old history") + } + if len(trimmed) != 4 || trimmed[0].Role != "system" || trimmed[1].Content != "summarize this video" { + t.Fatalf("expected system prefix and current turn only, got %#v", trimmed) + } + if tokens := estimateToolFollowUpInputTokens(llm.GenerateInput{}, trimmed); tokens > + int64(llm.EffectiveContextBudgetFromCapabilities("custom-model", capabilities)) { + t.Fatalf("expected trimmed follow-up within effective model budget, got %d tokens", tokens) + } +} + func TestSendMessageBillingDurationSeconds(t *testing.T) { if got := sendMessageBillingDurationSeconds(&SendMessageResult{DurationSeconds: 5}, 1200); got != 5 { t.Fatalf("expected explicit duration seconds to win, got %d", got) diff --git a/backend/internal/application/conversation/service_tool_execution.go b/backend/internal/application/conversation/service_tool_execution.go index 104f8155..867330b1 100644 --- a/backend/internal/application/conversation/service_tool_execution.go +++ b/backend/internal/application/conversation/service_tool_execution.go @@ -4,6 +4,7 @@ import ( "context" "encoding/json" "fmt" + "io" "sort" "strings" "time" @@ -13,31 +14,20 @@ import ( "github.com/DEEIX-AI/DEEIX-Chat/backend/internal/infra/mcp" ) -const ( - toolResultModelBudgetChars = 12000 - toolResultReferenceThresholdChars = 50000 - toolResultReferencePreviewChars = 2000 - toolResultAggregateBudgetChars = 200000 -) - -const ( - toolResultBudgetRetainedIn = "server-side tool call record" - toolResultOpaquePreviewMax = 512 -) - type executeAssistantToolCallsInput struct { - UserID uint - ConversationID uint - MessageID uint - RequestID string - RunID string - ToolCalls []llm.ToolCall - ToolCallLimit int - TraceRecorder *messageTraceRecorder - ToolNameMap map[string]string - MCPConfigs map[string]mcp.CallConfig - ToolSchemas map[string]json.RawMessage - Ledger *toolExecutionLedger + UserID uint + ConversationID uint + MessageID uint + RequestID string + RunID string + ToolCalls []llm.ToolCall + ToolCallLimit int + TraceRecorder *messageTraceRecorder + ToolNameMap map[string]string + MCPConfigs map[string]mcp.CallConfig + ToolSchemas map[string]json.RawMessage + Ledger *toolExecutionLedger + ResultTokenBudget int64 } type executeAssistantToolCallsResult struct { @@ -107,7 +97,7 @@ func (s *Service) executeAssistantToolCalls(ctx context.Context, input executeAs row.ErrorJSON = toolNotEnabledForRunMessage(modelToolName) slots[i] = toolExecutionSlot{ row: row, - result: buildToolResultForModel(row, modelToolName, false), + result: buildToolResultForModel(row, modelToolName), } if fatalErr == nil { fatalErr = fmt.Errorf("model requested tool %q, but it is not enabled for this run", modelToolName) @@ -124,7 +114,7 @@ func (s *Service) executeAssistantToolCalls(ctx context.Context, input executeAs row.ErrorJSON = validationErr.Error() slots[i] = toolExecutionSlot{ row: row, - result: buildToolResultForModel(row, modelToolName, false), + result: buildToolResultForModel(row, modelToolName), } if input.Ledger != nil { input.Ledger.store(row.ToolName, row.InputJSON, toolExecutionRecord{row: row, result: slots[i].result}) @@ -137,7 +127,7 @@ func (s *Service) executeAssistantToolCalls(ctx context.Context, input executeAs if previous, ok := input.Ledger.lookup(row.ToolName, row.InputJSON); ok { slot := buildRepeatedToolSlot(row, modelToolName, previous) persisted := s.persistToolCallResult(ctx, &slot.row) - slot.result = buildToolResultForModel(slot.row, modelToolName, persisted) + slot.result = buildToolResultForModel(slot.row, modelToolName) slot.persisted = persisted slots[i] = slot continue @@ -168,7 +158,7 @@ func (s *Service) executeAssistantToolCalls(ctx context.Context, input executeAs } } persisted := s.persistToolCallResult(ctx, &row) - result := buildToolResultForModel(row, modelToolName, persisted) + result := buildToolResultForModel(row, modelToolName) slots[i] = toolExecutionSlot{ row: row, result: result, @@ -182,7 +172,7 @@ func (s *Service) executeAssistantToolCalls(ctx context.Context, input executeAs rows := make([]model.ToolCall, 0, len(slots)) toolResults := make([]llm.ToolResult, 0, len(slots)) persistedToolCallKeys := make(map[string]struct{}) - enforceToolResultAggregateBudget(slots) + enforceToolResultAggregateBudget(slots, input.ResultTokenBudget) for _, slot := range slots { rows = append(rows, slot.row) toolResults = append(toolResults, slot.result) @@ -268,51 +258,103 @@ func (s *Service) persistToolCallResult(ctx context.Context, row *model.ToolCall return row.ID > 0 } -func buildToolResultForModel(row model.ToolCall, modelToolName string, persisted bool) llm.ToolResult { +func buildToolResultForModel(row model.ToolCall, modelToolName string) llm.ToolResult { return llm.ToolResult{ ToolCallID: row.ToolCallID, ToolName: modelToolName, - OutputJSON: budgetToolOutputForModel(row, toolResultModelBudgetChars, persisted), + OutputJSON: modelToolOutputForModel(row.OutputJSON), Status: row.Status, - Error: row.ErrorJSON, + Error: modelToolOutputForModel(row.ErrorJSON), } } -func enforceToolResultAggregateBudget(slots []toolExecutionSlot) { - total := 0 - for _, slot := range slots { - total += len([]rune(strings.TrimSpace(slot.result.OutputJSON))) - } - if total <= toolResultAggregateBudgetChars { +// enforceToolResultAggregateBudget 在模型可用 token 内分配本批工具结果。 +// 小结果优先完整保留，剩余预算由较大结果均分，避免单个工具吞掉整批上下文。 +func enforceToolResultAggregateBudget(slots []toolExecutionSlot, maxTokens int64) { + if len(slots) == 0 { return } - candidates := make([]int, 0, len(slots)) + if maxTokens < 0 { + maxTokens = 0 + } + total := int64(0) + type candidate struct { + index int + tokens int64 + } + candidates := make([]candidate, 0, len(slots)) for index, slot := range slots { - if !slot.persisted || strings.TrimSpace(slot.row.OutputJSON) == "" || strings.HasPrefix(strings.TrimSpace(slot.result.OutputJSON), " right + return candidates[i].tokens < candidates[j].tokens }) - for _, index := range candidates { - if total <= toolResultAggregateBudgetChars { - break + + allocations := make(map[int]int64, len(candidates)) + remaining := maxTokens + for position, item := range candidates { + count := int64(len(candidates) - position) + share := remaining / count + if item.tokens <= share { + allocations[item.index] = item.tokens + remaining -= item.tokens + continue } - current := strings.TrimSpace(slots[index].result.OutputJSON) - replacement := buildPersistedToolOutputForModel(slots[index].row, toolResultReferencePreviewChars) - slots[index].result.OutputJSON = replacement - total = total - len([]rune(current)) + len([]rune(replacement)) + for next := position; next < len(candidates); next++ { + count = int64(len(candidates) - next) + share = remaining / count + allocations[candidates[next].index] = share + remaining -= share + } + break + } + + for index, tokenBudget := range allocations { + if toolResultModelTokens(slots[index].result) <= tokenBudget { + continue + } + applyToolResultTokenBudget(&slots[index].result, tokenBudget) } } +// toolResultModelTokens 估算单个工具结果实际占用的模型 token。 +func toolResultModelTokens(result llm.ToolResult) int64 { + return estimateTokens(result.OutputJSON) + estimateTokens(result.Error) +} + +// applyToolResultTokenBudget 按同一配额约束工具正文和错误信息。 +func applyToolResultTokenBudget(result *llm.ToolResult, maxTokens int64) { + if result == nil { + return + } + outputTokens := estimateTokens(result.OutputJSON) + errorTokens := estimateTokens(result.Error) + total := outputTokens + errorTokens + if total <= maxTokens { + return + } + if outputTokens == 0 { + result.Error = budgetToolOutputForModel(result.Error, maxTokens) + return + } + if errorTokens == 0 { + result.OutputJSON = budgetToolOutputForModel(result.OutputJSON, maxTokens) + return + } + outputBudget := maxTokens * outputTokens / total + result.OutputJSON = budgetToolOutputForModel(result.OutputJSON, outputBudget) + result.Error = budgetToolOutputForModel(result.Error, maxTokens-outputBudget) +} + func toolNotEnabledForRunMessage(toolName string) string { name := strings.TrimSpace(toolName) if name == "" { @@ -321,137 +363,95 @@ func toolNotEnabledForRunMessage(toolName string) string { return fmt.Sprintf("tool %s is not enabled for this run", name) } -func budgetToolOutputForModel(row model.ToolCall, maxChars int, persisted bool) string { - value := strings.TrimSpace(row.OutputJSON) - if value == "" || maxChars <= 0 || len([]rune(value)) <= maxChars { - return value - } - if persisted && len([]rune(value)) > toolResultReferenceThresholdChars { - return buildPersistedToolOutputForModel(row, toolResultReferencePreviewChars) - } - modelText, metadata := budgetToolOutputText(value, maxChars, persisted) - if truncated, _ := metadata["truncated_for_model"].(bool); !truncated { - return modelText - } - note := "\n\n[Tool result truncated for model context.]" - if persisted { - note = "\n\n[Tool result truncated for model context. The full result is retained in the " + toolResultBudgetRetainedIn + ".]" - } - payload := map[string]interface{}{ - "content": []map[string]string{{ - "type": "text", - "text": modelText + note, - }}, - "structuredContent": metadata, - } - if encoded, err := json.Marshal(payload); err == nil { - return string(encoded) - } - return modelText -} - -func buildPersistedToolOutputForModel(row model.ToolCall, previewChars int) string { - output := strings.TrimSpace(row.OutputJSON) - preview := firstCharsAtLineBoundary(output, previewChars) - size := len([]rune(output)) - toolCallID := strings.TrimSpace(row.ToolCallID) - runID := strings.TrimSpace(row.RunID) - toolName := strings.TrimSpace(row.ToolName) - var builder strings.Builder - builder.WriteString(`\n") - builder.WriteString(fmt.Sprintf("Output too large (%d characters). Full output is stored outside the model context in the conversation tool result store.\n\n", size)) - builder.WriteString(fmt.Sprintf("Preview (first %d characters):\n", previewChars)) - builder.WriteString(xmlEscapeText(preview)) - if len([]rune(output)) > len([]rune(preview)) { - builder.WriteString("\n...") - } - builder.WriteString("\n") - return builder.String() -} - -func firstCharsAtLineBoundary(value string, maxChars int) string { - text := strings.TrimSpace(value) - if maxChars <= 0 { +// modelToolOutputForModel 保留可读文本，并从 JSON 中移除不适合进入模型上下文的不透明载荷。 +func modelToolOutputForModel(raw string) string { + value := strings.TrimSpace(raw) + if value == "" { return "" } - runes := []rune(text) - if len(runes) <= maxChars { - return text - } - truncated := string(runes[:maxChars]) - if lastNewline := strings.LastIndex(truncated, "\n"); lastNewline > maxChars/2 { - truncated = truncated[:lastNewline] + if looksLikeOpaqueToolOutput(value) { + return opaqueToolOutputSummary(len([]rune(value))) } - return strings.TrimSpace(truncated) -} - -func budgetToolOutputText(value string, maxChars int, persisted bool) (string, map[string]interface{}) { - normalized := strings.TrimSpace(value) - contentType := "text" - if jsonText, ok := normalizedToolOutputJSON(normalized); ok { - normalized = jsonText - contentType = "json" + if len(value) < 1024 { + return value } - originalChars := len([]rune(value)) - normalizedChars := len([]rune(normalized)) - metadata := map[string]interface{}{ - "truncated_for_model": true, - "original_chars": originalChars, - "model_chars": maxChars, - "content_type": contentType, - "selection": "head_tail", + decoder := json.NewDecoder(strings.NewReader(value)) + decoder.UseNumber() + var payload interface{} + if err := decoder.Decode(&payload); err != nil { + return value } - if persisted { - metadata["retained_in"] = toolResultBudgetRetainedIn + if err := decoder.Decode(&struct{}{}); err != io.EOF { + return value } - if contentType != "json" && looksLikeOpaqueToolOutput(normalized) { - contentType = "opaque" - metadata["content_type"] = contentType - metadata["selection"] = "metadata_preview" - return opaqueToolOutputPreview(normalized, originalChars), metadata + sanitized, changed := sanitizeOpaqueToolOutputJSON(payload) + if !changed { + return value } - if normalizedChars <= maxChars { - metadata["model_chars"] = normalizedChars - metadata["selection"] = "normalized" - metadata["truncated_for_model"] = false - return normalized, metadata + encoded, err := json.Marshal(sanitized) + if err != nil { + return value } - return headTailToolOutput(normalized, maxChars), metadata + return string(encoded) } -func normalizedToolOutputJSON(value string) (string, bool) { - var payload interface{} - if err := json.Unmarshal([]byte(value), &payload); err != nil { - return "", false +// budgetToolOutputForModel 仅在文本超过本批 token 配额时保留头尾片段。 +func budgetToolOutputForModel(value string, maxTokens int64) string { + text := strings.TrimSpace(value) + if text == "" || estimateTokens(text) <= maxTokens { + return text } - formatted, err := json.Marshal(payload) - if err != nil { - return "", false + return headTailToolOutputByTokens(text, maxTokens) +} + +// sanitizeOpaqueToolOutputJSON 递归替换 JSON 内的 base64 等大块不透明字符串。 +func sanitizeOpaqueToolOutputJSON(value interface{}) (interface{}, bool) { + switch item := value.(type) { + case string: + if looksLikeOpaqueToolOutput(item) { + return opaqueToolOutputSummary(len([]rune(item))), true + } + return item, false + case []interface{}: + changed := false + for index, child := range item { + sanitized, childChanged := sanitizeOpaqueToolOutputJSON(child) + if childChanged { + item[index] = sanitized + } + changed = changed || childChanged + } + return item, changed + case map[string]interface{}: + changed := false + for key, child := range item { + sanitized, childChanged := sanitizeOpaqueToolOutputJSON(child) + if childChanged { + item[key] = sanitized + } + changed = changed || childChanged + } + return item, changed + default: + return value, false } - text := strings.TrimSpace(string(formatted)) - return text, text != "" } +// looksLikeOpaqueToolOutput 识别 data URI 和高密度 base64 风格内容。 func looksLikeOpaqueToolOutput(value string) bool { text := strings.TrimSpace(value) runes := []rune(text) - if len(runes) < 1024 || strings.ContainsAny(text, " \n\t{}[],:") { + if len(runes) < 1024 { + return false + } + prefix := text + if len(prefix) > 128 { + prefix = prefix[:128] + } + if strings.HasPrefix(strings.ToLower(prefix), "data:") && strings.Contains(strings.ToLower(prefix), ";base64,") { + return true + } + if strings.ContainsAny(text, " \n\t{}[],:") { return false } base64ish := 0 @@ -466,9 +466,77 @@ func looksLikeOpaqueToolOutput(value string) bool { return float64(base64ish)/float64(len(runes)) > 0.95 } -func opaqueToolOutputPreview(value string, originalChars int) string { - preview := headTailToolOutput(value, toolResultOpaquePreviewMax) - return fmt.Sprintf("Large opaque tool result omitted from model context.\nOriginal characters: %d\nPreview:\n%s", originalChars, preview) +// opaqueToolOutputSummary 为被移除的不透明载荷生成可读说明。 +func opaqueToolOutputSummary(originalChars int) string { + return fmt.Sprintf("[Opaque tool payload omitted from model context: %d characters]", originalChars) +} + +// headTailToolOutputByTokens 按项目 token 估算保留文本头尾。 +func headTailToolOutputByTokens(value string, maxTokens int64) string { + text := strings.TrimSpace(value) + if text == "" || maxTokens <= 0 { + return "" + } + if estimateTokens(text) <= maxTokens { + return text + } + runes := []rune(text) + marker := fmt.Sprintf("\n\n[... %d characters omitted to fit the model context ...]\n\n", len(runes)) + contentTokens := maxTokens - estimateTokens(marker) - 4 + if contentTokens <= 0 { + summary := "[Tool result omitted to fit the model context]" + if estimateTokens(summary) <= maxTokens { + return summary + } + return toolOutputPrefixByTokenBudget(runes, maxTokens) + } + + headUnits := contentTokens * 6 + tailUnits := contentTokens*12 - headUnits + headEnd := toolOutputPrefixRuneCount(runes, headUnits) + tailStart := toolOutputSuffixRuneIndex(runes[headEnd:], tailUnits) + headEnd + omitted := tailStart - headEnd + marker = fmt.Sprintf("\n\n[... %d characters omitted to fit the model context ...]\n\n", omitted) + return strings.TrimSpace(string(runes[:headEnd])) + marker + strings.TrimSpace(string(runes[tailStart:])) +} + +// toolOutputPrefixByTokenBudget 在预算不足以容纳截断标记时保留安全前缀。 +func toolOutputPrefixByTokenBudget(runes []rune, maxTokens int64) string { + return strings.TrimSpace(string(runes[:toolOutputPrefixRuneCount(runes, maxTokens*12)])) +} + +// toolOutputPrefixRuneCount 返回给定估算单位内可保留的前缀长度。 +func toolOutputPrefixRuneCount(runes []rune, maxUnits int64) int { + used := int64(0) + for index, r := range runes { + units := toolOutputRuneTokenUnits(r) + if used+units > maxUnits { + return index + } + used += units + } + return len(runes) +} + +// toolOutputSuffixRuneIndex 返回给定估算单位内可保留的后缀起点。 +func toolOutputSuffixRuneIndex(runes []rune, maxUnits int64) int { + used := int64(0) + for index := len(runes) - 1; index >= 0; index-- { + units := toolOutputRuneTokenUnits(runes[index]) + if used+units > maxUnits { + return index + 1 + } + used += units + } + return 0 +} + +// toolOutputRuneTokenUnits 使用十二分之一 token 表示单个字符的估算权重。 +func toolOutputRuneTokenUnits(r rune) int64 { + if isCJKRune(r) { + return 8 + } + return 3 } func headTailToolOutput(value string, maxChars int) string { diff --git a/backend/internal/application/conversation/service_trace_tool_test.go b/backend/internal/application/conversation/service_trace_tool_test.go index 3c8ab372..532aeef8 100644 --- a/backend/internal/application/conversation/service_trace_tool_test.go +++ b/backend/internal/application/conversation/service_trace_tool_test.go @@ -549,75 +549,78 @@ func TestToolExecutionLedgerNormalizesArguments(t *testing.T) { } } -func TestBudgetToolOutputForModelKeepsSmallResults(t *testing.T) { - raw := `{"content":[{"type":"text","text":"small result"}]}` - if got := budgetToolOutputForModel(model.ToolCall{OutputJSON: raw}, 100, false); got != raw { - t.Fatalf("expected small tool result to stay unchanged, got %q", got) +func TestBudgetToolOutputForModelKeepsLargeReadableResultWhenItFits(t *testing.T) { + subtitle := "HEAD " + strings.Repeat("subtitle content ", 4000) + "TAIL" + raw := `{"content":[{"type":"text","text":"` + subtitle + `"}]}` + prepared := modelToolOutputForModel(raw) + if prepared != raw { + t.Fatalf("expected large readable MCP JSON to remain unchanged, got %d chars", len(prepared)) } -} - -func TestBudgetToolOutputForModelKeepsNormalizedJSONWhenItFits(t *testing.T) { - raw := "{\n \"ok\": true,\n \"items\": [\n 1,\n 2\n ]\n}" - got := budgetToolOutputForModel(model.ToolCall{OutputJSON: raw}, 32, false) - if got != `{"items":[1,2],"ok":true}` { - t.Fatalf("expected normalized JSON to fit without truncation envelope, got %q", got) + if got := budgetToolOutputForModel(prepared, 20_000); got != raw { + t.Fatalf("expected readable result to remain complete within token budget, got %d chars", len(got)) } } -func TestBudgetToolOutputForModelWrapsLargeResults(t *testing.T) { - raw := `{"content":[{"type":"text","text":"` + strings.Repeat("a", 80) + `TAIL"}]}` - got := budgetToolOutputForModel(model.ToolCall{OutputJSON: raw}, 80, false) - if !strings.Contains(got, "truncated_for_model") { - t.Fatalf("expected budgeted result marker, got %q", got) +func TestBudgetToolOutputForModelTruncatesByTokenBudget(t *testing.T) { + raw := "HEAD " + strings.Repeat("subtitle content ", 1000) + "TAIL" + got := budgetToolOutputForModel(raw, 800) + if !strings.Contains(got, "omitted to fit the model context") { + t.Fatalf("expected model-context marker, got %q", got) } - if strings.Contains(got, "server-side tool call record") { - t.Fatalf("did not expect retention note without persistence, got %q", got) + if !strings.Contains(got, "HEAD") || !strings.Contains(got, "TAIL") { + t.Fatalf("expected token budget to preserve head and tail, got %q", got) } - if !strings.Contains(got, "TAIL") { - t.Fatalf("expected budgeted model result to preserve tail context, got %q", got) - } - if !strings.Contains(got, "head_tail") { - t.Fatalf("expected budget metadata to describe head/tail selection, got %q", got) + if tokens := estimateTokens(got); tokens > 800 { + t.Fatalf("expected result within token budget, got %d tokens", tokens) } } -func TestBudgetToolOutputForModelOmitsOpaqueSingleLinePayload(t *testing.T) { - raw := strings.Repeat("A", 4096) - got := budgetToolOutputForModel(model.ToolCall{OutputJSON: raw}, 800, false) - if !strings.Contains(got, "Large opaque tool result omitted") { - t.Fatalf("expected opaque payload notice, got %q", got) - } - if !strings.Contains(got, `"content_type":"opaque"`) { - t.Fatalf("expected opaque content type metadata, got %q", got) +func TestBudgetToolOutputForModelTruncatesCJKWithinTokenBudget(t *testing.T) { + raw := "开头 " + strings.Repeat("字幕内容", 2000) + " 结尾" + got := budgetToolOutputForModel(raw, 500) + if !strings.Contains(got, "开头") || !strings.Contains(got, "结尾") { + t.Fatalf("expected CJK result to preserve head and tail") } - if strings.Count(got, strings.Repeat("A", 512)) > 1 { - t.Fatalf("expected opaque payload to be bounded, got %d chars", len(got)) + if tokens := estimateTokens(got); tokens > 500 { + t.Fatalf("expected CJK result within token budget, got %d tokens", tokens) } } -func TestBudgetToolOutputForModelUsesPersistedReferenceForLargeStoredResult(t *testing.T) { - raw := "HEAD\n" + strings.Repeat("x", toolResultReferenceThresholdChars) + "\nTAIL" - row := model.ToolCall{ - ToolCallID: "call_large", - ToolName: "fetch_large", - RunID: "run_1", - OutputJSON: raw, +func TestModelToolOutputForModelOmitsNestedOpaquePayload(t *testing.T) { + raw := `{"content":[{"type":"text","text":"subtitle"},{"type":"image","data":"` + strings.Repeat("A", 4096) + `"}]}` + got := modelToolOutputForModel(raw) + if !strings.Contains(got, "subtitle") || !strings.Contains(got, "Opaque tool payload omitted") { + t.Fatalf("expected readable text and opaque payload summary, got %q", got) } - got := budgetToolOutputForModel(row, toolResultModelBudgetChars, true) - if !strings.HasPrefix(got, " 1000 { + t.Fatalf("expected aggregate model output within budget, got %d tokens", total) } - if replaced == 0 { - t.Fatalf("expected at least one aggregate replacement, got %#v", slots) + if slots[2].result.OutputJSON != small { + t.Fatalf("expected small result to remain complete, got %q", slots[2].result.OutputJSON) } - if total > toolResultAggregateBudgetChars { - t.Fatalf("expected aggregate model-visible output under budget, got %d", total) + for _, index := range []int{0, 1} { + if !strings.Contains(slots[index].result.OutputJSON, "HEAD") || !strings.Contains(slots[index].result.OutputJSON, "TAIL") { + t.Fatalf("expected large result %d to retain head and tail", index) + } } }