feat(agent): 按部署 agent 聚合模型用量,闭合 #9 用量聚合 (#22)

#9 的审计脱敏半部已由 PR#11 完成。据评审报告(L95 usage/billing 结构完整、
L112 监控 mock metrics 属 AM)核实:HM 侧通用用量聚合并不缺,真正缺的是
agent 维度用量视图。

- model.SumAgentUsage(userId, tokenName, start, end):按 agent 铸币 token 名
  'agent:<deployment_id>' 聚合 consume 日志(quota/prompt/completion/调用数),
  COALESCE+COUNT/SUM 跨 SQLite/MySQL/PG。
- GET /api/heicode/agents/:deployment_id/usage:返回该 agent 用量 + quota_per_unit
  (raw,调用方换算,与 /api/heicode/self 同契约),支持 ?start=&end= 时间窗。
- 测试 model/agent_usage_test.go:聚合正确、排除他人/非消费/不同 token、时间窗、空名。

go build/vet 干净;测试 PASS。

Refs #9 (审计脱敏半部 PR#11 已完成;监控 mock metrics 属 AM 侧)

Co-authored-by: chenchen <chenchen@xinghanlab.com>
Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
zsbgnw12
2026-06-09 14:27:48 +08:00
committed by GitHub
co-authored by chenchen Claude Opus 4.8
parent aa4041f604
commit af094dde27
4 changed files with 104 additions and 0 deletions
@@ -444,6 +444,33 @@ func HeicodeGetAgentStatus(c *gin.Context) {
})
}
// HeicodeGetAgentUsage: GET /api/heicode/agents/:deployment_id/usage
// Per-agent model-usage rollup (#9): the agent calls HM /v1/* with its minted
// token named "agent:<deployment_id>", so its consumption is the sum of consume
// logs under that token name. Optional ?start=&end= unix-second window. Returns
// raw quota + quota_per_unit (caller converts, same contract as /api/heicode/self).
func HeicodeGetAgentUsage(c *gin.Context) {
row, ok := findUserTemplateAgent(c)
if !ok {
return
}
start, _ := strconv.ParseInt(strings.TrimSpace(c.Query("start")), 10, 64)
end, _ := strconv.ParseInt(strings.TrimSpace(c.Query("end")), 10, 64)
usage, err := model.SumAgentUsage(c.GetInt("id"), "agent:"+row.DeploymentID, start, end)
if err != nil {
agentError(c, "DEPLOYMENT_CONFLICT", "failed to aggregate agent usage")
return
}
common.ApiSuccess(c, gin.H{
"agent_id": row.DeploymentID,
"quota": usage.Quota,
"prompt_tokens": usage.PromptTokens,
"completion_tokens": usage.CompletionTokens,
"call_count": usage.CallCount,
"quota_per_unit": common.QuotaPerUnit,
})
}
// HeicodeStopAgent: POST /api/heicode/agents/:deployment_id/stop
func HeicodeStopAgent(c *gin.Context) {
row, ok := findUserTemplateAgent(c)
+41
View File
@@ -0,0 +1,41 @@
package model
import (
"testing"
"github.com/stretchr/testify/require"
)
// SumAgentUsage rolls up a single agent token's consume logs (issue #9). logs is
// migrated by the package TestMain; LOG_DB == DB in tests.
func TestSumAgentUsage(t *testing.T) {
tok := "agent:dep_usage_test"
mk := func(uid int, tname string, typ, quota, pt, ct int) {
require.NoError(t, LOG_DB.Create(&Log{
UserId: uid, TokenName: tname, Type: typ,
Quota: quota, PromptTokens: pt, CompletionTokens: ct, CreatedAt: 1700000000,
}).Error)
}
mk(970001, tok, LogTypeConsume, 100, 10, 5)
mk(970001, tok, LogTypeConsume, 200, 20, 15)
mk(970001, "agent:other", LogTypeConsume, 999, 99, 99) // different token — excluded
mk(970001, tok, LogTypeManage, 500, 0, 0) // non-consume — excluded
mk(970002, tok, LogTypeConsume, 777, 7, 7) // different user — excluded by user filter
u, err := SumAgentUsage(970001, tok, 0, 0)
require.NoError(t, err)
require.EqualValues(t, 300, u.Quota)
require.EqualValues(t, 30, u.PromptTokens)
require.EqualValues(t, 20, u.CompletionTokens)
require.EqualValues(t, 2, u.CallCount)
// time window excludes the rows (all stamped at 1700000000)
windowed, err := SumAgentUsage(970001, tok, 1800000000, 1900000000)
require.NoError(t, err)
require.EqualValues(t, 0, windowed.CallCount)
// empty token name -> zero, no error
z, err := SumAgentUsage(970001, "", 0, 0)
require.NoError(t, err)
require.EqualValues(t, 0, z.CallCount)
}
+35
View File
@@ -432,6 +432,41 @@ type Stat struct {
Tpm int `json:"tpm"`
}
// AgentUsage is the per-deployed-agent usage rollup (issue #9): a template agent
// calls HM /v1/* with its own minted token named "agent:<deployment_id>", so its
// consumption is the sum of consume logs under that token name.
type AgentUsage struct {
Quota int64 `json:"quota"`
PromptTokens int64 `json:"prompt_tokens"`
CompletionTokens int64 `json:"completion_tokens"`
CallCount int64 `json:"call_count"`
}
// SumAgentUsage aggregates consume-log usage for a single agent token (and user,
// when userId>0) over an optional time window. COALESCE keeps SUM non-null on
// empty sets; COUNT/SUM are portable across SQLite/MySQL/PostgreSQL.
func SumAgentUsage(userId int, tokenName string, startTimestamp int64, endTimestamp int64) (AgentUsage, error) {
var u AgentUsage
if tokenName == "" {
return u, nil
}
tx := LOG_DB.Table("logs").
Select("COALESCE(SUM(quota),0) AS quota, COALESCE(SUM(prompt_tokens),0) AS prompt_tokens, COALESCE(SUM(completion_tokens),0) AS completion_tokens, COUNT(*) AS call_count").
Where("type = ?", LogTypeConsume).
Where("token_name = ?", tokenName)
if userId > 0 {
tx = tx.Where("user_id = ?", userId)
}
if startTimestamp != 0 {
tx = tx.Where("created_at >= ?", startTimestamp)
}
if endTimestamp != 0 {
tx = tx.Where("created_at <= ?", endTimestamp)
}
err := tx.Scan(&u).Error
return u, err
}
func SumUsedQuota(logType int, startTimestamp int64, endTimestamp int64, modelName string, username string, tokenName string, channel int, group string) (stat Stat, err error) {
tx := LOG_DB.Table("logs").Select("sum(quota) quota")
+1
View File
@@ -535,6 +535,7 @@ func SetApiRouter(router *gin.Engine) {
heicodeAgentRoute.GET("/agents", controller.HeicodeListAgents)
heicodeAgentRoute.GET("/agents/:deployment_id", controller.HeicodeGetAgent)
heicodeAgentRoute.GET("/agents/:deployment_id/status", controller.HeicodeGetAgentStatus)
heicodeAgentRoute.GET("/agents/:deployment_id/usage", controller.HeicodeGetAgentUsage)
heicodeAgentRoute.POST("/agents/:deployment_id/stop", controller.HeicodeStopAgent)
heicodeAgentRoute.DELETE("/agents/:deployment_id", controller.HeicodeDeleteAgent)
}