fix: count uploaded file content in context token accounting

PromptTokenText now reflects the actual downstream context cost: the uploaded IGNORE.txt file content plus the neutral live prompt, instead of only the pre-split prompt text.
This commit is contained in:
shern-point
2026-04-30 01:12:35 +08:00
parent 78fdd63470
commit ba80052a26
2 changed files with 16 additions and 26 deletions

View File

@@ -35,11 +35,6 @@ func (s Service) ApplyCurrentInputFile(ctx context.Context, a *auth.RequestAuth,
if strings.TrimSpace(fileText) == "" {
return stdReq, errors.New("current user input file produced empty transcript")
}
originalPromptTokenText := strings.TrimSpace(stdReq.PromptTokenText)
if originalPromptTokenText == "" {
originalPromptTokenText = strings.TrimSpace(stdReq.FinalPrompt)
}
result, err := s.DS.UploadFile(ctx, a, dsclient.UploadFileRequest{
Filename: currentInputFilename,
ContentType: currentInputContentType,
@@ -65,7 +60,9 @@ func (s Service) ApplyCurrentInputFile(ctx context.Context, a *auth.RequestAuth,
stdReq.CurrentInputFileApplied = true
stdReq.RefFileIDs = prependUniqueRefFileID(stdReq.RefFileIDs, fileID)
stdReq.FinalPrompt, stdReq.ToolNames = promptcompat.BuildOpenAIPrompt(messages, stdReq.ToolsRaw, "", stdReq.ToolChoice, stdReq.Thinking)
stdReq.PromptTokenText = originalPromptTokenText
// Token accounting must reflect the actual downstream context:
// the uploaded IGNORE.txt file content + the neutral live prompt.
stdReq.PromptTokenText = fileText + "\n" + stdReq.FinalPrompt
return stdReq, nil
}

View File

@@ -327,20 +327,24 @@ func TestApplyCurrentInputFilePreservesFullContextPromptForTokenCounting(t *test
if err != nil {
t.Fatalf("normalize failed: %v", err)
}
originalPrompt := stdReq.FinalPrompt
out, err := h.applyCurrentInputFile(context.Background(), &auth.RequestAuth{DeepSeekToken: "token"}, stdReq)
if err != nil {
t.Fatalf("apply current input file failed: %v", err)
}
if out.PromptTokenText != originalPrompt {
t.Fatalf("expected prompt token text to preserve original prompt, got %q want %q", out.PromptTokenText, originalPrompt)
}
if out.FinalPrompt == originalPrompt {
if out.FinalPrompt == stdReq.FinalPrompt {
t.Fatalf("expected live prompt to be rewritten after current input file")
}
// PromptTokenText must include the uploaded file content (which contains the full context)
// plus the neutral live prompt — reflecting the actual downstream token cost.
if !strings.Contains(out.PromptTokenText, "first user turn") || !strings.Contains(out.PromptTokenText, "latest user turn") {
t.Fatalf("expected prompt token text to retain full context, got %q", out.PromptTokenText)
t.Fatalf("expected prompt token text to contain file context with full conversation, got %q", out.PromptTokenText)
}
if !strings.Contains(out.PromptTokenText, "[file content end]") || !strings.Contains(out.PromptTokenText, "[file name]: IGNORE") {
t.Fatalf("expected prompt token text to include IGNORE.txt file wrapper, got %q", out.PromptTokenText)
}
if !strings.Contains(out.PromptTokenText, "Answer the latest user request directly.") {
t.Fatalf("expected prompt token text to also include neutral live prompt, got %q", out.PromptTokenText)
}
if strings.Contains(out.FinalPrompt, "first user turn") || strings.Contains(out.FinalPrompt, "latest user turn") {
t.Fatalf("expected live prompt to hide original turns, got %q", out.FinalPrompt)
@@ -487,15 +491,9 @@ func TestChatCompletionsCurrentInputFileUploadsContextAndKeepsNeutralPrompt(t *t
}
usage, _ := body["usage"].(map[string]any)
promptTokens := int(usage["prompt_tokens"].(float64))
originalPrompt, _ := promptcompat.BuildOpenAIPrompt(historySplitTestMessages(), nil, "", defaultToolChoicePolicy(), true)
expectedMin := util.CountPromptTokens(originalPrompt, "deepseek-v4-flash")
neutralPrompt := promptText
neutralCount := util.CountPromptTokens(neutralPrompt, "deepseek-v4-flash")
if promptTokens != expectedMin {
t.Fatalf("expected prompt_tokens from full original context, got=%d want=%d", promptTokens, expectedMin)
}
neutralCount := util.CountPromptTokens(promptText, "deepseek-v4-flash")
if promptTokens <= neutralCount {
t.Fatalf("expected prompt_tokens to exceed neutral live prompt count, got=%d neutral=%d", promptTokens, neutralCount)
t.Fatalf("expected prompt_tokens to exceed neutral live prompt count (includes file context), got=%d neutral=%d", promptTokens, neutralCount)
}
}
@@ -545,14 +543,9 @@ func TestResponsesCurrentInputFileUploadsContextAndKeepsNeutralPrompt(t *testing
}
usage, _ := body["usage"].(map[string]any)
inputTokens := int(usage["input_tokens"].(float64))
originalPrompt, _ := promptcompat.BuildOpenAIPrompt(historySplitTestMessages(), nil, "", defaultToolChoicePolicy(), true)
expectedMin := util.CountPromptTokens(originalPrompt, "deepseek-v4-flash")
neutralCount := util.CountPromptTokens(promptText, "deepseek-v4-flash")
if inputTokens != expectedMin {
t.Fatalf("expected input_tokens from full original context, got=%d want=%d", inputTokens, expectedMin)
}
if inputTokens <= neutralCount {
t.Fatalf("expected input_tokens to exceed neutral live prompt count, got=%d neutral=%d", inputTokens, neutralCount)
t.Fatalf("expected input_tokens to exceed neutral live prompt count (includes file context), got=%d neutral=%d", inputTokens, neutralCount)
}
}