diff --git a/docs/guides/configuration.md b/docs/guides/configuration.md index c1b08047b..efdea0e3f 100644 --- a/docs/guides/configuration.md +++ b/docs/guides/configuration.md @@ -269,6 +269,7 @@ chat_endpoint_path: /chat/completions discovery_endpoint_path: /models generate_max_retries: 5 generate_idle_timeout_sec: 300 +prompt_cache: auto ``` 新增的生成链路控制字段含义如下： @@ -276,6 +277,23 @@ generate_idle_timeout_sec: 300 - `generate_max_retries`：额外重试次数，不含首次尝试；未填写时默认使用 `5`，显式填写 `0` 表示关闭生成重试，且必须 `<= 20`。 - `generate_start_timeout_sec`：写在 `config.yaml` 顶层，从发请求到收到首个有效流 payload 的最长等待窗口；`<= 0` 时回退默认值 `90`。 - `generate_idle_timeout_sec`：首包后连续没有任何新 payload 的最长空闲窗口；`<= 0` 时回退默认值 `300`。 +- `prompt_cache`：是否发送 prompt 缓存标记，只接受 `auto` 与 `off`。留空时按 driver 取默认值：`anthropic` 默认 `auto`，其余 driver 默认 `off`。关闭后请求内容不变，但不再携带缓存断点与缓存路由键，用于规避缓存写入溢价或排障。 + +### prompt 缓存 + +NeoCode 把系统提示固定切分为"稳定前缀"与"易变尾部"两段，稳定前缀由角色定义与工作区规则文件（如 `AGENTS.md`）组成。各 adapter 按厂商能力映射： + +| Driver | 映射方式 | +|--------|----------| +| `anthropic` | 系统提示拆为两个 text block，在稳定块末尾与最后一个工具上放置 `cache_control: ephemeral`；仅当稳定前缀估算达到 1024 token 时才打标 | +| `openaicompat` | `chat completions` 与 `responses` 两种模式均发送 `prompt_cache_key`（由模型、稳定前缀与工具名集合哈希生成），提升同前缀请求的路由命中率 | +| `gemini` | 依赖隐式缓存，不发送额外参数 | +| `deepseek` / `qwen` / `glm` / `mimo` / `minimax` | 不发送路由键：各自有独立的上下文缓存机制，且对未知字段校验更严格 | + +默认值按 driver 区分。`cache_control` 是 Anthropic 的官方协议字段，端点不会拒绝，因此默认开启。`openaicompat` 服务任意兼容网关，`prompt_cache_key` 属于可选扩展参数，在未确认端点支持前发送可能让原本可用的网关返回 400，因此默认关闭；需要在该 provider 的 `provider.yaml` 中显式写 `prompt_cache: auto` 声明该端点支持。该设置支持写入 provider 文件并随配置读写往返保留。 + +memo 索引、todo、git 状态等随会话变化的内容位于易变尾部，不会进入缓存前缀。usage 统计中的 `cache_read_input_tokens` 与 `prompt_tokens_details.cached_tokens` 会归一化到统一字段，缓存命中部分计入总输入预算，不会导致 compact 判定低估。 + 启动时会自动把缺失的 `generate_start_timeout_sec` 规范化写回 `config.yaml`，避免磁盘配置与运行时默认值不一致。 diff --git a/docs/runtime-provider-event-flow.md b/docs/runtime-provider-event-flow.md index 9d98fcaa2..16c24bafe 100644 --- a/docs/runtime-provider-event-flow.md +++ b/docs/runtime-provider-event-flow.md @@ -120,6 +120,10 @@ provider 返回后，runtime 会执行显式的账本调和。 - `output_tokens` - `output_source` - `has_unknown_usage` +- `cache_read_input_tokens`（可选，prompt 缓存命中读取的输入 token） +- `cache_write_input_tokens`（可选，写入 prompt 缓存的输入 token，仅 Anthropic 显式缓存有该口径） + +`input_tokens` 按“含缓存的总输入”统计（等于未命中输入 + 缓存命中 + 缓存写入），因此开启 prompt 缓存后预算累计与 compact 判定不会低估。 规则： diff --git a/internal/config/prompt_cache_test.go b/internal/config/prompt_cache_test.go new file mode 100644 index 000000000..ad21144fc --- /dev/null +++ b/internal/config/prompt_cache_test.go @@ -0,0 +1,48 @@ +package config + +import ( + "strings" + "testing" + + "neo-code/internal/provider" +) + +func TestValidateOptionalPromptCache(t *testing.T) { + t.Parallel() + + valid := []string{"", "auto", "off", " AUTO ", "Off"} + for _, value := range valid { + if err := validateOptionalPromptCache(value); err != nil { + t.Fatalf("validateOptionalPromptCache(%q) unexpected error: %v", value, err) + } + } + + if err := validateOptionalPromptCache("always"); err == nil { + t.Fatal("expected error for unsupported prompt_cache value") + } else if !strings.Contains(err.Error(), provider.PromptCacheModeAuto) { + t.Fatalf("error should mention accepted values, got %v", err) + } +} + +func TestProviderConfigPromptCacheReachesRuntimeConfig(t *testing.T) { + t.Parallel() + + resolved := ResolvedProviderConfig{ + ProviderConfig: ProviderConfig{ + Name: "cached-provider", + Driver: "openaicompat", + BaseURL: "https://llm.example.com/v1", + Model: "server-default", + APIKeyEnv: "CACHED_PROVIDER_KEY", + PromptCache: "off", + }, + } + + runtimeCfg, err := resolved.ToRuntimeConfig() + if err != nil { + t.Fatalf("ToRuntimeConfig() error = %v", err) + } + if runtimeCfg.PromptCacheMode != provider.PromptCacheModeOff { + t.Fatalf("PromptCacheMode = %q, want %q", runtimeCfg.PromptCacheMode, provider.PromptCacheModeOff) + } +} diff --git a/internal/config/provider.go b/internal/config/provider.go index 494c89947..51f6aad7a 100644 --- a/internal/config/provider.go +++ b/internal/config/provider.go @@ -22,20 +22,22 @@ const ( ) type ProviderConfig struct { - Name string `yaml:"name"` - Driver string `yaml:"driver"` - BaseURL string `yaml:"base_url"` - Model string `yaml:"model"` - APIKeyEnv string `yaml:"api_key_env"` - GenerateMaxRetries int `yaml:"generate_max_retries,omitempty"` - GenerateMaxRetriesSet bool `yaml:"-"` - GenerateIdleTimeoutSec int `yaml:"generate_idle_timeout_sec,omitempty"` - ModelSource string `yaml:"-"` - ChatAPIMode string `yaml:"-"` - ChatEndpointPath string `yaml:"-"` - DiscoveryEndpointPath string `yaml:"-"` - Models []providertypes.ModelDescriptor `yaml:"-"` - Source ProviderSource `yaml:"-"` + Name string `yaml:"name"` + Driver string `yaml:"driver"` + BaseURL string `yaml:"base_url"` + Model string `yaml:"model"` + APIKeyEnv string `yaml:"api_key_env"` + GenerateMaxRetries int `yaml:"generate_max_retries,omitempty"` + GenerateMaxRetriesSet bool `yaml:"-"` + GenerateIdleTimeoutSec int `yaml:"generate_idle_timeout_sec,omitempty"` + // PromptCache 控制是否发送 prompt 缓存标记，取值为 auto 或 off，留空视为 auto。 + PromptCache string `yaml:"prompt_cache,omitempty"` + ModelSource string `yaml:"-"` + ChatAPIMode string `yaml:"-"` + ChatEndpointPath string `yaml:"-"` + DiscoveryEndpointPath string `yaml:"-"` + Models []providertypes.ModelDescriptor `yaml:"-"` + Source ProviderSource `yaml:"-"` } type ResolvedProviderConfig struct { @@ -94,6 +96,9 @@ func (p ProviderConfig) Validate() error { if err := validateOptionalGenerateMaxRetries(p.GenerateMaxRetries); err != nil { return fmt.Errorf("provider %q: %w", p.Name, err) } + if err := validateOptionalPromptCache(p.PromptCache); err != nil { + return fmt.Errorf("provider %q: %w", p.Name, err) + } if err := validateOptionalGenerateDurationSeconds("generate_idle_timeout_sec", p.GenerateIdleTimeoutSec); err != nil { return fmt.Errorf("provider %q: %w", p.Name, err) } @@ -134,6 +139,17 @@ func validateOptionalNonNegativeGenerateControl(field string, value int) error { return nil } +// validateOptionalPromptCache 校验 prompt 缓存开关，只接受 auto 与 off，空值表示使用默认值。 +func validateOptionalPromptCache(value string) error { + normalized := strings.ToLower(strings.TrimSpace(value)) + switch normalized { + case "", provider.PromptCacheModeAuto, provider.PromptCacheModeOff: + return nil + default: + return fmt.Errorf("prompt_cache must be %q or %q", provider.PromptCacheModeAuto, provider.PromptCacheModeOff) + } +} + // validateOptionalGenerateMaxRetries 校验额外重试次数，防止超大值导致生成阶段重试循环过长。 func validateOptionalGenerateMaxRetries(value int) error { if err := validateOptionalNonNegativeGenerateControl("generate_max_retries", value); err != nil { @@ -309,9 +325,11 @@ func (p ResolvedProviderConfig) ToRuntimeConfig() (provider.RuntimeConfig, error ChatAPIMode: chatAPIMode, ChatEndpointPath: chatEndpointPath, DiscoveryEndpointPath: discoveryEndpointPath, - GenerateMaxRetries: p.resolveGenerateMaxRetries(), - GenerateStartTimeout: provider.NormalizeGenerateStartTimeout(time.Duration(p.GenerateStartTimeoutSec) * time.Second), - GenerateIdleTimeout: provider.NormalizeGenerateIdleTimeout(time.Duration(p.GenerateIdleTimeoutSec) * time.Second), + // 保留显式取值（含空值），由 provider 层按 driver 决定默认模式。 + PromptCacheMode: strings.ToLower(strings.TrimSpace(p.PromptCache)), + GenerateMaxRetries: p.resolveGenerateMaxRetries(), + GenerateStartTimeout: provider.NormalizeGenerateStartTimeout(time.Duration(p.GenerateStartTimeoutSec) * time.Second), + GenerateIdleTimeout: provider.NormalizeGenerateIdleTimeout(time.Duration(p.GenerateIdleTimeoutSec) * time.Second), }, nil } diff --git a/internal/config/provider_custom_normalize.go b/internal/config/provider_custom_normalize.go index f18967b7b..33f5f3c4a 100644 --- a/internal/config/provider_custom_normalize.go +++ b/internal/config/provider_custom_normalize.go @@ -21,6 +21,7 @@ func NormalizeCustomProviderInput(input SaveCustomProviderInput) (SaveCustomProv GenerateMaxRetriesSet: input.GenerateMaxRetriesSet || input.GenerateMaxRetries > 0, GenerateIdleTimeoutSec: normalizeOptionalGenerateInt(input.GenerateIdleTimeoutSec), DiscoveryEndpointPath: strings.TrimSpace(input.DiscoveryEndpointPath), + PromptCache: strings.ToLower(strings.TrimSpace(input.PromptCache)), } if err := validateCustomProviderName(normalized.Name); err != nil { @@ -121,6 +122,7 @@ func validateNormalizedCustomProviderInput(input SaveCustomProviderInput) error ChatAPIMode: input.ChatAPIMode, ChatEndpointPath: input.ChatEndpointPath, DiscoveryEndpointPath: input.DiscoveryEndpointPath, + PromptCache: input.PromptCache, Models: input.Models, Source: ProviderSourceCustom, } diff --git a/internal/config/provider_loader.go b/internal/config/provider_loader.go index 628ff574d..dc9fa14b0 100644 --- a/internal/config/provider_loader.go +++ b/internal/config/provider_loader.go @@ -21,17 +21,19 @@ const ( ) type customProviderFile struct { - Name string `yaml:"name"` - Driver string `yaml:"driver"` - APIKeyEnv string `yaml:"api_key_env"` - GenerateMaxRetries *int `yaml:"generate_max_retries,omitempty"` - GenerateIdleTimeoutSec int `yaml:"generate_idle_timeout_sec,omitempty"` - ModelSource string `yaml:"model_source,omitempty"` - ChatAPIMode string `yaml:"chat_api_mode,omitempty"` - BaseURL string `yaml:"base_url,omitempty"` - ChatEndpointPath string `yaml:"chat_endpoint_path,omitempty"` - DiscoveryEndpointPath string `yaml:"discovery_endpoint_path,omitempty"` - Models []customProviderModelFile `yaml:"models,omitempty"` + Name string `yaml:"name"` + Driver string `yaml:"driver"` + APIKeyEnv string `yaml:"api_key_env"` + GenerateMaxRetries *int `yaml:"generate_max_retries,omitempty"` + GenerateIdleTimeoutSec int `yaml:"generate_idle_timeout_sec,omitempty"` + ModelSource string `yaml:"model_source,omitempty"` + ChatAPIMode string `yaml:"chat_api_mode,omitempty"` + BaseURL string `yaml:"base_url,omitempty"` + ChatEndpointPath string `yaml:"chat_endpoint_path,omitempty"` + DiscoveryEndpointPath string `yaml:"discovery_endpoint_path,omitempty"` + // PromptCache 控制是否发送 prompt 缓存标记，取值 auto 或 off。 + PromptCache string `yaml:"prompt_cache,omitempty"` + Models []customProviderModelFile `yaml:"models,omitempty"` } type customProviderModelFile struct { @@ -120,6 +122,7 @@ func loadCustomProvider(providerDir string) (ProviderConfig, error) { ChatAPIMode: strings.TrimSpace(file.ChatAPIMode), ChatEndpointPath: strings.TrimSpace(file.ChatEndpointPath), DiscoveryEndpointPath: strings.TrimSpace(file.DiscoveryEndpointPath), + PromptCache: strings.TrimSpace(file.PromptCache), Models: models, }) if err != nil { @@ -138,6 +141,7 @@ func loadCustomProvider(providerDir string) (ProviderConfig, error) { ChatAPIMode: normalizedInput.ChatAPIMode, ChatEndpointPath: normalizedInput.ChatEndpointPath, DiscoveryEndpointPath: normalizedInput.DiscoveryEndpointPath, + PromptCache: normalizedInput.PromptCache, Models: normalizedInput.Models, Source: ProviderSourceCustom, } @@ -194,6 +198,7 @@ type SaveCustomProviderInput struct { GenerateIdleTimeoutSec int DiscoveryEndpointPath string ModelSource string + PromptCache string Models []providertypes.ModelDescriptor } @@ -217,6 +222,7 @@ func SaveCustomProviderWithModels(baseDir string, input SaveCustomProviderInput) GenerateIdleTimeoutSec: normalizedInput.GenerateIdleTimeoutSec, ModelSource: normalizedInput.ModelSource, ChatAPIMode: normalizedInput.ChatAPIMode, + PromptCache: normalizedInput.PromptCache, } cfg.BaseURL = normalizedInput.BaseURL diff --git a/internal/config/provider_prompt_cache_test.go b/internal/config/provider_prompt_cache_test.go new file mode 100644 index 000000000..a7df617af --- /dev/null +++ b/internal/config/provider_prompt_cache_test.go @@ -0,0 +1,119 @@ +package config + +import ( + "context" + "os" + "path/filepath" + "strings" + "testing" + + "neo-code/internal/provider" +) + +func TestCustomProviderPromptCacheRoundTrip(t *testing.T) { + t.Parallel() + + baseDir := t.TempDir() + const providerName = "cached-gateway" + if err := SaveCustomProviderWithModels(baseDir, SaveCustomProviderInput{ + Name: providerName, + Driver: provider.DriverOpenAICompat, + BaseURL: "https://llm.example.com/v1", + APIKeyEnv: "CACHED_GATEWAY_KEY", + ModelSource: ModelSourceDiscover, + DiscoveryEndpointPath: provider.DiscoveryEndpointPathModels, + PromptCache: "OFF", + }); err != nil { + t.Fatalf("SaveCustomProviderWithModels() error = %v", err) + } + + cfg, err := loadCustomProvider(filepath.Join(baseDir, providersDirName, providerName)) + if err != nil { + t.Fatalf("loadCustomProvider() error = %v", err) + } + if cfg.PromptCache != provider.PromptCacheModeOff { + t.Fatalf("PromptCache = %q, want %q after round trip", cfg.PromptCache, provider.PromptCacheModeOff) + } + + runtimeCfg, err := cfg.Resolve() + if err != nil { + t.Fatalf("Resolve() error = %v", err) + } + providerCfg, err := runtimeCfg.ToRuntimeConfig() + if err != nil { + t.Fatalf("ToRuntimeConfig() error = %v", err) + } + if providerCfg.PromptCacheMode != provider.PromptCacheModeOff { + t.Fatalf("PromptCacheMode = %q, want %q", providerCfg.PromptCacheMode, provider.PromptCacheModeOff) + } +} + +func TestCustomProviderFileAcceptsPromptCacheKey(t *testing.T) { + t.Parallel() + + loader := NewLoader(t.TempDir(), testDefaultConfig()) + customDir := filepath.Join(loader.BaseDir(), providersDirName, "company-gateway") + if err := os.MkdirAll(customDir, 0o755); err != nil { + t.Fatalf("mkdir custom provider dir: %v", err) + } + + // 该键曾被 KnownFields(true) 拒绝：结构体缺少 prompt_cache 字段时，按文档写配置会直接加载失败。 + providerYAML := ` +name: company-gateway +driver: openaicompat +base_url: https://llm.example.com/v1 +api_key_env: COMPANY_GATEWAY_API_KEY +discovery_endpoint_path: /models +prompt_cache: off +` + if err := os.WriteFile(filepath.Join(customDir, customProviderConfigName), []byte(strings.TrimSpace(providerYAML)+"\n"), 0o644); err != nil { + t.Fatalf("write provider.yaml: %v", err) + } + + cfg, err := loader.Load(context.Background()) + if err != nil { + t.Fatalf("Load() error = %v", err) + } + found := false + for _, p := range cfg.Providers { + if p.Name == "company-gateway" { + found = true + if p.PromptCache != provider.PromptCacheModeOff { + t.Fatalf("PromptCache = %q, want off", p.PromptCache) + } + } + } + if !found { + t.Fatal("custom provider company-gateway not loaded") + } +} + +func TestCustomProviderFileRejectsInvalidPromptCache(t *testing.T) { + t.Parallel() + + loader := NewLoader(t.TempDir(), testDefaultConfig()) + customDir := filepath.Join(loader.BaseDir(), providersDirName, "bad-cache-gateway") + if err := os.MkdirAll(customDir, 0o755); err != nil { + t.Fatalf("mkdir custom provider dir: %v", err) + } + + providerYAML := ` +name: bad-cache-gateway +driver: openaicompat +base_url: https://llm.example.com/v1 +api_key_env: BAD_CACHE_GATEWAY_KEY +discovery_endpoint_path: /models +prompt_cache: always +` + if err := os.WriteFile(filepath.Join(customDir, customProviderConfigName), []byte(strings.TrimSpace(providerYAML)+"\n"), 0o644); err != nil { + t.Fatalf("write provider.yaml: %v", err) + } + + _, err := loader.Load(context.Background()) + if err == nil { + t.Fatal("expected load to reject unsupported prompt_cache value") + } + if !strings.Contains(err.Error(), "prompt_cache") { + t.Fatalf("error should mention prompt_cache, got %v", err) + } +} diff --git a/internal/context/builder.go b/internal/context/builder.go index 131e8a226..bb315d98c 100644 --- a/internal/context/builder.go +++ b/internal/context/builder.go @@ -11,24 +11,26 @@ type DefaultBuilder struct { trimPolicy messageTrimPolicy } -// newStablePromptSources 返回稳定提示词来源列表，适合作为缓存前缀。 -// extra 会追加到 stable 中（如 memo 持久记忆索引）。 -func newStablePromptSources(extra ...SectionSource) []promptSectionSource { - sources := []promptSectionSource{ +// newStablePromptSources 返回稳定提示词来源列表，构成 prompt 缓存的可复用前缀。 +// 只有"进程生命周期内不变、且随工作区规则文件变化"的来源才允许进入本列表： +// 任何随会话写入变化的来源一旦进入稳定段，都会让整个缓存前缀失效。 +func newStablePromptSources() []promptSectionSource { + return []promptSectionSource{ corePromptSource{}, newRulesPromptSource(nil), } +} + +// newDynamicPromptSources 返回动态提示词来源列表，随任务进度、会话状态变化。 +// extra 来源（如 memo 索引、会话记忆）排在最前：它们随写入变化，属于动态段而非缓存前缀。 +func newDynamicPromptSources(extra ...SectionSource) []promptSectionSource { + sources := make([]promptSectionSource, 0, len(extra)+7) for _, src := range extra { if src != nil { sources = append(sources, src) } } - return sources -} - -// newDynamicPromptSources 返回动态提示词来源列表，随任务进度、会话状态变化。 -func newDynamicPromptSources() []promptSectionSource { - return []promptSectionSource{ + return append(sources, capabilitiesSource{}, taskStateSource{}, planModeContextSource{}, @@ -36,14 +38,15 @@ func newDynamicPromptSources() []promptSectionSource { skillPromptSource{}, repositoryContextSource{}, &systemStateSource{}, - } + ) } // NewConfiguredBuilder 基于可选 SectionSource 列表构建上下文构建器，是推荐的统一构造入口。 +// 传入的来源统一归入动态段，避免易变内容污染缓存前缀。 func NewConfiguredBuilder(sources ...SectionSource) Builder { return &DefaultBuilder{ - stablePromptSources: newStablePromptSources(sources...), - dynamicPromptSources: newDynamicPromptSources(), + stablePromptSources: newStablePromptSources(), + dynamicPromptSources: newDynamicPromptSources(sources...), trimPolicy: spanMessageTrimPolicy{}, } } diff --git a/internal/context/builder_test.go b/internal/context/builder_test.go index c76fe256f..3ecf44b60 100644 --- a/internal/context/builder_test.go +++ b/internal/context/builder_test.go @@ -909,7 +909,7 @@ func TestDefaultBuilderBuildTodoChangeDoesNotChangeStablePrompt(t *testing.T) { } } -func TestDefaultBuilderBuildMemoIsStable(t *testing.T) { +func TestDefaultBuilderBuildMemoStaysInDynamicSegment(t *testing.T) { t.Parallel() builder := NewConfiguredBuilder(stubPromptSectionSource{ @@ -928,11 +928,13 @@ func TestDefaultBuilderBuildMemoIsStable(t *testing.T) { t.Fatalf("Build() error = %v", err) } - if !strings.Contains(result.StableSystemPrompt, "## memo") { - t.Fatalf("expected memo in StableSystemPrompt, got %q", result.StableSystemPrompt) + // memo 索引随会话写入变化，必须留在动态段：一旦进入稳定段， + // 每次 memo 变更都会让 prompt 缓存前缀整体失效。 + if strings.Contains(result.StableSystemPrompt, "## memo") { + t.Fatalf("did not expect memo in StableSystemPrompt, got %q", result.StableSystemPrompt) } - if strings.Contains(result.DynamicSystemPrompt, "## memo") { - t.Fatalf("did not expect memo in DynamicSystemPrompt, got %q", result.DynamicSystemPrompt) + if !strings.Contains(result.DynamicSystemPrompt, "## memo") { + t.Fatalf("expected memo in DynamicSystemPrompt, got %q", result.DynamicSystemPrompt) } } diff --git a/internal/provider/anthropic/provider.go b/internal/provider/anthropic/provider.go index 8d8092f8e..50b21c6fa 100644 --- a/internal/provider/anthropic/provider.go +++ b/internal/provider/anthropic/provider.go @@ -125,6 +125,12 @@ func (p *Provider) generateOnce( usage.InputTokens = int(variant.Message.Usage.InputTokens) usage.InputObserved = true } + if variant.Message.Usage.JSON.CacheReadInputTokens.Valid() { + usage.CacheReadInputTokens = int(variant.Message.Usage.CacheReadInputTokens) + } + if variant.Message.Usage.JSON.CacheCreationInputTokens.Valid() { + usage.CacheWriteInputTokens = int(variant.Message.Usage.CacheCreationInputTokens) + } if variant.Message.Usage.JSON.OutputTokens.Valid() { usage.OutputTokens = int(variant.Message.Usage.OutputTokens) usage.OutputObserved = true @@ -200,6 +206,12 @@ func (p *Provider) generateOnce( usage.InputTokens = int(variant.Usage.InputTokens) usage.InputObserved = true } + if variant.Usage.JSON.CacheReadInputTokens.Valid() { + usage.CacheReadInputTokens = int(variant.Usage.CacheReadInputTokens) + } + if variant.Usage.JSON.CacheCreationInputTokens.Valid() { + usage.CacheWriteInputTokens = int(variant.Usage.CacheCreationInputTokens) + } } } if streamErr := streamReader.Err(); streamErr != nil { @@ -223,7 +235,8 @@ func (p *Provider) generateOnce( } } if usage.TotalTokens <= 0 && (usage.InputObserved || usage.OutputObserved) { - usage.TotalTokens = usage.InputTokens + usage.OutputTokens + // 缓存命中与写入部分不计入 InputTokens，总量必须按 TotalInputTokens 汇总。 + usage.TotalTokens = usage.TotalInputTokens() + usage.OutputTokens } if !usage.InputObserved && !usage.OutputObserved { return provider.EmitMessageDone(ctx, events, finishReason, nil) diff --git a/internal/provider/anthropic/request.go b/internal/provider/anthropic/request.go index 61263f444..c7397d473 100644 --- a/internal/provider/anthropic/request.go +++ b/internal/provider/anthropic/request.go @@ -17,6 +17,14 @@ import ( const defaultMaxTokens = 4096 +const ( + // anthropicMinimumCacheableTokens 是 Anthropic 显式缓存的最小可缓存长度。 + // 低于该阈值时打标不会产生命中，却仍需支付缓存写入溢价，因此必须前置拦截。 + anthropicMinimumCacheableTokens = 1024 + // anthropicMaxCacheBreakpoints 是单次请求允许的最大缓存断点数。 + anthropicMaxCacheBreakpoints = 4 +) + // BuildRequest 将通用 GenerateRequest 直接转换为 Anthropic SDK 入参。 func BuildRequest(ctx context.Context, cfg provider.RuntimeConfig, req providertypes.GenerateRequest) (anthropic.MessageNewParams, error) { model := strings.TrimSpace(req.Model) @@ -27,13 +35,18 @@ func BuildRequest(ctx context.Context, cfg provider.RuntimeConfig, req providert return anthropic.MessageNewParams{}, errors.New(errorPrefix + "model is empty") } + // adapter 自报 driver：anthropic 的 cache_control 是官方字段，默认开启。 + cacheEnabled := provider.PromptCacheEnabled(cfg.PromptCacheMode, provider.DriverAnthropic) + breakpoints := 0 + params := anthropic.MessageNewParams{ MaxTokens: int64(defaultMaxTokens), Model: anthropic.Model(model), Messages: make([]anthropic.MessageParam, 0, len(req.Messages)), } - if strings.TrimSpace(req.SystemPrompt) != "" { - params.System = []anthropic.TextBlockParam{{Text: req.SystemPrompt}} + if blocks, used := buildAnthropicSystemBlocks(req, cacheEnabled); len(blocks) > 0 { + params.System = blocks + breakpoints += used } assetPolicy := session.NormalizeAssetPolicy(cfg.SessionAssetPolicy) @@ -72,11 +85,67 @@ func BuildRequest(ctx context.Context, cfg provider.RuntimeConfig, req providert } params.Tools = append(params.Tools, converted) } + // 把断点打在最后一个工具上，使工具定义独立于系统提示被缓存： + // 系统提示因规则文件变更而失效时，工具前缀仍可命中。 + if cacheEnabled && breakpoints < anthropicMaxCacheBreakpoints && toolSpecsCacheable(req.Tools) { + if last := ¶ms.Tools[len(params.Tools)-1]; last.OfTool != nil { + last.OfTool.CacheControl = anthropic.NewCacheControlEphemeralParam() + breakpoints++ + } + } } return params, nil } +// buildAnthropicSystemBlocks 把系统提示拆成稳定前缀与易变尾部两个 text block， +// 并在稳定块末尾放置缓存断点。两块的文本拼接后严格等于 SystemPrompt。 +// 返回值 used 表示本次实际放置的断点数（0 或 1）。 +func buildAnthropicSystemBlocks(req providertypes.GenerateRequest, cacheEnabled bool) ([]anthropic.TextBlockParam, int) { + if strings.TrimSpace(req.SystemPrompt) == "" { + return nil, 0 + } + + prefix := req.ResolvedStablePrefix() + if prefix == "" { + return []anthropic.TextBlockParam{{Text: req.SystemPrompt}}, 0 + } + suffix := req.VolatileSystemSuffix() + + // 仅为达到最小可缓存长度且显式启用时才打标，避免支付写入溢价却拿不到命中。 + cacheable := cacheEnabled && provider.EstimateTextTokens(prefix) >= anthropicMinimumCacheableTokens + + if !cacheable { + blocks := []anthropic.TextBlockParam{{Text: prefix}} + if suffix != "" { + blocks = append(blocks, anthropic.TextBlockParam{Text: suffix}) + } + return blocks, 0 + } + + stableBlock := anthropic.TextBlockParam{ + Text: prefix, + CacheControl: anthropic.NewCacheControlEphemeralParam(), + } + if suffix == "" { + return []anthropic.TextBlockParam{stableBlock}, 1 + } + return []anthropic.TextBlockParam{stableBlock, {Text: suffix}}, 1 +} + +// toolSpecsCacheable 估算工具定义的规模是否达到最小可缓存长度。 +// 仅做粗粒度字节累加，不做序列化，避免每次请求都构造完整 JSON 文本。 +func toolSpecsCacheable(specs []providertypes.ToolSpec) bool { + size := 0 + for _, spec := range specs { + size += len(spec.Name) + len(spec.Description) + for key, value := range spec.Schema { + size += len(key) + len(fmt.Sprint(value)) + } + } + return provider.EstimateByteCountTokens(size) >= anthropicMinimumCacheableTokens +} + // toAnthropicMessageWithBudget 将通用消息映射为 Anthropic SDK MessageParam，并记录 session_asset 消耗字节数。 func toAnthropicMessageWithBudget( ctx context.Context, diff --git a/internal/provider/anthropic/request_cache_test.go b/internal/provider/anthropic/request_cache_test.go new file mode 100644 index 000000000..8cd5cf359 --- /dev/null +++ b/internal/provider/anthropic/request_cache_test.go @@ -0,0 +1,160 @@ +package anthropic + +import ( + "strings" + "testing" + + "neo-code/internal/provider" + providertypes "neo-code/internal/provider/types" +) + +// longStablePrefix 生成一段足以达到最小可缓存长度的稳定前缀。 +func longStablePrefix() string { + return strings.Repeat("stable-system-prompt-line\n", 200) +} + +func TestBuildAnthropicSystemBlocks(t *testing.T) { + t.Parallel() + + t.Run("空系统提示返回 nil", func(t *testing.T) { + t.Parallel() + blocks, used := buildAnthropicSystemBlocks(providertypes.GenerateRequest{}, true) + if blocks != nil || used != 0 { + t.Fatalf("blocks = %v, used = %d, want nil and 0", blocks, used) + } + }) + + t.Run("无稳定前缀时单块且不打标", func(t *testing.T) { + t.Parallel() + req := providertypes.GenerateRequest{SystemPrompt: "whole prompt"} + blocks, used := buildAnthropicSystemBlocks(req, true) + if len(blocks) != 1 || blocks[0].Text != "whole prompt" { + t.Fatalf("blocks = %v, want single block", blocks) + } + if used != 0 || blocks[0].CacheControl.Type != "" { + t.Fatalf("expected no cache breakpoint, got used=%d control=%v", used, blocks[0].CacheControl) + } + }) + + t.Run("前缀过短时不打标", func(t *testing.T) { + t.Parallel() + req := providertypes.GenerateRequest{ + SystemPrompt: "short" + "\n\n" + "tail", + SystemStablePrefix: "short", + } + blocks, used := buildAnthropicSystemBlocks(req, true) + if used != 0 { + t.Fatalf("used = %d, want 0 for short prefix", used) + } + if len(blocks) != 2 || blocks[0].Text != "short" || blocks[1].Text != "\n\ntail" { + t.Fatalf("blocks = %v, want split without cache control", blocks) + } + }) + + t.Run("够长的前缀打标并保留尾部", func(t *testing.T) { + t.Parallel() + prefix := longStablePrefix() + req := providertypes.GenerateRequest{ + SystemPrompt: prefix + "\n\ndynamic", + SystemStablePrefix: prefix, + } + blocks, used := buildAnthropicSystemBlocks(req, true) + if used != 1 { + t.Fatalf("used = %d, want 1", used) + } + if len(blocks) != 2 { + t.Fatalf("len(blocks) = %d, want 2", len(blocks)) + } + if blocks[0].CacheControl.Type != "ephemeral" { + t.Fatalf("first block cache control = %q, want ephemeral", blocks[0].CacheControl.Type) + } + if blocks[1].CacheControl.Type != "" { + t.Fatal("tail block must not carry cache control") + } + if got := blocks[0].Text + blocks[1].Text; got != req.SystemPrompt { + t.Fatal("split blocks must reconstruct the original system prompt") + } + }) + + t.Run("整体稳定时单块打标", func(t *testing.T) { + t.Parallel() + prefix := longStablePrefix() + req := providertypes.GenerateRequest{SystemPrompt: prefix, SystemStablePrefix: prefix} + blocks, used := buildAnthropicSystemBlocks(req, true) + if used != 1 || len(blocks) != 1 || blocks[0].CacheControl.Type != "ephemeral" { + t.Fatalf("blocks = %v, used = %d, want single cached block", blocks, used) + } + }) + + t.Run("禁用缓存时不打标但仍分段", func(t *testing.T) { + t.Parallel() + prefix := longStablePrefix() + req := providertypes.GenerateRequest{ + SystemPrompt: prefix + "\n\ndynamic", + SystemStablePrefix: prefix, + } + blocks, used := buildAnthropicSystemBlocks(req, false) + if used != 0 { + t.Fatalf("used = %d, want 0 when cache disabled", used) + } + if len(blocks) != 2 || blocks[0].CacheControl.Type != "" { + t.Fatalf("blocks = %v, want split without cache control", blocks) + } + }) +} + +func TestBuildRequestAppliesPromptCache(t *testing.T) { + t.Parallel() + + prefix := longStablePrefix() + req := providertypes.GenerateRequest{ + Model: "claude-sonnet-4", + SystemPrompt: prefix + "\n\ndynamic", + SystemStablePrefix: prefix, + Tools: []providertypes.ToolSpec{ + {Name: "read_file", Description: strings.Repeat("read a file ", 300)}, + {Name: "grep", Description: strings.Repeat("search text ", 300)}, + }, + } + + params, err := BuildRequest(t.Context(), provider.RuntimeConfig{DefaultModel: "claude-sonnet-4"}, req) + if err != nil { + t.Fatalf("BuildRequest() error = %v", err) + } + if len(params.System) != 2 { + t.Fatalf("len(System) = %d, want 2", len(params.System)) + } + if params.System[0].CacheControl.Type != "ephemeral" { + t.Fatal("system stable block should carry cache control") + } + if len(params.Tools) != 2 { + t.Fatalf("len(Tools) = %d, want 2", len(params.Tools)) + } + lastTool := params.Tools[len(params.Tools)-1] + if lastTool.OfTool == nil || lastTool.OfTool.CacheControl.Type != "ephemeral" { + t.Fatal("last tool should carry a cache breakpoint") + } +} + +func TestBuildRequestCacheDisabled(t *testing.T) { + t.Parallel() + + prefix := longStablePrefix() + req := providertypes.GenerateRequest{ + Model: "claude-sonnet-4", + SystemPrompt: prefix + "\n\ndynamic", + SystemStablePrefix: prefix, + } + params, err := BuildRequest(t.Context(), provider.RuntimeConfig{ + DefaultModel: "claude-sonnet-4", + PromptCacheMode: provider.PromptCacheModeOff, + }, req) + if err != nil { + t.Fatalf("BuildRequest() error = %v", err) + } + for i, block := range params.System { + if block.CacheControl.Type != "" { + t.Fatalf("system block %d must not carry cache control when disabled", i) + } + } +} diff --git a/internal/provider/contracts.go b/internal/provider/contracts.go index 32b8cc37b..9353708ff 100644 --- a/internal/provider/contracts.go +++ b/internal/provider/contracts.go @@ -28,9 +28,11 @@ type RuntimeConfig struct { ChatAPIMode string ChatEndpointPath string DiscoveryEndpointPath string - GenerateMaxRetries int - GenerateStartTimeout time.Duration - GenerateIdleTimeout time.Duration + // PromptCacheMode 控制是否发送 prompt 缓存标记，取值参见 NormalizePromptCacheMode。 + PromptCacheMode string + GenerateMaxRetries int + GenerateStartTimeout time.Duration + GenerateIdleTimeout time.Duration } // ResolveAPIKeyValue 在 provider 即将发起请求前解析当前配置引用的 API Key。 diff --git a/internal/provider/estimate.go b/internal/provider/estimate.go index e94b59063..c818e793c 100644 --- a/internal/provider/estimate.go +++ b/internal/provider/estimate.go @@ -35,7 +35,16 @@ func EstimateTextTokens(text string) int { if text == "" { return 0 } - return int(math.Ceil(float64(len([]byte(text))) / 4.0 * localEstimateSlack)) + return EstimateByteCountTokens(len([]byte(text))) +} + +// EstimateByteCountTokens 按字节数做保守放大的本地 token 估算。 +// 供不必构造完整文本、只需按规模做阈值判断的场景复用。 +func EstimateByteCountTokens(byteCount int) int { + if byteCount <= 0 { + return 0 + } + return int(math.Ceil(float64(byteCount) / 4.0 * localEstimateSlack)) } // RequestContainsImagePart 判断请求中是否包含图片分片，供 provider 选择多模态投影估算路径。 diff --git a/internal/provider/openaicompat/chatcompletions/adapter.go b/internal/provider/openaicompat/chatcompletions/adapter.go index 986eb1a4a..0fb7141e8 100644 --- a/internal/provider/openaicompat/chatcompletions/adapter.go +++ b/internal/provider/openaicompat/chatcompletions/adapter.go @@ -100,9 +100,12 @@ type streamChunk struct { } type streamUsage struct { - PromptTokens int `json:"prompt_tokens"` - CompletionTokens int `json:"completion_tokens"` - TotalTokens int `json:"total_tokens"` + PromptTokens int `json:"prompt_tokens"` + CompletionTokens int `json:"completion_tokens"` + TotalTokens int `json:"total_tokens"` + PromptTokensDetails struct { + CachedTokens int `json:"cached_tokens"` + } `json:"prompt_tokens_details"` } // ConsumeStream 消费 chat/completions SSE 流并发出统一流式事件，兼容弱 SSE 事件分隔格式。 @@ -258,26 +261,53 @@ func extractLegacyStreamUsage(usage *providertypes.Usage, raw *streamUsage) { if raw == nil { return } + uncached, cached := normalizeCachedPromptTokens(raw.PromptTokens, raw.PromptTokensDetails.CachedTokens) *usage = providertypes.Usage{ - InputTokens: raw.PromptTokens, - OutputTokens: raw.CompletionTokens, - TotalTokens: raw.TotalTokens, - InputObserved: true, - OutputObserved: true, + InputTokens: uncached, + OutputTokens: raw.CompletionTokens, + TotalTokens: raw.TotalTokens, + InputObserved: true, + OutputObserved: true, + CacheReadInputTokens: cached, } } // extractStreamUsage 将 OpenAI usage 覆盖到统一 token 统计。 +// OpenAI 的 prompt_tokens 已包含缓存命中部分，而 Anthropic 的 input_tokens 不含， +// 因此这里拆出未命中部分，把两家的 InputTokens 统一为"未命中缓存的输入"语义。 func extractStreamUsage(usage *providertypes.Usage, raw openai.CompletionUsage) { + uncached, cached := normalizeCachedPromptTokens(int(raw.PromptTokens), int(raw.PromptTokensDetails.CachedTokens)) *usage = providertypes.Usage{ - InputTokens: int(raw.PromptTokens), - OutputTokens: int(raw.CompletionTokens), - TotalTokens: int(raw.TotalTokens), - InputObserved: true, - OutputObserved: true, + InputTokens: uncached, + OutputTokens: int(raw.CompletionTokens), + TotalTokens: int(raw.TotalTokens), + InputObserved: true, + OutputObserved: true, + CacheReadInputTokens: cached, } } +// normalizeCachedPromptTokens 把厂商上报的缓存命中数收敛到合法区间，并返回拆出后的未命中输入。 +// cached 必须落在 [0, prompt]：若只对减法结果做下限裁剪而保留超额 cached， +// TotalInputTokens() 会高于厂商上报的总量（prompt=100、cached=150 时会得到 150）， +// 而该值会被账本当作已观测输入持久化。 +func normalizeCachedPromptTokens(promptTokens int, cachedTokens int) (uncached int, cached int) { + prompt := clampNonNegative(promptTokens) + cached = clampNonNegative(cachedTokens) + if cached > prompt { + cached = prompt + } + return prompt - cached, cached +} + +// clampNonNegative 把可能因厂商数据异常而出现的负值收敛为 0。 +func clampNonNegative(value int) int { + if value < 0 { + return 0 + } + return value +} + // doneUsagePtr 在 message_done 事件中按 usage 观测状态返回 payload，未观测时返回 nil。 func doneUsagePtr(usage providertypes.Usage) *providertypes.Usage { if !usage.InputObserved && !usage.OutputObserved { diff --git a/internal/provider/openaicompat/chatcompletions/prompt_cache_test.go b/internal/provider/openaicompat/chatcompletions/prompt_cache_test.go new file mode 100644 index 000000000..0f9a5b8b9 --- /dev/null +++ b/internal/provider/openaicompat/chatcompletions/prompt_cache_test.go @@ -0,0 +1,90 @@ +package chatcompletions + +import ( + "testing" + + "neo-code/internal/provider" + providertypes "neo-code/internal/provider/types" +) + +func TestComputePromptCacheKey(t *testing.T) { + t.Parallel() + + base := providertypes.GenerateRequest{ + SystemPrompt: "stable\n\ndynamic", + SystemStablePrefix: "stable", + Tools: []providertypes.ToolSpec{ + {Name: "read_file"}, + {Name: "grep"}, + }, + } + + key := computePromptCacheKey("gpt-5", base) + if len(key) != promptCacheKeyHexLength { + t.Fatalf("len(key) = %d, want %d", len(key), promptCacheKeyHexLength) + } + if got := computePromptCacheKey("gpt-5", base); got != key { + t.Fatal("same input must produce a stable key") + } + + if got := computePromptCacheKey("gpt-5-mini", base); got == key { + t.Fatal("model change must change the key") + } + + withExtraTool := base + withExtraTool.Tools = append(append([]providertypes.ToolSpec{}, base.Tools...), providertypes.ToolSpec{Name: "bash"}) + if got := computePromptCacheKey("gpt-5", withExtraTool); got == key { + t.Fatal("tool set change must change the key") + } + + noPrefix := base + noPrefix.SystemStablePrefix = "" + if got := computePromptCacheKey("gpt-5", noPrefix); got != "" { + t.Fatalf("key without stable prefix = %q, want empty", got) + } +} + +func TestBuildRequestGatesPromptCacheKeyByDriver(t *testing.T) { + t.Parallel() + + req := providertypes.GenerateRequest{ + Model: "gpt-5", + SystemPrompt: "stable\n\ndynamic", + SystemStablePrefix: "stable", + } + + cases := []struct { + name string + driver string + mode string + want bool + }{ + // 通用驱动服务任意网关，未显式声明支持前默认不发送，避免端点 400 回归。 + {name: "通用兼容驱动默认不发送", driver: provider.DriverOpenAICompat, mode: "", want: false}, + {name: "显式 auto 后发送", driver: provider.DriverOpenAICompat, mode: provider.PromptCacheModeAuto, want: true}, + {name: "显式 off 不发送", driver: provider.DriverOpenAICompat, mode: provider.PromptCacheModeOff, want: false}, + // 自研实现是硬约束，显式 auto 也不发送。 + {name: "DeepSeek 显式 auto 也不发送", driver: provider.DriverDeepSeek, mode: provider.PromptCacheModeAuto, want: false}, + {name: "Qwen 显式 auto 也不发送", driver: provider.DriverQwen, mode: provider.PromptCacheModeAuto, want: false}, + } + + for _, tc := range cases { + tc := tc + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + payload, err := BuildRequest(t.Context(), provider.RuntimeConfig{ + Driver: tc.driver, + PromptCacheMode: tc.mode, + }, req) + if err != nil { + t.Fatalf("BuildRequest() error = %v", err) + } + if tc.want && payload.PromptCacheKey == "" { + t.Fatal("expected prompt_cache_key to be set") + } + if !tc.want && payload.PromptCacheKey != "" { + t.Fatalf("expected prompt_cache_key to stay empty, got %q", payload.PromptCacheKey) + } + }) + } +} diff --git a/internal/provider/openaicompat/chatcompletions/request.go b/internal/provider/openaicompat/chatcompletions/request.go index 6e07e8915..293c21f39 100644 --- a/internal/provider/openaicompat/chatcompletions/request.go +++ b/internal/provider/openaicompat/chatcompletions/request.go @@ -15,12 +15,24 @@ import ( "neo-code/internal/session" ) +// errorPrefix 是 chat/completions 适配层统一的错误前缀。 const errorPrefix = "openaicompat provider: " +// maxSessionAssetsTotalBytes 是单次请求允许携带的会话附件总字节上限。 const maxSessionAssetsTotalBytes = provider.MaxSessionAssetsTotalBytes +// htmlErrorSnippetMaxRunes 是错误响应中截取 HTML 片段的字符上限。 const htmlErrorSnippetMaxRunes = 320 +// promptCacheKeyHexLength 与 provider.ComputePromptCacheKey 的输出长度保持一致。 +const promptCacheKeyHexLength = provider.PromptCacheKeyHexLength + +// computePromptCacheKey 委托 provider 层统一实现，保证 chat completions 与 +// responses 两种模式下同一前缀产生相同路由键。 +func computePromptCacheKey(model string, req providertypes.GenerateRequest) string { + return provider.ComputePromptCacheKey(model, req.ResolvedStablePrefix(), provider.PromptCacheToolNames(req.Tools)) +} + // BuildRequest 将 provider.GenerateRequest 转换为 Chat Completions 请求结构。 // 模型优先取 req.Model，其次使用配置中的默认模型。 func BuildRequest(ctx context.Context, cfg provider.RuntimeConfig, req providertypes.GenerateRequest) (Request, error) { @@ -47,6 +59,14 @@ func BuildRequest(ctx context.Context, cfg provider.RuntimeConfig, req providert }) } + // OpenAI 的前缀缓存是自动的，但需要稳定前缀 + 稳定的路由键才能获得高命中率。 + // 该驱动服务任意兼容网关，prompt_cache_key 属可选扩展参数，因此默认不发送， + // 需在 provider 配置中显式写 prompt_cache: auto 声明端点支持； + // 自研 HTTP 厂商（DeepSeek / Qwen / GLM 等）由硬约束排除。 + if provider.ShouldSendPromptCacheKey(cfg.PromptCacheMode, cfg.Driver) { + payload.PromptCacheKey = computePromptCacheKey(model, req) + } + var usedSessionAssetBytes int64 for _, message := range req.Messages { remainingSessionAssetBytes := requestBudget.MaxSessionAssetsTotalBytes - usedSessionAssetBytes diff --git a/internal/provider/openaicompat/chatcompletions/types.go b/internal/provider/openaicompat/chatcompletions/types.go index 18ce33d5c..25ee850ad 100644 --- a/internal/provider/openaicompat/chatcompletions/types.go +++ b/internal/provider/openaicompat/chatcompletions/types.go @@ -11,6 +11,9 @@ type Request struct { ToolChoice string `json:"tool_choice,omitempty"` Stream bool `json:"stream"` ReasoningEffort string `json:"reasoning_effort,omitempty"` + // PromptCacheKey 是缓存路由键，作用是让相同前缀的请求尽量落到同一后端， + // 提升 OpenAI 自动前缀缓存的命中率；为空时不发送该字段。 + PromptCacheKey string `json:"prompt_cache_key,omitempty"` } // Message 表示 OpenAI 协议中的消息格式。 diff --git a/internal/provider/openaicompat/chatcompletions/usage_cache_test.go b/internal/provider/openaicompat/chatcompletions/usage_cache_test.go new file mode 100644 index 000000000..e25b67dca --- /dev/null +++ b/internal/provider/openaicompat/chatcompletions/usage_cache_test.go @@ -0,0 +1,123 @@ +package chatcompletions + +import ( + "testing" + + "github.com/openai/openai-go/v3" + + providertypes "neo-code/internal/provider/types" +) + +func TestExtractStreamUsageNormalizesCachedTokens(t *testing.T) { + t.Parallel() + + raw := openai.CompletionUsage{ + PromptTokens: 1000, + CompletionTokens: 120, + TotalTokens: 1120, + } + raw.PromptTokensDetails.CachedTokens = 600 + + var usage providertypes.Usage + extractStreamUsage(&usage, raw) + + if usage.InputTokens != 400 { + t.Fatalf("InputTokens = %d, want 400 (prompt tokens minus cached)", usage.InputTokens) + } + if usage.CacheReadInputTokens != 600 { + t.Fatalf("CacheReadInputTokens = %d, want 600", usage.CacheReadInputTokens) + } + if usage.TotalInputTokens() != 1000 { + t.Fatalf("TotalInputTokens() = %d, want 1000 without double counting", usage.TotalInputTokens()) + } + if !usage.InputObserved || !usage.OutputObserved { + t.Fatal("observed flags must stay true") + } +} + +func TestExtractLegacyStreamUsageNormalizesCachedTokens(t *testing.T) { + t.Parallel() + + raw := &streamUsage{PromptTokens: 500, CompletionTokens: 40, TotalTokens: 540} + raw.PromptTokensDetails.CachedTokens = 500 + + var usage providertypes.Usage + extractLegacyStreamUsage(&usage, raw) + + if usage.InputTokens != 0 { + t.Fatalf("InputTokens = %d, want 0 when everything is cached", usage.InputTokens) + } + if usage.CacheReadInputTokens != 500 { + t.Fatalf("CacheReadInputTokens = %d, want 500", usage.CacheReadInputTokens) + } + if usage.TotalInputTokens() != 500 { + t.Fatalf("TotalInputTokens() = %d, want 500", usage.TotalInputTokens()) + } +} + +func TestExtractStreamUsageBoundsExcessiveCachedTokens(t *testing.T) { + t.Parallel() + + raw := openai.CompletionUsage{PromptTokens: 100, CompletionTokens: 10, TotalTokens: 110} + // 厂商返回异常数据时 cached 可能大于 prompt。此时 cached 需被收敛到 prompt， + // 否则 TotalInputTokens() 会高于厂商上报的总量。 + raw.PromptTokensDetails.CachedTokens = 150 + + var usage providertypes.Usage + extractStreamUsage(&usage, raw) + + if usage.InputTokens != 0 { + t.Fatalf("InputTokens = %d, want 0 after clamping", usage.InputTokens) + } + if usage.CacheReadInputTokens != 100 { + t.Fatalf("CacheReadInputTokens = %d, want 100 (bounded by prompt tokens)", usage.CacheReadInputTokens) + } + if usage.TotalInputTokens() != 100 { + t.Fatalf("TotalInputTokens() = %d, want 100 (must not exceed reported prompt tokens)", usage.TotalInputTokens()) + } +} + +func TestExtractLegacyStreamUsageBoundsExcessiveCachedTokens(t *testing.T) { + t.Parallel() + + raw := &streamUsage{PromptTokens: 40, CompletionTokens: 5, TotalTokens: 45} + raw.PromptTokensDetails.CachedTokens = 90 + + var usage providertypes.Usage + extractLegacyStreamUsage(&usage, raw) + + if usage.CacheReadInputTokens != 40 { + t.Fatalf("CacheReadInputTokens = %d, want 40", usage.CacheReadInputTokens) + } + if usage.TotalInputTokens() != 40 { + t.Fatalf("TotalInputTokens() = %d, want 40", usage.TotalInputTokens()) + } +} + +func TestNormalizeCachedPromptTokens(t *testing.T) { + t.Parallel() + + cases := []struct { + prompt int + cached int + wantUncached int + wantCached int + }{ + {prompt: 1000, cached: 600, wantUncached: 400, wantCached: 600}, + {prompt: 1000, cached: 0, wantUncached: 1000, wantCached: 0}, + {prompt: 1000, cached: 1000, wantUncached: 0, wantCached: 1000}, + {prompt: 100, cached: 150, wantUncached: 0, wantCached: 100}, + {prompt: -5, cached: 10, wantUncached: 0, wantCached: 0}, + {prompt: 10, cached: -5, wantUncached: 10, wantCached: 0}, + } + for _, tc := range cases { + uncached, cached := normalizeCachedPromptTokens(tc.prompt, tc.cached) + if uncached != tc.wantUncached || cached != tc.wantCached { + t.Fatalf("normalizeCachedPromptTokens(%d, %d) = (%d, %d), want (%d, %d)", + tc.prompt, tc.cached, uncached, cached, tc.wantUncached, tc.wantCached) + } + if uncached+cached > max(tc.prompt, 0) { + t.Fatalf("sum %d exceeds clamped prompt %d", uncached+cached, tc.prompt) + } + } +} diff --git a/internal/provider/openaicompat/generate_sdk.go b/internal/provider/openaicompat/generate_sdk.go index 3bd0a6712..5073bd3d6 100644 --- a/internal/provider/openaicompat/generate_sdk.go +++ b/internal/provider/openaicompat/generate_sdk.go @@ -102,6 +102,9 @@ func convertToChatCompletionParams(req chatcompletions.Request) openai.ChatCompl IncludeUsage: openai.Bool(true), }, } + if key := strings.TrimSpace(req.PromptCacheKey); key != "" { + params.PromptCacheKey = openai.String(key) + } messages := make([]openai.ChatCompletionMessageParamUnion, 0, len(req.Messages)) for _, msg := range req.Messages { diff --git a/internal/provider/openaicompat/responses/prompt_cache_test.go b/internal/provider/openaicompat/responses/prompt_cache_test.go new file mode 100644 index 000000000..fed88d4c1 --- /dev/null +++ b/internal/provider/openaicompat/responses/prompt_cache_test.go @@ -0,0 +1,75 @@ +package responses + +import ( + "testing" + + "neo-code/internal/provider" + providertypes "neo-code/internal/provider/types" +) + +func TestBuildRequestPromptCacheKey(t *testing.T) { + t.Parallel() + + req := providertypes.GenerateRequest{ + Model: "gpt-5", + SystemPrompt: "stable\n\ndynamic", + SystemStablePrefix: "stable", + Tools: []providertypes.ToolSpec{{Name: "read_file"}, {Name: "grep"}}, + } + + t.Run("默认不发送", func(t *testing.T) { + t.Parallel() + payload, err := BuildRequest(t.Context(), provider.RuntimeConfig{Driver: provider.DriverOpenAICompat}, req) + if err != nil { + t.Fatalf("BuildRequest() error = %v", err) + } + if payload.PromptCacheKey != "" { + t.Fatalf("PromptCacheKey = %q, want empty by default", payload.PromptCacheKey) + } + }) + + t.Run("显式 auto 后发送且与 chat completions 一致", func(t *testing.T) { + t.Parallel() + payload, err := BuildRequest(t.Context(), provider.RuntimeConfig{ + Driver: provider.DriverOpenAICompat, + PromptCacheMode: provider.PromptCacheModeAuto, + }, req) + if err != nil { + t.Fatalf("BuildRequest() error = %v", err) + } + want := provider.ComputePromptCacheKey("gpt-5", "stable", []string{"read_file", "grep"}) + if payload.PromptCacheKey != want { + t.Fatalf("PromptCacheKey = %q, want %q", payload.PromptCacheKey, want) + } + }) + + t.Run("自研厂商恒不发送", func(t *testing.T) { + t.Parallel() + payload, err := BuildRequest(t.Context(), provider.RuntimeConfig{ + Driver: provider.DriverDeepSeek, + PromptCacheMode: provider.PromptCacheModeAuto, + }, req) + if err != nil { + t.Fatalf("BuildRequest() error = %v", err) + } + if payload.PromptCacheKey != "" { + t.Fatalf("PromptCacheKey = %q, want empty for self-implemented driver", payload.PromptCacheKey) + } + }) + + t.Run("无稳定前缀不发送", func(t *testing.T) { + t.Parallel() + noPrefix := req + noPrefix.SystemStablePrefix = "" + payload, err := BuildRequest(t.Context(), provider.RuntimeConfig{ + Driver: provider.DriverOpenAICompat, + PromptCacheMode: provider.PromptCacheModeAuto, + }, noPrefix) + if err != nil { + t.Fatalf("BuildRequest() error = %v", err) + } + if payload.PromptCacheKey != "" { + t.Fatalf("PromptCacheKey = %q, want empty without stable prefix", payload.PromptCacheKey) + } + }) +} diff --git a/internal/provider/openaicompat/responses/request.go b/internal/provider/openaicompat/responses/request.go index 3ee9553c6..1ded9e7e4 100644 --- a/internal/provider/openaicompat/responses/request.go +++ b/internal/provider/openaicompat/responses/request.go @@ -34,6 +34,16 @@ func BuildRequest(ctx context.Context, cfg provider.RuntimeConfig, req providert payload.Instructions = req.SystemPrompt } + // 与 chat completions 共用同一套键计算，保证同一前缀在两种模式下路由一致。 + // 该驱动服务任意兼容网关，默认不发送，需显式 prompt_cache: auto 声明端点支持。 + if provider.ShouldSendPromptCacheKey(cfg.PromptCacheMode, cfg.Driver) { + payload.PromptCacheKey = provider.ComputePromptCacheKey( + model, + req.ResolvedStablePrefix(), + provider.PromptCacheToolNames(req.Tools), + ) + } + assetPolicy := session.NormalizeAssetPolicy(cfg.SessionAssetPolicy) requestBudget := provider.NormalizeRequestAssetBudget(cfg.RequestAssetBudget, assetPolicy.MaxSessionAssetBytes) var usedSessionAssetBytes int64 diff --git a/internal/provider/openaicompat/responses/types.go b/internal/provider/openaicompat/responses/types.go index cc07f8bc7..0abeaa26e 100644 --- a/internal/provider/openaicompat/responses/types.go +++ b/internal/provider/openaicompat/responses/types.go @@ -8,6 +8,9 @@ type Request struct { Tools []ToolDefinition `json:"tools,omitempty"` ToolChoice string `json:"tool_choice,omitempty"` Stream bool `json:"stream"` + // PromptCacheKey 是缓存路由键，作用是让相同前缀的请求尽量落到同一后端， + // 提升自动前缀缓存的命中率；为空时不发送该字段。 + PromptCacheKey string `json:"prompt_cache_key,omitempty"` } // InputItem 表示 Responses API 中输入项（消息、函数调用、函数调用输出）。 diff --git a/internal/provider/prompt_cache.go b/internal/provider/prompt_cache.go new file mode 100644 index 000000000..13ad3f7e6 --- /dev/null +++ b/internal/provider/prompt_cache.go @@ -0,0 +1,113 @@ +package provider + +import ( + "crypto/sha256" + "encoding/hex" + "strings" + + providertypes "neo-code/internal/provider/types" +) + +// PromptCacheMode 控制 provider 是否在请求中携带 prompt 缓存标记。 +const ( + // PromptCacheModeAuto 表示按各 adapter 的能力启用缓存标记。 + PromptCacheModeAuto = "auto" + // PromptCacheModeOff 表示完全禁用缓存标记，用于规避缓存写入溢价或排障。 + PromptCacheModeOff = "off" +) + +// NormalizePromptCacheMode 归一化缓存模式配置，空值回落为 auto，未知值同样回落为 auto。 +// 仅用于不区分 driver 的场景；需要按 driver 取默认值时使用 ResolvePromptCacheMode。 +func NormalizePromptCacheMode(mode string) string { + switch strings.ToLower(strings.TrimSpace(mode)) { + case PromptCacheModeOff, "false", "disabled", "none": + return PromptCacheModeOff + default: + return PromptCacheModeAuto + } +} + +// DefaultPromptCacheMode 返回指定 driver 在未显式配置时的缓存模式。 +// 仅 anthropic 默认开启：cache_control 是官方协议字段，不会被端点拒绝。 +// openaicompat 服务任意兼容网关，prompt_cache_key 属于可选扩展参数，在未确认端点 +// 支持前默认发送会让原本可用的网关返回 400，因此默认关闭，由用户在 provider 配置中 +// 显式写 prompt_cache: auto 声明该端点支持。 +func DefaultPromptCacheMode(driver string) string { + if strings.EqualFold(strings.TrimSpace(driver), DriverAnthropic) { + return PromptCacheModeAuto + } + return PromptCacheModeOff +} + +// ResolvePromptCacheMode 结合显式配置与 driver 默认值，得出最终缓存模式。 +// 配置为空或无法识别时使用 driver 默认值，避免把拼写错误静默解释成"开启"。 +func ResolvePromptCacheMode(configuredMode string, driver string) string { + switch strings.ToLower(strings.TrimSpace(configuredMode)) { + case PromptCacheModeAuto: + return PromptCacheModeAuto + case PromptCacheModeOff, "false", "disabled", "none": + return PromptCacheModeOff + default: + return DefaultPromptCacheMode(driver) + } +} + +// PromptCacheEnabled 判断给定配置与 driver 下是否允许发送缓存标记。 +func PromptCacheEnabled(configuredMode string, driver string) bool { + return ResolvePromptCacheMode(configuredMode, driver) == PromptCacheModeAuto +} + +// SupportsPromptCacheKey 判断该 driver 是否允许发送 OpenAI 的 prompt_cache_key 路由键。 +// 五个内置厂商走的是自研 HTTP 实现，各自有独立的上下文缓存机制或对未知字段校验更严格， +// 因此永不发送该字段：DeepSeek 按前缀全局命中，Qwen / GLM 走厂商自有接口。 +// 该判断是硬约束，优先于显式 prompt_cache: auto 配置。 +func SupportsPromptCacheKey(driver string) bool { + switch strings.ToLower(strings.TrimSpace(driver)) { + case DriverDeepSeek, DriverQwen, DriverGLM, DriverMiMo, DriverMiniMax: + return false + default: + return true + } +} + +// ShouldSendPromptCacheKey 综合配置、driver 默认值与硬约束，决定是否发送缓存路由键。 +func ShouldSendPromptCacheKey(configuredMode string, driver string) bool { + if !SupportsPromptCacheKey(driver) { + return false + } + return PromptCacheEnabled(configuredMode, driver) +} + +// PromptCacheKeyHexLength 是缓存路由键的十六进制长度，远低于各厂商 64 字符的上限。 +const PromptCacheKeyHexLength = 32 + +// promptCacheKeyHexLength 是内部别名，保持 key 生成处可读。 +const promptCacheKeyHexLength = PromptCacheKeyHexLength + +// ComputePromptCacheKey 基于模型、稳定前缀与工具名集合生成缓存路由键。 +// 三者任一变化都应产生新键：不同前缀共享同一键会把请求错误地路由到一起。 +// 无稳定前缀时返回空串，调用方据此不发送该字段。 +func ComputePromptCacheKey(model string, stablePrefix string, toolNames []string) string { + if strings.TrimSpace(stablePrefix) == "" { + return "" + } + + digest := sha256.New() + digest.Write([]byte(strings.TrimSpace(model))) + digest.Write([]byte{0}) + digest.Write([]byte(stablePrefix)) + for _, name := range toolNames { + digest.Write([]byte{0}) + digest.Write([]byte(strings.TrimSpace(name))) + } + return hex.EncodeToString(digest.Sum(nil))[:promptCacheKeyHexLength] +} + +// PromptCacheToolNames 提取工具名列表，供 ComputePromptCacheKey 参与键计算。 +func PromptCacheToolNames(specs []providertypes.ToolSpec) []string { + names := make([]string, 0, len(specs)) + for _, spec := range specs { + names = append(names, spec.Name) + } + return names +} diff --git a/internal/provider/prompt_cache_test.go b/internal/provider/prompt_cache_test.go new file mode 100644 index 000000000..8da45452e --- /dev/null +++ b/internal/provider/prompt_cache_test.go @@ -0,0 +1,124 @@ +package provider + +import "testing" + +func TestNormalizePromptCacheMode(t *testing.T) { + t.Parallel() + + cases := map[string]string{ + "": PromptCacheModeAuto, + "auto": PromptCacheModeAuto, + " AUTO ": PromptCacheModeAuto, + "off": PromptCacheModeOff, + "Off": PromptCacheModeOff, + "false": PromptCacheModeOff, + "disabled": PromptCacheModeOff, + "none": PromptCacheModeOff, + "unknown": PromptCacheModeAuto, + } + for input, want := range cases { + if got := NormalizePromptCacheMode(input); got != want { + t.Fatalf("NormalizePromptCacheMode(%q) = %q, want %q", input, got, want) + } + } +} + +func TestDefaultPromptCacheMode(t *testing.T) { + t.Parallel() + + // 仅 anthropic 默认开启：cache_control 是官方协议字段。 + if got := DefaultPromptCacheMode(DriverAnthropic); got != PromptCacheModeAuto { + t.Fatalf("anthropic default = %q, want auto", got) + } + // openaicompat 服务任意兼容网关，未确认支持前不发送，避免 400 回归。 + if got := DefaultPromptCacheMode(DriverOpenAICompat); got != PromptCacheModeOff { + t.Fatalf("openaicompat default = %q, want off", got) + } + if got := DefaultPromptCacheMode(""); got != PromptCacheModeOff { + t.Fatalf("empty driver default = %q, want off", got) + } + if got := DefaultPromptCacheMode(" ANTHROPIC "); got != PromptCacheModeAuto { + t.Fatalf("driver matching must ignore case and spaces, got %q", got) + } +} + +func TestResolvePromptCacheMode(t *testing.T) { + t.Parallel() + + cases := []struct { + configured string + driver string + want string + }{ + {configured: "auto", driver: DriverOpenAICompat, want: PromptCacheModeAuto}, + {configured: "off", driver: DriverAnthropic, want: PromptCacheModeOff}, + {configured: "", driver: DriverAnthropic, want: PromptCacheModeAuto}, + {configured: "", driver: DriverOpenAICompat, want: PromptCacheModeOff}, + // 拼写错误不应被静默解释成开启，回落 driver 默认值。 + {configured: "always", driver: DriverAnthropic, want: PromptCacheModeAuto}, + {configured: "always", driver: DriverOpenAICompat, want: PromptCacheModeOff}, + } + for _, tc := range cases { + if got := ResolvePromptCacheMode(tc.configured, tc.driver); got != tc.want { + t.Fatalf("ResolvePromptCacheMode(%q, %q) = %q, want %q", tc.configured, tc.driver, got, tc.want) + } + } +} + +func TestPromptCacheEnabled(t *testing.T) { + t.Parallel() + + if !PromptCacheEnabled("", DriverAnthropic) { + t.Fatal("anthropic should be enabled by default") + } + if PromptCacheEnabled("", DriverOpenAICompat) { + t.Fatal("openaicompat should stay disabled until explicitly enabled") + } + if !PromptCacheEnabled(PromptCacheModeAuto, DriverOpenAICompat) { + t.Fatal("explicit auto should enable openaicompat") + } + if PromptCacheEnabled(PromptCacheModeOff, DriverAnthropic) { + t.Fatal("off mode should disable anthropic") + } +} + +func TestSupportsPromptCacheKey(t *testing.T) { + t.Parallel() + + // 自研 HTTP 实现的厂商不发送 OpenAI 的缓存路由键。 + unsupported := []string{DriverDeepSeek, DriverQwen, DriverGLM, DriverMiMo, DriverMiniMax} + for _, driver := range unsupported { + if SupportsPromptCacheKey(driver) { + t.Fatalf("driver %q should not send prompt_cache_key", driver) + } + } + if SupportsPromptCacheKey(" DEEPSEEK ") { + t.Fatal("driver matching must ignore case and surrounding spaces") + } + + eligible := []string{DriverOpenAICompat, "", "custom-gateway"} + for _, driver := range eligible { + if !SupportsPromptCacheKey(driver) { + t.Fatalf("driver %q should be eligible for prompt_cache_key", driver) + } + } +} + +func TestShouldSendPromptCacheKey(t *testing.T) { + t.Parallel() + + // 硬约束优先：即使显式开启，自研厂商也不发送。 + if ShouldSendPromptCacheKey(PromptCacheModeAuto, DriverDeepSeek) { + t.Fatal("deepseek must never send prompt_cache_key even when explicitly enabled") + } + // 默认关闭：未显式声明支持的兼容网关不发送。 + if ShouldSendPromptCacheKey("", DriverOpenAICompat) { + t.Fatal("openaicompat must stay disabled by default") + } + if !ShouldSendPromptCacheKey(PromptCacheModeAuto, DriverOpenAICompat) { + t.Fatal("explicit auto should enable openaicompat") + } + if ShouldSendPromptCacheKey(PromptCacheModeOff, DriverOpenAICompat) { + t.Fatal("off mode should disable") + } +} diff --git a/internal/provider/types/request.go b/internal/provider/types/request.go index e7a21e4d0..c4a16cf40 100644 --- a/internal/provider/types/request.go +++ b/internal/provider/types/request.go @@ -3,6 +3,7 @@ package types import ( "context" "io" + "strings" ) // SessionAssetReader 定义 provider 请求阶段读取会话附件内容的最小能力。 @@ -19,10 +20,36 @@ type ThinkingConfig struct { // GenerateRequest 是 provider.Generate() 的请求参数。 type GenerateRequest struct { - Model string `json:"model"` - SystemPrompt string `json:"system_prompt"` + Model string `json:"model"` + // SystemPrompt 是渲染完成的完整系统提示，由稳定前缀与易变尾部按序拼接而成。 + SystemPrompt string `json:"system_prompt"` + // SystemStablePrefix 是 SystemPrompt 中长期稳定的缓存前缀部分。 + // 约束：必须是 SystemPrompt 的真实前缀；不满足时 adapter 必须丢弃该提示，不得打缓存标记。 + SystemStablePrefix string `json:"system_stable_prefix,omitempty"` Messages []Message `json:"messages"` Tools []ToolSpec `json:"tools,omitempty"` ThinkingConfig *ThinkingConfig `json:"thinking_config,omitempty"` SessionAssetReader SessionAssetReader `json:"-"` } + +// ResolvedStablePrefix 返回可用于缓存打标的稳定前缀。 +// 当提示为空、或不是 SystemPrompt 的真实前缀时返回空串，调用方据此降级为不打标。 +func (r GenerateRequest) ResolvedStablePrefix() string { + if r.SystemStablePrefix == "" { + return "" + } + if !strings.HasPrefix(r.SystemPrompt, r.SystemStablePrefix) { + return "" + } + return r.SystemStablePrefix +} + +// VolatileSystemSuffix 返回系统提示中稳定前缀之后的部分，按原样保留分隔符与尾部内容。 +// 与 ResolvedStablePrefix 拼接后严格等于 SystemPrompt，保证分段发送不改变实际文本。 +func (r GenerateRequest) VolatileSystemSuffix() string { + prefix := r.ResolvedStablePrefix() + if prefix == "" { + return r.SystemPrompt + } + return r.SystemPrompt[len(prefix):] +} diff --git a/internal/provider/types/request_test.go b/internal/provider/types/request_test.go new file mode 100644 index 000000000..aa6e2573a --- /dev/null +++ b/internal/provider/types/request_test.go @@ -0,0 +1,101 @@ +package types + +import "testing" + +func TestGenerateRequestResolvedStablePrefix(t *testing.T) { + t.Parallel() + + cases := []struct { + name string + req GenerateRequest + want string + }{ + { + name: "空前缀返回空", + req: GenerateRequest{SystemPrompt: "stable dynamic"}, + want: "", + }, + { + name: "真实前缀被接受", + req: GenerateRequest{ + SystemPrompt: "stable\n\ndynamic", + SystemStablePrefix: "stable", + }, + want: "stable", + }, + { + name: "非真实前缀被拒", + req: GenerateRequest{ + SystemPrompt: "stable\n\ndynamic", + SystemStablePrefix: "other", + }, + want: "", + }, + { + name: "整段作为前缀", + req: GenerateRequest{ + SystemPrompt: "all-stable", + SystemStablePrefix: "all-stable", + }, + want: "all-stable", + }, + } + + for _, tc := range cases { + tc := tc + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + if got := tc.req.ResolvedStablePrefix(); got != tc.want { + t.Fatalf("ResolvedStablePrefix() = %q, want %q", got, tc.want) + } + }) + } +} + +func TestGenerateRequestVolatileSystemSuffixPreservesBytes(t *testing.T) { + t.Parallel() + + const stable = "## 角色\n你是编码助手。" + const dynamic = "## 当前任务\n修复登录。" + req := GenerateRequest{ + SystemPrompt: stable + "\n\n" + dynamic, + SystemStablePrefix: stable, + } + + if got, want := req.VolatileSystemSuffix(), "\n\n"+dynamic; got != want { + t.Fatalf("VolatileSystemSuffix() = %q, want %q", got, want) + } + if got := req.ResolvedStablePrefix() + req.VolatileSystemSuffix(); got != req.SystemPrompt { + t.Fatalf("prefix + suffix = %q, want %q", got, req.SystemPrompt) + } +} + +func TestGenerateRequestVolatileSystemSuffixWithoutPrefix(t *testing.T) { + t.Parallel() + + req := GenerateRequest{SystemPrompt: "whole"} + if got := req.VolatileSystemSuffix(); got != "whole" { + t.Fatalf("VolatileSystemSuffix() = %q, want %q", got, "whole") + } +} + +func TestUsageTotalInputTokensAndHitRate(t *testing.T) { + t.Parallel() + + usage := Usage{ + InputTokens: 500, + CacheReadInputTokens: 1500, + CacheWriteInputTokens: 0, + OutputTokens: 200, + } + if got, want := usage.TotalInputTokens(), 2000; got != want { + t.Fatalf("TotalInputTokens() = %d, want %d", got, want) + } + if got, want := usage.CacheHitRate(), 0.75; got != want { + t.Fatalf("CacheHitRate() = %v, want %v", got, want) + } + + if got := (Usage{}).CacheHitRate(); got != 0 { + t.Fatalf("CacheHitRate() on empty usage = %v, want 0", got) + } +} diff --git a/internal/provider/types/usage.go b/internal/provider/types/usage.go index c8c76e260..663b8ecb7 100644 --- a/internal/provider/types/usage.go +++ b/internal/provider/types/usage.go @@ -7,6 +7,26 @@ type Usage struct { TotalTokens int `json:"total_tokens"` InputObserved bool `json:"input_observed"` OutputObserved bool `json:"output_observed"` + // CacheReadInputTokens 记录命中 prompt 缓存读取的输入 token 数。 + // Anthropic 的 cache_read_input_tokens、OpenAI 的 cached_tokens 均映射到此字段。 + CacheReadInputTokens int `json:"cache_read_input_tokens,omitempty"` + // CacheWriteInputTokens 记录写入 prompt 缓存的输入 token 数，仅 Anthropic 显式缓存有该口径。 + CacheWriteInputTokens int `json:"cache_write_input_tokens,omitempty"` +} + +// TotalInputTokens 返回计入预算的总输入 token，包含缓存命中与缓存写入部分。 +// Anthropic 的 input_tokens 不含缓存相关计数，预算与 compact 必须使用本方法以免低估。 +func (u Usage) TotalInputTokens() int { + return u.InputTokens + u.CacheReadInputTokens + u.CacheWriteInputTokens +} + +// CacheHitRate 返回缓存命中的输入占比，无可统计输入时返回 0。 +func (u Usage) CacheHitRate() float64 { + total := u.TotalInputTokens() + if total <= 0 { + return 0 + } + return float64(u.CacheReadInputTokens) / float64(total) } // BudgetEstimate 描述 provider 对冻结请求输入 token 的估算结果。 diff --git a/internal/runtime/budget_models.go b/internal/runtime/budget_models.go index 12c8ec94e..21597b3c9 100644 --- a/internal/runtime/budget_models.go +++ b/internal/runtime/budget_models.go @@ -33,6 +33,9 @@ type TurnBudgetUsageObservation struct { OutputTokens int InputObserved bool OutputObserved bool + // CacheReadInputTokens 与 CacheWriteInputTokens 是 prompt 缓存的命中与写入明细。 + CacheReadInputTokens int + CacheWriteInputTokens int } // turnProviderOutput 汇总一次 provider 调用返回的 assistant 消息与预算 usage observation。 @@ -48,6 +51,9 @@ type ledgerReconcileResult struct { outputTokens int outputSource string hasUnknownUsage bool + // cacheReadTokens 与 cacheWriteTokens 透传观测到的缓存明细，供对账事件展示。 + cacheReadTokens int + cacheWriteTokens int } // newTurnBudgetSnapshot 构造本次发送尝试的冻结预算快照。 @@ -113,12 +119,16 @@ func newTurnBudgetUsageObservation( outputTokens int, inputObserved bool, outputObserved bool, + cacheReadTokens int, + cacheWriteTokens int, ) TurnBudgetUsageObservation { return TurnBudgetUsageObservation{ - ID: id, - InputTokens: inputTokens, - OutputTokens: outputTokens, - InputObserved: inputObserved, - OutputObserved: outputObserved, + ID: id, + InputTokens: inputTokens, + OutputTokens: outputTokens, + InputObserved: inputObserved, + OutputObserved: outputObserved, + CacheReadInputTokens: cacheReadTokens, + CacheWriteInputTokens: cacheWriteTokens, } } diff --git a/internal/runtime/events.go b/internal/runtime/events.go index 24c17fb4b..3fd76779c 100644 --- a/internal/runtime/events.go +++ b/internal/runtime/events.go @@ -117,6 +117,9 @@ type LedgerReconciledPayload struct { OutputTokens int `json:"output_tokens"` OutputSource string `json:"output_source"` HasUnknownUsage bool `json:"has_unknown_usage"` + // CacheReadInputTokens 与 CacheWriteInputTokens 暴露 prompt 缓存命中与写入明细。 + CacheReadInputTokens int `json:"cache_read_input_tokens,omitempty"` + CacheWriteInputTokens int `json:"cache_write_input_tokens,omitempty"` } // newBudgetCheckedPayload 将预算决策对象展开为对外事件 payload，保持可观测字段稳定。 @@ -152,13 +155,15 @@ func newLedgerReconciledPayload( result ledgerReconcileResult, ) LedgerReconciledPayload { return LedgerReconciledPayload{ - AttemptSeq: observation.ID.AttemptSeq, - RequestHash: observation.ID.RequestHash, - InputTokens: result.inputTokens, - InputSource: result.inputSource, - OutputTokens: result.outputTokens, - OutputSource: result.outputSource, - HasUnknownUsage: result.hasUnknownUsage, + AttemptSeq: observation.ID.AttemptSeq, + RequestHash: observation.ID.RequestHash, + InputTokens: result.inputTokens, + InputSource: result.inputSource, + OutputTokens: result.outputTokens, + OutputSource: result.outputSource, + HasUnknownUsage: result.hasUnknownUsage, + CacheReadInputTokens: result.cacheReadTokens, + CacheWriteInputTokens: result.cacheWriteTokens, } } @@ -460,6 +465,10 @@ type TokenUsagePayload struct { HasUnknownUsage bool `json:"has_unknown_usage,omitempty"` SessionInputTokens int `json:"session_input_tokens"` SessionOutputTokens int `json:"session_output_tokens"` + // CacheReadInputTokens 与 CacheWriteInputTokens 暴露 prompt 缓存命中与写入明细， + // 供客户端计算命中率；为空表示本轮没有发生缓存读写。 + CacheReadInputTokens int `json:"cache_read_input_tokens,omitempty"` + CacheWriteInputTokens int `json:"cache_write_input_tokens,omitempty"` } // CheckpointCreatedPayload 描述 checkpoint 创建成功事件。 diff --git a/internal/runtime/ledger_cache_payload_test.go b/internal/runtime/ledger_cache_payload_test.go new file mode 100644 index 000000000..4a0850609 --- /dev/null +++ b/internal/runtime/ledger_cache_payload_test.go @@ -0,0 +1,39 @@ +package runtime + +import ( + "testing" + + "neo-code/internal/runtime/controlplane" +) + +func TestNewLedgerReconciledPayloadCarriesCacheUsage(t *testing.T) { + t.Parallel() + + observation := TurnBudgetUsageObservation{ + ID: controlplane.TurnBudgetID{AttemptSeq: 2, RequestHash: "hash-1"}, + InputTokens: 1200, + OutputTokens: 80, + InputObserved: true, + OutputObserved: true, + CacheReadInputTokens: 900, + CacheWriteInputTokens: 0, + } + result := ledgerReconcileResult{ + inputTokens: 1200, + inputSource: usageSourceObserved, + outputTokens: 80, + outputSource: usageSourceObserved, + cacheReadTokens: 900, + } + + payload := newLedgerReconciledPayload(observation, result) + if payload.CacheReadInputTokens != 900 { + t.Fatalf("CacheReadInputTokens = %d, want 900", payload.CacheReadInputTokens) + } + if payload.CacheWriteInputTokens != 0 { + t.Fatalf("CacheWriteInputTokens = %d, want 0", payload.CacheWriteInputTokens) + } + if payload.InputTokens != 1200 { + t.Fatalf("InputTokens = %d, want 1200", payload.InputTokens) + } +} diff --git a/internal/runtime/prompt_cache_usage_flow_test.go b/internal/runtime/prompt_cache_usage_flow_test.go new file mode 100644 index 000000000..1963aa895 --- /dev/null +++ b/internal/runtime/prompt_cache_usage_flow_test.go @@ -0,0 +1,124 @@ +package runtime + +import ( + "context" + "testing" + + providertypes "neo-code/internal/provider/types" + "neo-code/internal/runtime/controlplane" + "neo-code/internal/runtime/streaming" +) + +// cacheUsageProvider 返回一次携带缓存读写的 message_done，用于验证 usage 穿过生成链路。 +type cacheUsageProvider struct { + usage *providertypes.Usage +} + +func (p *cacheUsageProvider) EstimateInputTokens( + ctx context.Context, + req providertypes.GenerateRequest, +) (providertypes.BudgetEstimate, error) { + return providertypes.BudgetEstimate{EstimatedInputTokens: 1}, nil +} + +func (p *cacheUsageProvider) Generate( + ctx context.Context, + req providertypes.GenerateRequest, + events chan<- providertypes.StreamEvent, +) error { + events <- providertypes.NewTextDeltaStreamEvent("done") + events <- providertypes.NewMessageDoneStreamEvent("end_turn", p.usage) + return nil +} + +func TestGenerateStreamingMessageSumsCachedInputTokens(t *testing.T) { + t.Parallel() + + // Anthropic 口径：input_tokens 不含缓存部分，缓存读写单独上报。 + usage := &providertypes.Usage{ + InputTokens: 500, + OutputTokens: 80, + InputObserved: true, + OutputObserved: true, + CacheReadInputTokens: 1800, + CacheWriteInputTokens: 1200, + } + + result := generateStreamingMessage( + context.Background(), + &cacheUsageProvider{usage: usage}, + providertypes.GenerateRequest{Model: "claude-sonnet-4"}, + streaming.Hooks{}, + ) + if result.err != nil { + t.Fatalf("generateStreamingMessage() err = %v", result.err) + } + + // 关键断言：必须累计含缓存的总输入，否则账本低估会让 compact 永不触发。 + if result.inputTokens != 3500 { + t.Fatalf("inputTokens = %d, want 3500 (500 + 1800 + 1200)", result.inputTokens) + } + if result.cacheReadTokens != 1800 { + t.Fatalf("cacheReadTokens = %d, want 1800", result.cacheReadTokens) + } + if result.cacheWriteTokens != 1200 { + t.Fatalf("cacheWriteTokens = %d, want 1200", result.cacheWriteTokens) + } + if !result.inputObserved || !result.outputObserved { + t.Fatal("observed flags must be preserved") + } +} + +func TestReconcileLedgerPreservesCacheUsage(t *testing.T) { + t.Parallel() + + service := &Service{} + id := controlplane.TurnBudgetID{AttemptSeq: 1, RequestHash: "hash"} + observation := TurnBudgetUsageObservation{ + ID: id, + InputTokens: 3500, + OutputTokens: 80, + InputObserved: true, + OutputObserved: true, + CacheReadInputTokens: 1800, + CacheWriteInputTokens: 1200, + } + decision := controlplane.TurnBudgetDecision{ + ID: id, + EstimatedInputTokens: 9999, + } + + result, err := service.reconcileLedger(nil, decision, observation) + if err != nil { + t.Fatalf("reconcileLedger() error = %v", err) + } + // 已观测时必须使用观测值而不是估算值，否则缓存写读会被估算覆盖。 + if result.inputTokens != 3500 { + t.Fatalf("inputTokens = %d, want 3500 from observation", result.inputTokens) + } + if result.cacheReadTokens != 1800 || result.cacheWriteTokens != 1200 { + t.Fatalf("cache detail = (%d, %d), want (1800, 1200)", result.cacheReadTokens, result.cacheWriteTokens) + } + + payload := newLedgerReconciledPayload(observation, result) + if payload.CacheReadInputTokens != 1800 || payload.CacheWriteInputTokens != 1200 { + t.Fatalf("ledger payload cache detail = (%d, %d), want (1800, 1200)", + payload.CacheReadInputTokens, payload.CacheWriteInputTokens) + } + + tokenPayload := TokenUsagePayload{ + InputTokens: result.inputTokens, + OutputTokens: result.outputTokens, + CacheReadInputTokens: result.cacheReadTokens, + CacheWriteInputTokens: result.cacheWriteTokens, + } + + // token_usage 事件同样必须携带缓存明细，否则客户端无法计算命中率。 + if tokenPayload.CacheReadInputTokens != 1800 || tokenPayload.CacheWriteInputTokens != 1200 { + t.Fatalf("token usage payload cache detail = (%d, %d), want (1800, 1200)", + tokenPayload.CacheReadInputTokens, tokenPayload.CacheWriteInputTokens) + } + if tokenPayload.InputTokens != 3500 { + t.Fatalf("token usage payload InputTokens = %d, want 3500", tokenPayload.InputTokens) + } +} diff --git a/internal/runtime/provider_stream.go b/internal/runtime/provider_stream.go index d559de321..f3b9d8038 100644 --- a/internal/runtime/provider_stream.go +++ b/internal/runtime/provider_stream.go @@ -17,7 +17,11 @@ type streamGenerateResult struct { outputTokens int inputObserved bool outputObserved bool - err error + // cacheReadTokens 与 cacheWriteTokens 记录 prompt 缓存命中与写入的输入 token 数， + // 仅用于可观测性，不计入 inputTokens（inputTokens 已是含缓存的输入总量）。 + cacheReadTokens int + cacheWriteTokens int + err error } // generateStreamingMessage 负责执行一次基于流式事件的生成调用，并收敛最终 assistant 消息与 usage。 @@ -41,10 +45,14 @@ func generateStreamingMessage( hooksCopy := hooks hooksCopy.OnMessageDone = func(payload providertypes.MessageDonePayload) { if payload.Usage != nil { - outcome.inputTokens = payload.Usage.InputTokens + // 必须使用 TotalInputTokens：厂商原始 input_tokens 不含缓存命中与写入部分， + // 直接取值会让预算累计与 compact 判定在开启 prompt 缓存后系统性低估。 + outcome.inputTokens = payload.Usage.TotalInputTokens() outcome.outputTokens = payload.Usage.OutputTokens outcome.inputObserved = payload.Usage.InputObserved outcome.outputObserved = payload.Usage.OutputObserved + outcome.cacheReadTokens = payload.Usage.CacheReadInputTokens + outcome.cacheWriteTokens = payload.Usage.CacheWriteInputTokens } if userOnMessageDone != nil { userOnMessageDone(payload) diff --git a/internal/runtime/run.go b/internal/runtime/run.go index 82cd9ed30..6e7f08ffa 100644 --- a/internal/runtime/run.go +++ b/internal/runtime/run.go @@ -615,6 +615,9 @@ func (s *Service) prepareTurnBudgetSnapshot(ctx context.Context, state *runState repeatLimit := resolveRepeatCycleStreakLimit(cfg.Runtime) var systemPrompt string + // systemStablePrefix 记录 systemPrompt 中可用于 prompt 缓存打标的稳定前缀。 + // 仅在能够证明它为最终文本的真实前缀时赋值，其余情况保持为空以触发 adapter 降级。 + var systemStablePrefix string if builtContext.DynamicSystemPrompt != "" { // New style: reminders append to dynamicPrompt, then join with stable prompt := withProgressReminder(builtContext.DynamicSystemPrompt, score) @@ -625,6 +628,12 @@ func (s *Service) prepareTurnBudgetSnapshot(ctx context.Context, state *runState prompt = mergeEphemeralHookNotificationIntoSystemPrompt(prompt, notificationHint) } systemPrompt = joinRuntimeSystemPromptParts(builtContext.StableSystemPrompt, prompt) + // joinRuntimeSystemPromptParts 会裁剪两端空白，故取裁剪后的稳定段作为前缀候选， + // 并在确认构成真实前缀后才采用，避免裁剪差异导致缓存标记指向错误的字节边界。 + stableCandidate := strings.TrimSpace(builtContext.StableSystemPrompt) + if stableCandidate != "" && strings.HasPrefix(systemPrompt, stableCandidate) { + systemStablePrefix = stableCandidate + } } else { // Legacy fallback: reminders append to SystemPrompt directly systemPrompt = withProgressReminder(builtContext.SystemPrompt, score) @@ -660,6 +669,7 @@ func (s *Service) prepareTurnBudgetSnapshot(ctx context.Context, state *runState request := providertypes.GenerateRequest{ Model: model, SystemPrompt: systemPrompt, + SystemStablePrefix: systemStablePrefix, Messages: requestMessages, Tools: toolSpecs, ThinkingConfig: thinkingCfg, @@ -773,6 +783,8 @@ func (s *Service) callProvider( streamOutcome.outputTokens, streamOutcome.inputObserved, streamOutcome.outputObserved, + streamOutcome.cacheReadTokens, + streamOutcome.cacheWriteTokens, ), }, nil } @@ -831,13 +843,15 @@ func (s *Service) emitTokenUsage(ctx context.Context, state *runState, result le return } s.emitRunScoped(ctx, EventTokenUsage, state, TokenUsagePayload{ - InputTokens: result.inputTokens, - OutputTokens: result.outputTokens, - InputSource: result.inputSource, - OutputSource: result.outputSource, - HasUnknownUsage: result.hasUnknownUsage, - SessionInputTokens: state.session.TokenInputTotal, - SessionOutputTokens: state.session.TokenOutputTotal, + InputTokens: result.inputTokens, + OutputTokens: result.outputTokens, + InputSource: result.inputSource, + OutputSource: result.outputSource, + HasUnknownUsage: result.hasUnknownUsage, + SessionInputTokens: state.session.TokenInputTotal, + SessionOutputTokens: state.session.TokenOutputTotal, + CacheReadInputTokens: result.cacheReadTokens, + CacheWriteInputTokens: result.cacheWriteTokens, }) } @@ -1054,8 +1068,10 @@ func (s *Service) reconcileLedger( return ledgerReconcileResult{}, fmt.Errorf("runtime: turn budget id mismatch between decision and usage observation") } reconciled := ledgerReconcileResult{ - inputSource: usageSourceUnknown, - outputSource: usageSourceUnknown, + inputSource: usageSourceUnknown, + outputSource: usageSourceUnknown, + cacheReadTokens: observation.CacheReadInputTokens, + cacheWriteTokens: observation.CacheWriteInputTokens, } if observation.InputObserved { reconciled.inputTokens = observation.InputTokens