Skip to content

feat: support Claude messages token counting #5596

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
lyzgeorge wants to merge 2 commits into
base: main
Choose a base branch
from
Open

feat: support Claude messages token counting #5596

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
ae7ac65
fix(subscription): apply gateway price to plan payments
Jun 18, 2026
dd76d7b
feat: support Claude messages token counting
Jun 18, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
28 changes: 16 additions & 12 deletions controller/relay.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -150,21 +150,25 @@ func Relay(c *gin.Context, relayFormat types.RelayFormat) {

relayInfo.SetEstimatePromptTokens(tokens)

priceData, err := helper.ModelPriceHelper(c, relayInfo, tokens, meta)
if err != nil {
newAPIError = types.NewError(err, types.ErrorCodeModelPriceError, types.ErrOptionWithStatusCode(http.StatusBadRequest))
return
}

// common.SetContextKey(c, constant.ContextKeyTokenCountMeta, meta)

if priceData.FreeModel {
logger.LogInfo(c, fmt.Sprintf("模型 %s 免费，跳过预扣费", relayInfo.OriginModelName))
if relaycommon.IsClaudeCountTokensRequest(relayInfo) {
relayInfo.FinalPreConsumedQuota = 0
} else {
newAPIError = service.PreConsumeBilling(c, priceData.QuotaToPreConsume, relayInfo)
if newAPIError != nil {
priceData, err := helper.ModelPriceHelper(c, relayInfo, tokens, meta)
if err != nil {
newAPIError = types.NewError(err, types.ErrorCodeModelPriceError, types.ErrOptionWithStatusCode(http.StatusBadRequest))
return
}

// common.SetContextKey(c, constant.ContextKeyTokenCountMeta, meta)

if priceData.FreeModel {
logger.LogInfo(c, fmt.Sprintf("模型 %s 免费，跳过预扣费", relayInfo.OriginModelName))
} else {
newAPIError = service.PreConsumeBilling(c, priceData.QuotaToPreConsume, relayInfo)
if newAPIError != nil {
return
}
}
}

defer func() {
Expand Down
14 changes: 12 additions & 2 deletions controller/subscription_payment_epay.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -84,10 +84,20 @@ func SubscriptionRequestEpay(c *gin.Context) {
return
}

group, err := model.GetUserGroup(userId, true)
if err != nil {
c.JSON(http.StatusOK, gin.H{"message": "error", "data": "获取用户分组失败"})
return
}
payMoney := getPayMoney(int64(plan.PriceAmount), group)

@coderabbitai coderabbitai Bot Jun 18, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟠 Major

🧩 Analysis chain

🏁 Script executed:

#!/bin/bash
# Verify plan.PriceAmount type and getPayMoney signature
echo "=== Check SubscriptionPlan model for PriceAmount type ==="
rg -n "PriceAmount" --type go -C3 | head -60

echo ""
echo "=== Check getPayMoney signature ==="
ast-grep --pattern 'func getPayMoney($$$) $_ {
  $$$
}'

Repository: QuantumNous/new-api

Length of output: 5059

🏁 Script executed:

#!/bin/bash
# Verify the actual line 92 in subscription_payment_epay.go
echo "=== Lines 85-105 in subscription_payment_epay.go ==="
sed -n '85,105p' controller/subscription_payment_epay.go

echo ""
echo "=== Check SubscriptionOrder Money field type ==="
rg -n "type SubscriptionOrder struct" -A 20 model/subscription.go | head -30

echo ""
echo "=== Search for cents convention patterns ==="
rg -n "int64.*100\|cents\|\* 100" --type go controller/ | head -20

Repository: QuantumNous/new-api

Length of output: 1651

🏁 Script executed:

#!/bin/bash
# Check how other payment methods handle plan.PriceAmount
echo "=== Check other payment method handlers (balance, stripe, etc) ==="
rg -n "payMoney\|Money.*plan.PriceAmount\|Money.*plan\.Price" --type go controller/ -B2 -A2 | head -50

echo ""
echo "=== Check if there's a float64-based payment function ==="
rg -n "func.*Payment.*float64\|func.*Money.*float64" --type go controller/ | grep -v test

echo ""
echo "=== Look at balance payment implementation for comparison ==="
rg -n "PaymentMethodBalance\|PaymentProviderBalance" --type go -A10 controller/subscription*.go | head -60

Repository: QuantumNous/new-api

Length of output: 252

🏁 Script executed:

#!/bin/bash
# Check the full subscription payment file for other methods
echo "=== Check subscription_payment files ==="
ls -la controller/subscription_payment*.go

echo ""
echo "=== Check balance payment method in detail ==="
rg -n "handleSubscriptionBalancePay\|PaymentMethodBalance" --type go -B5 -A15 model/subscription.go | head -80

echo ""
echo "=== Check how getPayMoney is used elsewhere ==="
rg -n "getPayMoney" --type go -B2 -A2

Repository: QuantumNous/new-api

Length of output: 1689

🏁 Script executed:

#!/bin/bash
# Check other subscription payment methods for how they handle plan.PriceAmount
echo "=== Check subscription_payment_stripe.go ==="
sed -n '1,120p' controller/subscription_payment_stripe.go | grep -A5 -B5 "plan.PriceAmount\|Money.*plan"

echo ""
echo "=== Check subscription_payment_creem.go ==="
sed -n '1,120p' controller/subscription_payment_creem.go | grep -A5 -B5 "plan.PriceAmount\|Money.*plan"

echo ""
echo "=== Check subscription_payment_waffo_pancake.go ==="
sed -n '1,120p' controller/subscription_payment_waffo_pancake.go | grep -A5 -B5 "plan.PriceAmount\|Money.*plan"

Repository: QuantumNous/new-api

Length of output: 1532

🏁 Script executed:

#!/bin/bash
# Check getWaffoPancakePayMoney function
echo "=== Check getWaffoPancakePayMoney function ==="
rg -n "func getWaffoPancakePayMoney" --type go -A20

echo ""
echo "=== Check complete epay payment file to understand context ==="
cat -n controller/subscription_payment_epay.go | sed -n '1,120p'

Repository: QuantumNous/new-api

Length of output: 5439

🏁 Script executed:

#!/bin/bash
# Final verification: Check if topup Amount is in cents or dollars
echo "=== Check TopupRequest structure ==="
rg -n "type TopupRequest\|type.*Request struct" --type go controller/topup.go -A10 | head -40

echo ""
echo "=== Check how req.Amount is used in topup vs subscription ==="
echo "In topup.go context:"
rg -n "req.Amount" controller/topup.go | head -10

echo ""
echo "=== Verify if there's a cents convention in topup ==="
rg -n "Amount.*100\|cents" --type go controller/topup.go | head -10

Repository: QuantumNous/new-api

Length of output: 1132

int64() truncation loses fractional price amounts—affects billing accuracy

plan.PriceAmount is float64 (stored as decimal(10,6) in the database), but casting to int64 truncates decimals. A plan priced at $9.99 becomes $9, undercharging by $0.99.

The issue appears in both subscription payment methods:

  • subscription_payment_epay.go line 92
  • subscription_payment_waffo_pancake.go (similar pattern)

Other payment methods (Stripe, Creem) avoid this by using plan.PriceAmount directly without conversion.

Consider either:

  1. Passing cents: int64(plan.PriceAmount * 100) and adjusting getPayMoney to work in cents, or
  2. Creating a float-aware helper for subscription pricing that preserves precision (similar to how Stripe/Creem handle it).
🤖 Prompt for AI Agents 
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@controller/subscription_payment_epay.go` at line 92, The int64 cast in the
getPayMoney call is truncating the fractional cents from plan.PriceAmount (a
float64 representing decimal currency amounts). Instead of casting directly to
int64, multiply plan.PriceAmount by 100 to convert dollars to cents, then cast
to int64, ensuring full precision is preserved throughout the payment
calculation. Make the corresponding adjustment to the getPayMoney function
signature and implementation to handle the input as cents rather than dollars.
Apply this same fix to the identical pattern in
subscription_payment_waffo_pancake.go to ensure consistency across all payment
methods.

if payMoney < 0.01 {
common.ApiErrorMsg(c, "套餐金额过低")
return
}
order := &model.SubscriptionOrder{
UserId: userId,
PlanId: plan.Id,
Money: plan.PriceAmount,
Money: payMoney,
TradeNo: tradeNo,
PaymentMethod: req.PaymentMethod,
PaymentProvider: model.PaymentProviderEpay,
Expand All @@ -102,7 +112,7 @@ func SubscriptionRequestEpay(c *gin.Context) {
Type: req.PaymentMethod,
ServiceTradeNo: tradeNo,
Name: fmt.Sprintf("SUB:%s", plan.Title),
Money: strconv.FormatFloat(plan.PriceAmount, 'f', 2, 64),
Money: strconv.FormatFloat(payMoney, 'f', 2, 64),
Device: epay.PC,
NotifyUrl: notifyUrl,
ReturnUrl: returnUrl,
Expand Down
11 changes: 8 additions & 3 deletions controller/subscription_payment_waffo_pancake.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -79,10 +79,15 @@ func SubscriptionRequestWaffoPancakePay(c *gin.Context) {
// dispatch in WaffoPancakeWebhook.
tradeNo := fmt.Sprintf("WAFFO_PANCAKE_SUB-%d-%d-%s", userId, time.Now().UnixMilli(), randstr.String(6))

payMoney := getWaffoPancakePayMoney(int64(plan.PriceAmount), user.Group)
if payMoney < 0.01 {
c.JSON(http.StatusOK, gin.H{"message": "error", "data": "套餐金额过低"})
return
}
Comment on lines +82 to +86

@coderabbitai coderabbitai Bot Jun 18, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟠 Major | ⚡ Quick win

Same int64() truncation issue as epay handler

int64(plan.PriceAmount) truncates decimal prices (see comment on subscription_payment_epay.go line 92).

Additionally, the comparison < 0.01 differs from the existing topup_waffo_pancake.go which uses <= 0.01 (context snippet 2, line 45). For consistency within the same payment provider, consider aligning the threshold check.

🤖 Prompt for AI Agents 
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@controller/subscription_payment_waffo_pancake.go` around lines 82 - 86, In
the getWaffoPancakePayMoney function call, remove the int64() type cast on
plan.PriceAmount to prevent truncation of decimal values; pass plan.PriceAmount
directly to the function. Additionally, change the payment amount threshold
comparison from < 0.01 to <= 0.01 to align with the consistent threshold check
used in topup_waffo_pancake.go for the same payment provider.

order := &model.SubscriptionOrder{
UserId: userId,
PlanId: plan.Id,
Money: plan.PriceAmount,
Money: payMoney,
TradeNo: tradeNo,
PaymentMethod: model.PaymentMethodWaffoPancake,
PaymentProvider: model.PaymentProviderWaffoPancake,
Expand All @@ -100,7 +105,7 @@ func SubscriptionRequestWaffoPancakePay(c *gin.Context) {
ProductID: plan.WaffoPancakeProductId,
BuyerIdentity: service.WaffoPancakeBuyerIdentityFromUserID(user.Id),
PriceSnapshot: &service.WaffoPancakePriceSnapshot{
Amount: decimal.NewFromFloat(plan.PriceAmount).StringFixed(2),
Amount: decimal.NewFromFloat(payMoney).StringFixed(2),
TaxCategory: "saas",
},
BuyerEmail: getWaffoPancakeBuyerEmail(user),
Expand All @@ -114,7 +119,7 @@ func SubscriptionRequestWaffoPancakePay(c *gin.Context) {
c.JSON(http.StatusOK, gin.H{"message": "error", "data": "拉起支付失败"})
return
}
logger.LogInfo(c.Request.Context(), fmt.Sprintf("Waffo Pancake 订阅订单创建成功 user_id=%d plan_id=%d trade_no=%s session_id=%s money=%.2f", userId, plan.Id, tradeNo, session.SessionID, plan.PriceAmount))
logger.LogInfo(c.Request.Context(), fmt.Sprintf("Waffo Pancake 订阅订单创建成功 user_id=%d plan_id=%d trade_no=%s session_id=%s money=%.2f", userId, plan.Id, tradeNo, session.SessionID, payMoney))

c.JSON(http.StatusOK, gin.H{
"message": "success",
Expand Down
6 changes: 5 additions & 1 deletion relay/channel/claude/adaptor.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -42,7 +42,11 @@ func (a *Adaptor) Init(info *relaycommon.RelayInfo) {
}

func (a *Adaptor) GetRequestURL(info *relaycommon.RelayInfo) (string, error) {
requestURL := fmt.Sprintf("%s/v1/messages", info.ChannelBaseUrl)
requestPath := "/v1/messages"
if relaycommon.IsClaudeCountTokensRequest(info) {
requestPath = "/v1/messages/count_tokens"
}
requestURL := fmt.Sprintf("%s%s", info.ChannelBaseUrl, requestPath)
if !shouldAppendClaudeBetaQuery(info) {
return requestURL, nil
}
Expand Down
33 changes: 33 additions & 0 deletions relay/channel/claude/adaptor_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,33 @@
package claude

import (
"testing"

relaycommon "github.com/QuantumNous/new-api/relay/common"
relayconstant "github.com/QuantumNous/new-api/relay/constant"
"github.com/stretchr/testify/require"
)

func TestGetRequestURLUsesMessagesCountTokensPath(t *testing.T) {
adaptor := &Adaptor{}
info := &relaycommon.RelayInfo{
ChannelBaseUrl: "https://api.anthropic.com",
RelayMode: relayconstant.RelayModeClaudeCountTokens,
}
Comment on lines +14 to +16

@coderabbitai coderabbitai Bot Jun 18, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🔴 Critical

🧩 Analysis chain

🏁 Script executed:

#!/bin/bash
set -euo pipefail

echo "== RelayInfo and ChannelMeta declarations =="
rg -n --type=go 'type RelayInfo struct|type ChannelMeta struct|ChannelBaseUrl' relay/common/relay_info.go relay/common/channel_meta.go || true

echo
echo "== Invalid test literal usage =="
rg -nPU --type=go 'RelayInfo\s*{(?s).*?ChannelBaseUrl\s*:' relay/channel/claude/adaptor_test.go

Repository: QuantumNous/new-api

Length of output: 952

🏁 Script executed:

# View the RelayInfo and ChannelMeta struct definitions
sed -n '64,115p' relay/common/relay_info.go

Repository: QuantumNous/new-api

Length of output: 1580

🏁 Script executed:

# View the full test context around the problematic lines
sed -n '10,35p' relay/channel/claude/adaptor_test.go

Repository: QuantumNous/new-api

Length of output: 818

🏁 Script executed:

sed -n '89,120p' relay/common/relay_info.go

Repository: QuantumNous/new-api

Length of output: 1084

🏁 Script executed:

sed -n '120,150p' relay/common/relay_info.go

Repository: QuantumNous/new-api

Length of output: 1502

🏁 Script executed:

sed -n '150,180p' relay/common/relay_info.go

Repository: QuantumNous/new-api

Length of output: 1502

🏁 Script executed:

sed -n '180,210p' relay/common/relay_info.go

Repository: QuantumNous/new-api

Length of output: 1474

Fix invalid RelayInfo struct literals (compile blocker).

Lines 14 and 26 attempt to initialize ChannelBaseUrl directly on RelayInfo, but ChannelMeta is embedded as a pointer (*ChannelMeta), not a direct field. Struct literal syntax cannot set promoted fields of pointer-embedded types. Initialize ChannelMeta explicitly instead.

Proposed fix 
 func TestGetRequestURLUsesMessagesCountTokensPath(t *testing.T) {
 	adaptor := &Adaptor{}
 	info := &relaycommon.RelayInfo{
-		ChannelBaseUrl: "https://api.anthropic.com",
-		RelayMode:      relayconstant.RelayModeClaudeCountTokens,
+		ChannelMeta: &relaycommon.ChannelMeta{
+			ChannelBaseUrl: "https://api.anthropic.com",
+		},
+		RelayMode: relayconstant.RelayModeClaudeCountTokens,
 	}
@@
 func TestGetRequestURLKeepsMessagesPathForClaudeMessages(t *testing.T) {
 	adaptor := &Adaptor{}
 	info := &relaycommon.RelayInfo{
-		ChannelBaseUrl: "https://api.anthropic.com",
-		RequestURLPath:  "/v1/messages",
+		ChannelMeta: &relaycommon.ChannelMeta{
+			ChannelBaseUrl: "https://api.anthropic.com",
+		},
+		RequestURLPath: "/v1/messages",
 	}
🧰 Tools
🪛 golangci-lint (2.12.2)

[error] 14-14: : # github.com/QuantumNous/new-api/relay/channel/claude [github.com/QuantumNous/new-api/relay/channel/claude.test]
relay/channel/claude/adaptor_test.go:14:3: unknown field ChannelBaseUrl in struct literal of type "github.com/QuantumNous/new-api/relay/common".RelayInfo
relay/channel/claude/adaptor_test.go:26:3: unknown field ChannelBaseUrl in struct literal of type "github.com/QuantumNous/new-api/relay/common".RelayInfo

(typecheck)

🤖 Prompt for AI Agents 
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@relay/channel/claude/adaptor_test.go` around lines 14 - 16, The RelayInfo
struct literal is attempting to set ChannelBaseUrl directly, but since
ChannelMeta is embedded as a pointer (*ChannelMeta) rather than directly
embedded, Go's struct literal syntax cannot set its promoted fields this way.
Fix this by explicitly initializing the ChannelMeta field with a pointer to a
new ChannelMeta struct that contains the ChannelBaseUrl value. This applies to
both locations mentioned in the comment (around lines 14 and 26) where RelayInfo
struct literals are being constructed with ChannelBaseUrl and RelayMode fields.

Source: Linters/SAST tools


requestURL, err := adaptor.GetRequestURL(info)
require.NoError(t, err)
require.Equal(t, "https://api.anthropic.com/v1/messages/count_tokens", requestURL)
}

func TestGetRequestURLKeepsMessagesPathForClaudeMessages(t *testing.T) {
adaptor := &Adaptor{}
info := &relaycommon.RelayInfo{
ChannelBaseUrl: "https://api.anthropic.com",
RequestURLPath: "/v1/messages",
}

requestURL, err := adaptor.GetRequestURL(info)
require.NoError(t, err)
require.Equal(t, "https://api.anthropic.com/v1/messages", requestURL)
}
16 changes: 16 additions & 0 deletions relay/claude_handler.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -132,6 +132,12 @@ func ClaudeHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *typ
}
}

if relaycommon.IsClaudeCountTokensRequest(info) {
request.MaxTokens = nil
request.MaxTokensToSample = nil
request.Stream = nil
}

if !model_setting.GetGlobalSettings().PassThroughRequestEnabled &&
!info.ChannelSetting.PassThroughBodyEnabled &&
service.ShouldChatCompletionsUseResponsesGlobal(info.ChannelId, info.ChannelType, info.OriginModelName) {
Expand Down Expand Up @@ -211,6 +217,16 @@ func ClaudeHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *typ
}
}

if relaycommon.IsClaudeCountTokensRequest(info) {
respBody, err := io.ReadAll(httpResp.Body)
service.CloseResponseBodyGracefully(httpResp)
if err != nil {
return types.NewOpenAIError(err, types.ErrorCodeBadResponseBody, http.StatusInternalServerError)
}
service.IOCopyBytesGracefully(c, httpResp, respBody)
return nil
}

usage, newAPIError := adaptor.DoResponse(c, httpResp, info)
if newAPIError != nil {
// reset status code 重置状态码
Expand Down
19 changes: 19 additions & 0 deletions relay/common/relay_info.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@ import (
"encoding/json"
"errors"
"fmt"
"net/url"
"strconv"
"strings"
"time"
Expand Down Expand Up @@ -535,6 +536,24 @@ func cloneRequestHeaders(c *gin.Context) map[string]string {
return headers
}

func IsClaudeCountTokensRequestPath(path string) bool {
parsedPath := path
if parsedURL, err := url.Parse(path); err == nil && parsedURL.Path != "" {
parsedPath = parsedURL.Path
}
return parsedPath == "/v1/messages/count_tokens"
}

func IsClaudeCountTokensRequest(info *RelayInfo) bool {
if info == nil {
return false
}
if info.RelayMode == relayconstant.RelayModeClaudeCountTokens {
return true
}
return IsClaudeCountTokensRequestPath(info.RequestURLPath)
}

func GenRelayInfo(c *gin.Context, relayFormat types.RelayFormat, request dto.Request, ws *websocket.Conn) (*RelayInfo, error) {
var info *RelayInfo
var err error
Expand Down
8 changes: 8 additions & 0 deletions relay/common/relay_info_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@ package common
import (
"testing"

relayconstant "github.com/QuantumNous/new-api/relay/constant"
"github.com/QuantumNous/new-api/types"
"github.com/stretchr/testify/require"
)
Expand Down Expand Up @@ -38,3 +39,10 @@ func TestRelayInfoGetFinalRequestRelayFormatNilReceiver(t *testing.T) {
var info *RelayInfo
require.Equal(t, types.RelayFormat(""), info.GetFinalRequestRelayFormat())
}

func TestIsClaudeCountTokensRequest(t *testing.T) {
require.True(t, IsClaudeCountTokensRequest(&RelayInfo{RequestURLPath: "/v1/messages/count_tokens"}))
require.True(t, IsClaudeCountTokensRequest(&RelayInfo{RequestURLPath: "/v1/messages/count_tokens?beta=true"}))
require.True(t, IsClaudeCountTokensRequest(&RelayInfo{RelayMode: relayconstant.RelayModeClaudeCountTokens}))
require.False(t, IsClaudeCountTokensRequest(&RelayInfo{RequestURLPath: "/v1/messages"}))
}
4 changes: 4 additions & 0 deletions relay/constant/relay_mode.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -52,6 +52,8 @@ const (
RelayModeGemini

RelayModeResponsesCompact

RelayModeClaudeCountTokens
)

func Path2RelayMode(path string) int {
Expand Down Expand Up @@ -86,6 +88,8 @@ func Path2RelayMode(path string) int {
relayMode = RelayModeRerank
} else if strings.HasPrefix(path, "/v1/realtime") {
relayMode = RelayModeRealtime
} else if strings.HasPrefix(path, "/v1/messages/count_tokens") {
relayMode = RelayModeClaudeCountTokens
} else if strings.HasPrefix(path, "/v1beta/models") || strings.HasPrefix(path, "/v1/models") {
relayMode = RelayModeGemini
} else if strings.HasPrefix(path, "/mj") {
Expand Down
11 changes: 11 additions & 0 deletions relay/constant/relay_mode_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,11 @@
package constant

import (
"testing"

"github.com/stretchr/testify/require"
)

func TestPath2RelayModeClaudeCountTokens(t *testing.T) {
require.Equal(t, RelayModeClaudeCountTokens, Path2RelayMode("/v1/messages/count_tokens"))
}
3 changes: 3 additions & 0 deletions router/relay-router.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -85,6 +85,9 @@ func SetRelayRouter(router *gin.Engine) {
httpRouter.Use(middleware.Distribute())

// claude related routes
httpRouter.POST("/messages/count_tokens", func(c *gin.Context) {
controller.Relay(c, types.RelayFormatClaude)
})
httpRouter.POST("/messages", func(c *gin.Context) {
controller.Relay(c, types.RelayFormatClaude)
})
Expand Down
Loading