feat: publish sanitized custom.97 source

relay/common/tool_usage.go

@@ -1,6 +1,7 @@
 package common
 
 import (
+	"cmp"
 	"crypto/sha256"
 	"encoding/hex"
 	"fmt"
@@ -11,39 +12,44 @@ import (
 	"github.com/QuantumNous/new-api/setting/operation_setting"
 )
 
+// GoogleSearchGroundedPromptTool bills Gemini 2.5-and-older Google Search
+// grounding, charged once per grounded prompt.
+const GoogleSearchGroundedPromptTool = "google_search_grounded_prompt"
+
+// reservedBillableToolNames are hosted-tool keys without a built-in price
+// (Anthropic server tools reported in usage.server_tool_use). They and every
+// tool name the built-in price seed prices
+// (operation_setting.IsBuiltInToolPriceKey) bill hosted tools; a client
+// function or tool_use with the same name is never counted under them.
 var reservedBillableToolNames = map[string]struct{}{
-	dto.BuildInToolWebSearchPreview: {},
-	dto.BuildInToolWebSearch:        {},
-	dto.BuildInToolFileSearch:       {},
-	dto.BuildInToolGoogleSearch:     {},
-	dto.BuildInToolImageGeneration:  {},
+	"web_fetch":      {},
+	"code_execution": {},
+	"tool_search":    {},
 }
 
+// MaxBillableToolCallCount bounds every tool count at settlement, whatever its
+// source, before it becomes a surcharge multiplier (see
+// .agents/rules/billing.md).
+const MaxBillableToolCallCount = 10_000
+
 // CountBillableToolCall is the single entry point for per-call tool billing counts.
 // Built-in call types always count; custom function/tool_use names only count when priced.
 func (info *RelayInfo) CountBillableToolCall(itemType string, functionName string) {
 	if info == nil {
 		return
 	}
-	if info.ResponsesUsageInfo == nil {
-		info.ResponsesUsageInfo = &ResponsesUsageInfo{
-			BuiltInTools: make(map[string]*BuildInToolInfo),
-		}
-	}
-	if info.ResponsesUsageInfo.BuiltInTools == nil {
-		info.ResponsesUsageInfo.BuiltInTools = make(map[string]*BuildInToolInfo)
-	}
+	info.ensureBuiltInTools()
 
 	switch itemType {
 	case dto.BuildInCallWebSearchCall:
-		info.incrementBillableToolCall(resolveWebSearchToolName(info.ResponsesUsageInfo.BuiltInTools))
+		info.incrementBillableToolCall(cmp.Or(info.WebSearchBillingKey, resolveWebSearchToolName(info.ResponsesUsageInfo.BuiltInTools)))
 	case dto.BuildInCallFileSearchCall:
 		info.incrementBillableToolCall(dto.BuildInToolFileSearch)
 	case dto.BuildInCallFunctionCall, dto.BuildInCallToolUse:
 		if functionName == "" {
 			return
 		}
-		if _, reserved := reservedBillableToolNames[functionName]; reserved {
+		if _, reserved := reservedBillableToolNames[functionName]; reserved || operation_setting.IsBuiltInToolPriceKey(functionName) {
 			return
 		}
 		if operation_setting.GetToolPriceForModel(functionName, info.GetBillingModelName()) <= 0 {
@@ -53,6 +59,59 @@ func (info *RelayInfo) CountBillableToolCall(itemType string, functionName strin
 	}
 }
 
+// SetBillableToolCount records an upstream-reported count for one tool-price
+// key, replacing whatever per-item counting produced for the same key. Zero
+// is a valid value: a vendor billing field that reports zero suppresses the
+// protocol items counted under that key. Negative counts are ignored.
+func (info *RelayInfo) SetBillableToolCount(name string, count int) {
+	if info == nil || name == "" {
+		return
+	}
+	if count < 0 {
+		common.SysError(common.LogText("billable tool count ignored: tool=%s count=%d", name, count))
+		return
+	}
+	info.ensureBuiltInTools()
+	if existing, ok := info.ResponsesUsageInfo.BuiltInTools[name]; ok && existing != nil {
+		existing.CallCount = count
+		return
+	}
+	info.ResponsesUsageInfo.BuiltInTools[name] = &BuildInToolInfo{ToolName: name, CallCount: count}
+}
+
+// ApplyVendorToolUsage writes the vendor-billed counts of one upstream
+// response into BuiltInTools. Protocol items counted earlier for the same key
+// are replaced; keys the vendor does not mention are kept.
+func (info *RelayInfo) ApplyVendorToolUsage(raw []byte) {
+	if info == nil || info.VendorToolUsage == nil || len(raw) == 0 {
+		return
+	}
+	for name, count := range info.VendorToolUsage(raw) {
+		info.SetBillableToolCount(name, count)
+	}
+}
+
+// RecordRequestInferredToolCall bills one call of key that the request itself
+// proves: a forced hosted search on an upstream whose response reports no
+// search usage. The record belongs to the current channel attempt;
+// InitChannelMeta drops it when the request moves to another channel.
+func (info *RelayInfo) RecordRequestInferredToolCall(key string) {
+	if info == nil || key == "" {
+		return
+	}
+	info.SetBillableToolCount(key, 1)
+	info.requestInferredToolKey = key
+}
+
+func (info *RelayInfo) ensureBuiltInTools() {
+	if info.ResponsesUsageInfo == nil {
+		info.ResponsesUsageInfo = &ResponsesUsageInfo{}
+	}
+	if info.ResponsesUsageInfo.BuiltInTools == nil {
+		info.ResponsesUsageInfo.BuiltInTools = make(map[string]*BuildInToolInfo)
+	}
+}
+
 func resolveWebSearchToolName(tools map[string]*BuildInToolInfo) string {
 	if _, ok := tools[dto.BuildInToolWebSearchPreview]; ok {
 		return dto.BuildInToolWebSearchPreview
@@ -125,15 +184,6 @@ func (c *ImageGenerationCallCounter) Observe(item *dto.ResponsesOutput, outputIn
 	c.count++
 }
 
-// Reset clears pending observations (used when a terminal response fails).
-func (c *ImageGenerationCallCounter) Reset() {
-	if c == nil {
-		return
-	}
-	c.seen = nil
-	c.count = 0
-}
-
 // Count returns the deduplicated completed image output count before commit capping.
 func (c *ImageGenerationCallCounter) Count() int {
 	if c == nil {
@@ -148,14 +198,7 @@ func (c *ImageGenerationCallCounter) Commit(info *RelayInfo) {
 	if info == nil {
 		return
 	}
-	if info.ResponsesUsageInfo == nil {
-		info.ResponsesUsageInfo = &ResponsesUsageInfo{
-			BuiltInTools: make(map[string]*BuildInToolInfo),
-		}
-	}
-	if info.ResponsesUsageInfo.BuiltInTools == nil {
-		info.ResponsesUsageInfo.BuiltInTools = make(map[string]*BuildInToolInfo)
-	}
+	info.ensureBuiltInTools()
 
 	count := 0
 	if c != nil {
@@ -174,21 +217,3 @@ func (c *ImageGenerationCallCounter) Commit(info *RelayInfo) {
 		CallCount: count,
 	}
 }
-
-// IsNonBillableResponsesStatus reports terminal response statuses that must not
-// bill pending image_generation observations.
-func IsNonBillableResponsesStatus(status []byte) bool {
-	if len(status) == 0 {
-		return false
-	}
-	var s string
-	if err := common.Unmarshal(status, &s); err != nil {
-		return false
-	}
-	switch strings.ToLower(strings.TrimSpace(s)) {
-	case "failed", "cancelled", "canceled", "incomplete":
-		return true
-	default:
-		return false
-	}
-}