From af094dde279ccbe99df5e4fa30fcde21042601c9 Mon Sep 17 00:00:00 2001 From: zsbgnw12 <103713022+zsbgnw12@users.noreply.github.com> Date: Tue, 9 Jun 2026 14:27:48 +0800 Subject: [PATCH] =?UTF-8?q?feat(agent):=20=E6=8C=89=E9=83=A8=E7=BD=B2=20ag?= =?UTF-8?q?ent=20=E8=81=9A=E5=90=88=E6=A8=A1=E5=9E=8B=E7=94=A8=E9=87=8F,?= =?UTF-8?q?=E9=97=AD=E5=90=88=20#9=20=E7=94=A8=E9=87=8F=E8=81=9A=E5=90=88?= =?UTF-8?q?=20(#22)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit #9 的审计脱敏半部已由 PR#11 完成。据评审报告(L95 usage/billing 结构完整、 L112 监控 mock metrics 属 AM)核实:HM 侧通用用量聚合并不缺,真正缺的是 agent 维度用量视图。 - model.SumAgentUsage(userId, tokenName, start, end):按 agent 铸币 token 名 'agent:' 聚合 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 Co-authored-by: Claude Opus 4.8 --- heicode/controller/agent_template_handlers.go | 27 ++++++++++++ heicode/model/agent_usage_test.go | 41 +++++++++++++++++++ heicode/model/log.go | 35 ++++++++++++++++ heicode/router/api-router.go | 1 + 4 files changed, 104 insertions(+) create mode 100644 heicode/model/agent_usage_test.go diff --git a/heicode/controller/agent_template_handlers.go b/heicode/controller/agent_template_handlers.go index 8972fc7f..8d926097 100644 --- a/heicode/controller/agent_template_handlers.go +++ b/heicode/controller/agent_template_handlers.go @@ -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:", 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) diff --git a/heicode/model/agent_usage_test.go b/heicode/model/agent_usage_test.go new file mode 100644 index 00000000..af0c1165 --- /dev/null +++ b/heicode/model/agent_usage_test.go @@ -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) +} diff --git a/heicode/model/log.go b/heicode/model/log.go index 32693eea..da47471c 100644 --- a/heicode/model/log.go +++ b/heicode/model/log.go @@ -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:", 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") diff --git a/heicode/router/api-router.go b/heicode/router/api-router.go index bab035f4..8ac88abb 100644 --- a/heicode/router/api-router.go +++ b/heicode/router/api-router.go @@ -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) }