Browse Source
## 先修鉴权：此前公网接口完全没有身份校验
authFromRequest 只把 Authorization 塞进 context，**全仓没有一处读回来，也没有
任何 JWT 校验**；三个用户工具（get_user_tasks / allhelp_task / cancel_user_task）
的 uid 都是 mcp.WithString("uid", Required()) —— 由大模型填进来的普通参数。
而 mcp 是独立进程、独立端口 7300，不经过 gateway 那套 parseToken + isInWhiteList
（那只管 /api/* 和 /web/*），又必须对百炼公网可达。所以这不是「内网接口没做鉴权」，
是公网接口完全没有鉴权：知道 uid 就能读改任意人的数据。
新增 modules/mcp/auth.go：
- 从 Authorization 解 JWT 取 uid，与 gateway/core.go 同一套口径（RegisteredClaims，
uid 在 ID 字段）；兼容 "Bearer xxx" 前缀（gateway 收的是裸 token，标准 MCP
客户端会加前缀）。
- 参数里的 uid **只用于比对**，不一致直接拒。不静默改用会话 uid——调用方显然
误以为自己能指定用户，让它失败比让它以为成功了更安全。
- 解不出会话 uid 一律拒，**不回退到参数 uid**，那等于这层没做。
- TokenKey 未配置时拒绝所有用户数据类工具：配置漏了导致鉴权静默失效，
比工具不可用严重得多。TokenKey 必须与 gateway 同值，两边都读 GATEWAY_TOKEN_KEY。
存量三个工具同步改造，uid 参数保留但改为可选（兼容百炼后台已配好的工具定义）。
9 个测试守住这条线，其中三条是核心安全断言：参数 uid 与会话不符必须拒、
没有 token 必须拒而不是回退、未配 TokenKey 必须全拒。另有伪造签名、过期 token、
空 uid claim 的用例。
## 共享查询层：避免 home 与 mcp 各写一份
comm/memoryquery.go。mcp 从不 RpcCall、只裸查 MySQL，如果记忆项的过滤条件、
分类白名单、排序口径两边各写一遍，改了一边另一边不报错，只会让 EMAI 答的和
App 里显示的对不上——而且没人会发现。memory 模块的 listItems 也改走这里。
ApplyMemoryQuery 里 uid 为空时拼 "1 = 0" 而不是不加条件：少一个 uid 条件就是把
全库记忆项返回给调用方，而它有两个调用方，其中一个公网可达。
另提供 MemoryRelativeRange 把「本周/上月/最近7天」换算成日期区间给模型用——
模型对「今天几号」的认知来自对话上下文，很不可靠，让它自己算 start_date 经常
算错（SET_clock 把「下午 3:30」填成 03:30 是同一类问题）。
## 三个只读工具
get_memory_items / get_memory_stats / get_memory_report。
**只读是刻意的**：新增/修改/删除不放这里。写路径的校验、幂等（client_key）、
提醒时刻重算都在 home 的 memory 模块内，在 MCP 里复制一份必然漂移。
录入走端侧指令（tool_calls）、查询走 MCP，这个分工不变。
- 花销按币种分行返回，不做汇率换算；同时给「分」和「元」，模型念元不容易错，
留分是以免它想自己做加减时用浮点。
- get_memory_report 只返回已生成完的：还在生成中的 summary 是空的，
给出去只会让模型编一段话填空。
- 返回给模型的是精简字段，不是整行 DBMemoryItem——client_key/gen_round/extra
这些内部字段对回答问题没用，只占 token 还诱导模型去解释它们。
- date_certain=false 的项带上 date_is_inferred 标记，让模型别把推断出来的日期
当成用户确认过的安排来陈述。
- 模型填错分类时忽略该筛选当作「不限分类」，而不是整次查询失败——它多半会再填错一次。
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
main
Rodger-Wang
1 month ago
12 changed files with 848 additions and 35 deletions
@ -0,0 +1,134 @@ |
|||||
|
package comm |
||||
|
|
||||
|
import ( |
||||
|
"fmt" |
||||
|
"strings" |
||||
|
"time" |
||||
|
|
||||
|
"gorm.io/gorm" |
||||
|
) |
||||
|
|
||||
|
/* |
||||
|
记忆项的**共享查询层**。 |
||||
|
|
||||
|
放在 comm 是为了解决一个具体问题：`mcp` 是独立进程、独立服务，它的工具从不 |
||||
|
RpcCall、只裸查 MySQL（见 tool_get_user_tasks 的 |
||||
|
`mysql.Table(...).Where("uid = ?", uid)`）。如果 memory 的过滤条件、分类白名单、 |
||||
|
排序口径在 home 和 mcp 里各写一遍，改了一边另一边不会报错，只会让 EMAI 答的 |
||||
|
和 App 里显示的对不上——而且没人会发现。 |
||||
|
|
||||
|
所以两边都调这里。**新增过滤条件只改这一个文件。** |
||||
|
|
||||
|
⚠️ 这里只放「读」。写操作（新增/改/删）仍只在 home 的 memory 模块里， |
||||
|
MCP 侧的写工具经它的 HTTP 接口不了——那是刻意的：写路径的校验、幂等、 |
||||
|
提醒重算都在 memory 模块内，复制一份必然漂移。 |
||||
|
*/ |
||||
|
|
||||
|
// MemoryQuery 记忆项查询条件。零值表示不限制。
|
||||
|
type MemoryQuery struct { |
||||
|
Uid string |
||||
|
StartDate string // YYYY-MM-DD，含
|
||||
|
EndDate string // YYYY-MM-DD，含
|
||||
|
Categories []string |
||||
|
States []int32 |
||||
|
Limit int |
||||
|
Offset int |
||||
|
} |
||||
|
|
||||
|
// ApplyMemoryQuery 把条件套到一个已经 Table(memory_item) 过的 *gorm.DB 上。
|
||||
|
//
|
||||
|
// ⚠️ **uid 一定会被拼进去**，即使 q.Uid 为空也会拼一个不可能命中的条件——
|
||||
|
// 少一个 uid 条件就是把全库的记忆项返回给调用方，而这个函数有两个调用方，
|
||||
|
// 其中一个（MCP）是公网可达的。宁可返回空也不要返回别人的数据。
|
||||
|
func ApplyMemoryQuery(tx *gorm.DB, q MemoryQuery) *gorm.DB { |
||||
|
if strings.TrimSpace(q.Uid) == "" { |
||||
|
return tx.Where("1 = 0") |
||||
|
} |
||||
|
tx = tx.Where("uid = ?", q.Uid) |
||||
|
if q.StartDate != "" { |
||||
|
tx = tx.Where("happen_date >= ?", q.StartDate) |
||||
|
} |
||||
|
if q.EndDate != "" { |
||||
|
tx = tx.Where("happen_date <= ?", q.EndDate) |
||||
|
} |
||||
|
if len(q.Categories) > 0 { |
||||
|
tx = tx.Where("category in ?", q.Categories) |
||||
|
} |
||||
|
if len(q.States) > 0 { |
||||
|
tx = tx.Where("state in ?", q.States) |
||||
|
} |
||||
|
return tx |
||||
|
} |
||||
|
|
||||
|
// OrderMemoryItems 统一排序口径：先日期，再时刻（无时刻的排当天最后），最后 id。
|
||||
|
//
|
||||
|
// 两边必须一致，否则「App 里第一条」和「EMAI 说的第一条」会是不同的东西。
|
||||
|
func OrderMemoryItems(tx *gorm.DB) *gorm.DB { |
||||
|
return tx.Order("happen_date asc"). |
||||
|
Order("happen_time = '' asc"). |
||||
|
Order("happen_time asc"). |
||||
|
Order("id asc") |
||||
|
} |
||||
|
|
||||
|
// NormalizeMemoryCategories 过滤掉不在白名单里的分类。
|
||||
|
//
|
||||
|
// 返回 (合法分类, 第一个非法分类)。全部合法时第二个返回值为空串。
|
||||
|
// 调用方决定是拒绝还是忽略——MCP 侧倾向忽略（模型填错分类不该让整次查询失败），
|
||||
|
// HTTP 接口侧倾向拒绝（客户端填错是 bug，早点暴露）。
|
||||
|
func NormalizeMemoryCategories(in []string) ([]string, string) { |
||||
|
out := make([]string, 0, len(in)) |
||||
|
bad := "" |
||||
|
for _, c := range in { |
||||
|
c = strings.ToLower(strings.TrimSpace(c)) |
||||
|
if c == "" { |
||||
|
continue |
||||
|
} |
||||
|
if !IsValidMemoryCategory(c) { |
||||
|
if bad == "" { |
||||
|
bad = c |
||||
|
} |
||||
|
continue |
||||
|
} |
||||
|
out = append(out, c) |
||||
|
} |
||||
|
return out, bad |
||||
|
} |
||||
|
|
||||
|
// MemoryRelativeRange 把「最近 N 天 / 本周 / 本月」这类相对说法换算成日期区间。
|
||||
|
//
|
||||
|
// 给 MCP 工具用：模型问「我这个月花了多少」时不该自己算日期，
|
||||
|
// 它对「今天几号」的认知来自对话上下文，很不可靠。
|
||||
|
func MemoryRelativeRange(kind string, now time.Time) (start, end string, err error) { |
||||
|
switch strings.ToLower(strings.TrimSpace(kind)) { |
||||
|
case "today": |
||||
|
d := FormatMemoryDate(now) |
||||
|
return d, d, nil |
||||
|
case "yesterday": |
||||
|
d := FormatMemoryDate(now.AddDate(0, 0, -1)) |
||||
|
return d, d, nil |
||||
|
case "this_week": |
||||
|
s, e := MemoryWeekRange(now) |
||||
|
return FormatMemoryDate(s), FormatMemoryDate(e), nil |
||||
|
case "last_week": |
||||
|
s, e := MemoryWeekRange(now.AddDate(0, 0, -7)) |
||||
|
return FormatMemoryDate(s), FormatMemoryDate(e), nil |
||||
|
case "this_month": |
||||
|
s, e := MemoryMonthRange(now) |
||||
|
return FormatMemoryDate(s), FormatMemoryDate(e), nil |
||||
|
case "last_month": |
||||
|
prev := time.Date(now.Year(), now.Month(), 1, 0, 0, 0, 0, now.Location()).AddDate(0, 0, -1) |
||||
|
s, e := MemoryMonthRange(prev) |
||||
|
return FormatMemoryDate(s), FormatMemoryDate(e), nil |
||||
|
case "next_7_days": |
||||
|
return FormatMemoryDate(now), FormatMemoryDate(now.AddDate(0, 0, 7)), nil |
||||
|
case "last_7_days": |
||||
|
return FormatMemoryDate(now.AddDate(0, 0, -7)), FormatMemoryDate(now), nil |
||||
|
} |
||||
|
return "", "", fmt.Errorf("不支持的时间范围: %s", kind) |
||||
|
} |
||||
|
|
||||
|
// MemoryRelativeRanges 支持的相对范围，给工具描述用
|
||||
|
var MemoryRelativeRanges = []string{ |
||||
|
"today", "yesterday", "this_week", "last_week", |
||||
|
"this_month", "last_month", "next_7_days", "last_7_days", |
||||
|
} |
||||
@ -0,0 +1,64 @@ |
|||||
|
package comm |
||||
|
|
||||
|
import ( |
||||
|
"testing" |
||||
|
"time" |
||||
|
) |
||||
|
|
||||
|
func TestMemoryRelativeRange(t *testing.T) { |
||||
|
loc := time.UTC |
||||
|
// 2026-09-05 是周六
|
||||
|
now := time.Date(2026, 9, 5, 12, 0, 0, 0, loc) |
||||
|
|
||||
|
cases := []struct{ kind, start, end string }{ |
||||
|
{"today", "2026-09-05", "2026-09-05"}, |
||||
|
{"yesterday", "2026-09-04", "2026-09-04"}, |
||||
|
// ISO 周：周一起，所以本周是 08-31(周一) ~ 09-06(周日)
|
||||
|
{"this_week", "2026-08-31", "2026-09-06"}, |
||||
|
{"last_week", "2026-08-24", "2026-08-30"}, |
||||
|
{"this_month", "2026-09-01", "2026-09-30"}, |
||||
|
{"last_month", "2026-08-01", "2026-08-31"}, |
||||
|
{"next_7_days", "2026-09-05", "2026-09-12"}, |
||||
|
{"last_7_days", "2026-08-29", "2026-09-05"}, |
||||
|
} |
||||
|
for _, c := range cases { |
||||
|
s, e, err := MemoryRelativeRange(c.kind, now) |
||||
|
if err != nil { |
||||
|
t.Errorf("%s: 报错 %v", c.kind, err) |
||||
|
continue |
||||
|
} |
||||
|
if s != c.start || e != c.end { |
||||
|
t.Errorf("%s: 期望 %s~%s 得到 %s~%s", c.kind, c.start, c.end, s, e) |
||||
|
} |
||||
|
} |
||||
|
} |
||||
|
|
||||
|
func TestMemoryRelativeRange_Unknown(t *testing.T) { |
||||
|
if _, _, err := MemoryRelativeRange("last_fortnight", time.Now()); err == nil { |
||||
|
t.Error("不认识的范围必须报错，不能静默返回一个区间") |
||||
|
} |
||||
|
} |
||||
|
|
||||
|
func TestNormalizeMemoryCategories(t *testing.T) { |
||||
|
got, bad := NormalizeMemoryCategories([]string{"todo", "TODO", " idea ", "nonsense", ""}) |
||||
|
if bad != "nonsense" { |
||||
|
t.Errorf("应报出第一个非法分类 nonsense，得到 %q", bad) |
||||
|
} |
||||
|
// 大小写与空格都要归一
|
||||
|
want := []string{"todo", "todo", "idea"} |
||||
|
if len(got) != len(want) { |
||||
|
t.Fatalf("期望 %d 个合法分类，得到 %d: %v", len(want), len(got), got) |
||||
|
} |
||||
|
for i := range want { |
||||
|
if got[i] != want[i] { |
||||
|
t.Errorf("第 %d 个：期望 %s 得到 %s", i, want[i], got[i]) |
||||
|
} |
||||
|
} |
||||
|
} |
||||
|
|
||||
|
func TestNormalizeMemoryCategories_AllValid(t *testing.T) { |
||||
|
_, bad := NormalizeMemoryCategories([]string{"todo", "alarm", "idea", "expense"}) |
||||
|
if bad != "" { |
||||
|
t.Errorf("全部合法时不该报错，却得到 %q", bad) |
||||
|
} |
||||
|
} |
||||
@ -0,0 +1,105 @@ |
|||||
|
package mcp |
||||
|
|
||||
|
import ( |
||||
|
"context" |
||||
|
"fmt" |
||||
|
"strings" |
||||
|
|
||||
|
"github.com/golang-jwt/jwt/v5" |
||||
|
) |
||||
|
|
||||
|
/* |
||||
|
MCP 的会话身份。 |
||||
|
|
||||
|
## 这块此前是空的 |
||||
|
|
||||
|
`authFromRequest` 只把 Authorization 头塞进 context，**全仓没有一处读回来， |
||||
|
也没有任何校验**；三个用户工具（get_user_tasks / allhelp_task / cancel_user_task） |
||||
|
的 uid 都是 `mcp.WithString("uid", Required())` —— 由大模型填进来的普通参数。 |
||||
|
|
||||
|
而 MCP 是**独立进程、独立端口（7300）**，不经过 gateway 那套 parseToken + |
||||
|
isInWhiteList（那只管 /api/* 和 /web/*），又必须对百炼公网可达。 |
||||
|
也就是说这不是「内网接口没做鉴权」，是**公网接口完全没有鉴权**： |
||||
|
知道 uid 就能读改任意人的数据。 |
||||
|
|
||||
|
记忆中心的条目装的是花销金额、灵感原文、每日行程，比「提醒任务」敏感得多， |
||||
|
所以在挂 memory 工具之前先把这层补上。 |
||||
|
|
||||
|
## 口径 |
||||
|
|
||||
|
- 有 Authorization 且能解出 uid → 用**会话 uid**，忽略参数里的 uid； |
||||
|
- 参数里的 uid 与会话 uid 不一致 → 直接拒（不是静默改用会话 uid： |
||||
|
调用方显然误以为自己能指定用户，让它失败比让它以为成功了更安全）； |
||||
|
- 解不出会话 uid → 拒。**不降级回退到参数 uid** —— 那等于这层没做。 |
||||
|
|
||||
|
⚠️ TokenKey 必须与 gateway 的 `GATEWAY_TOKEN_KEY` 一致，否则所有 token 都验不过。 |
||||
|
两边都从同一个 .env 变量注入。 |
||||
|
*/ |
||||
|
|
||||
|
// authKey is a custom context key for storing the auth token.
|
||||
|
type authKey struct{} |
||||
|
|
||||
|
// withAuthKey adds an auth key to the context.
|
||||
|
func withAuthKey(ctx context.Context, auth string) context.Context { |
||||
|
return context.WithValue(ctx, authKey{}, auth) |
||||
|
} |
||||
|
|
||||
|
// tokenFromContext 取出请求带来的 Authorization 原始串
|
||||
|
func tokenFromContext(ctx context.Context) string { |
||||
|
v, _ := ctx.Value(authKey{}).(string) |
||||
|
return v |
||||
|
} |
||||
|
|
||||
|
// parseUID 从 JWT 里解出 uid。与 gateway/core.go 的 parseToken 同一套口径
|
||||
|
// （jwt.RegisteredClaims，uid 放在 ID 字段）。
|
||||
|
func parseUID(tokenString string, secretKey []byte) (string, error) { |
||||
|
tokenString = strings.TrimSpace(tokenString) |
||||
|
if tokenString == "" { |
||||
|
return "", fmt.Errorf("缺少 Authorization") |
||||
|
} |
||||
|
// 兼容 "Bearer xxx" 写法：gateway 那边收到的就是裸 token，
|
||||
|
// 但 MCP 客户端（百炼）常按标准加 Bearer 前缀。
|
||||
|
if len(tokenString) > 7 && strings.EqualFold(tokenString[:7], "bearer ") { |
||||
|
tokenString = strings.TrimSpace(tokenString[7:]) |
||||
|
} |
||||
|
parsed, err := jwt.ParseWithClaims(tokenString, &jwt.RegisteredClaims{}, |
||||
|
func(token *jwt.Token) (interface{}, error) { |
||||
|
if _, ok := token.Method.(*jwt.SigningMethodHMAC); !ok { |
||||
|
return nil, fmt.Errorf("unexpected signing method: %v", token.Header["alg"]) |
||||
|
} |
||||
|
return secretKey, nil |
||||
|
}) |
||||
|
if err != nil { |
||||
|
return "", err |
||||
|
} |
||||
|
claims, ok := parsed.Claims.(*jwt.RegisteredClaims) |
||||
|
if !ok || !parsed.Valid { |
||||
|
return "", fmt.Errorf("invalid token") |
||||
|
} |
||||
|
if claims.ID == "" { |
||||
|
return "", fmt.Errorf("token 里没有 uid") |
||||
|
} |
||||
|
return claims.ID, nil |
||||
|
} |
||||
|
|
||||
|
// ResolveUID 决定这次工具调用该用哪个 uid。
|
||||
|
//
|
||||
|
// argUID 是模型传上来的参数，只用于「与会话 uid 比对」，绝不作为回退值。
|
||||
|
func (this *Mcp) ResolveUID(ctx context.Context, argUID string) (string, error) { |
||||
|
key := this.options.TokenKey |
||||
|
if key == "" { |
||||
|
// 没配 TokenKey 说明部署时漏了 GATEWAY_TOKEN_KEY。
|
||||
|
// 这时**拒绝所有用户类工具**而不是放行：配置缺失导致鉴权静默失效
|
||||
|
// 是最糟的一种失败方式。
|
||||
|
return "", fmt.Errorf("服务端未配置 TokenKey，用户数据类工具不可用") |
||||
|
} |
||||
|
uid, err := parseUID(tokenFromContext(ctx), []byte(key)) |
||||
|
if err != nil { |
||||
|
return "", fmt.Errorf("身份校验失败: %v", err) |
||||
|
} |
||||
|
argUID = strings.TrimSpace(argUID) |
||||
|
if argUID != "" && argUID != uid { |
||||
|
return "", fmt.Errorf("uid 与当前会话不符，已拒绝") |
||||
|
} |
||||
|
return uid, nil |
||||
|
} |
||||
@ -0,0 +1,121 @@ |
|||||
|
package mcp |
||||
|
|
||||
|
import ( |
||||
|
"context" |
||||
|
"testing" |
||||
|
"time" |
||||
|
|
||||
|
"github.com/golang-jwt/jwt/v5" |
||||
|
) |
||||
|
|
||||
|
const testKey = "test-token-key-please-ignore" |
||||
|
|
||||
|
func mkToken(t *testing.T, uid string, key string, expired bool) string { |
||||
|
t.Helper() |
||||
|
exp := time.Now().Add(time.Hour) |
||||
|
if expired { |
||||
|
exp = time.Now().Add(-time.Hour) |
||||
|
} |
||||
|
tok := jwt.NewWithClaims(jwt.SigningMethodHS256, &jwt.RegisteredClaims{ |
||||
|
ID: uid, |
||||
|
ExpiresAt: jwt.NewNumericDate(exp), |
||||
|
}) |
||||
|
s, err := tok.SignedString([]byte(key)) |
||||
|
if err != nil { |
||||
|
t.Fatalf("签 token 失败: %v", err) |
||||
|
} |
||||
|
return s |
||||
|
} |
||||
|
|
||||
|
func mcpWith(key string) *Mcp { |
||||
|
return &Mcp{options: &Options{TokenKey: key}} |
||||
|
} |
||||
|
|
||||
|
// 正常路径：能从 token 解出 uid
|
||||
|
func TestResolveUID_FromToken(t *testing.T) { |
||||
|
m := mcpWith(testKey) |
||||
|
ctx := withAuthKey(context.Background(), mkToken(t, "u123", testKey, false)) |
||||
|
uid, err := m.ResolveUID(ctx, "") |
||||
|
if err != nil || uid != "u123" { |
||||
|
t.Fatalf("期望 u123，得到 %q err=%v", uid, err) |
||||
|
} |
||||
|
} |
||||
|
|
||||
|
// Bearer 前缀要能吃下：gateway 收的是裸 token，但标准 MCP 客户端会加前缀
|
||||
|
func TestResolveUID_BearerPrefix(t *testing.T) { |
||||
|
m := mcpWith(testKey) |
||||
|
ctx := withAuthKey(context.Background(), "Bearer "+mkToken(t, "u123", testKey, false)) |
||||
|
uid, err := m.ResolveUID(ctx, "") |
||||
|
if err != nil || uid != "u123" { |
||||
|
t.Fatalf("Bearer 前缀应被剥掉，得到 %q err=%v", uid, err) |
||||
|
} |
||||
|
} |
||||
|
|
||||
|
// **核心安全断言**：参数里的 uid 与会话不符必须拒绝，
|
||||
|
// 绝不能静默改用会话 uid（那会掩盖调用方的错误认知），更不能采信参数。
|
||||
|
func TestResolveUID_RejectsMismatchedArgUID(t *testing.T) { |
||||
|
m := mcpWith(testKey) |
||||
|
ctx := withAuthKey(context.Background(), mkToken(t, "u123", testKey, false)) |
||||
|
if _, err := m.ResolveUID(ctx, "someone_else"); err == nil { |
||||
|
t.Fatal("参数 uid 与会话不符时必须报错") |
||||
|
} |
||||
|
} |
||||
|
|
||||
|
// 参数 uid 与会话一致时放行（百炼后台已配好的工具定义仍会传 uid）
|
||||
|
func TestResolveUID_AllowsMatchingArgUID(t *testing.T) { |
||||
|
m := mcpWith(testKey) |
||||
|
ctx := withAuthKey(context.Background(), mkToken(t, "u123", testKey, false)) |
||||
|
uid, err := m.ResolveUID(ctx, "u123") |
||||
|
if err != nil || uid != "u123" { |
||||
|
t.Fatalf("一致的 uid 应放行，得到 %q err=%v", uid, err) |
||||
|
} |
||||
|
} |
||||
|
|
||||
|
// **核心安全断言**：没有 token 时必须拒绝，
|
||||
|
// 绝不能回退成「用参数里的 uid」——那等于这层鉴权没做。
|
||||
|
func TestResolveUID_NoTokenIsRejectedNotFallback(t *testing.T) { |
||||
|
m := mcpWith(testKey) |
||||
|
ctx := context.Background() |
||||
|
if uid, err := m.ResolveUID(ctx, "u123"); err == nil { |
||||
|
t.Fatalf("没有 token 时必须拒绝，却返回了 uid=%q", uid) |
||||
|
} |
||||
|
} |
||||
|
|
||||
|
// 签名不对（伪造 token）必须拒绝
|
||||
|
func TestResolveUID_RejectsWrongSignature(t *testing.T) { |
||||
|
m := mcpWith(testKey) |
||||
|
ctx := withAuthKey(context.Background(), mkToken(t, "u123", "another-key", false)) |
||||
|
if _, err := m.ResolveUID(ctx, ""); err == nil { |
||||
|
t.Fatal("签名不匹配的 token 必须拒绝") |
||||
|
} |
||||
|
} |
||||
|
|
||||
|
// 过期 token 必须拒绝
|
||||
|
func TestResolveUID_RejectsExpired(t *testing.T) { |
||||
|
m := mcpWith(testKey) |
||||
|
ctx := withAuthKey(context.Background(), mkToken(t, "u123", testKey, true)) |
||||
|
if _, err := m.ResolveUID(ctx, ""); err == nil { |
||||
|
t.Fatal("过期 token 必须拒绝") |
||||
|
} |
||||
|
} |
||||
|
|
||||
|
// **核心安全断言**：服务端没配 TokenKey 时一律拒绝，不是放行。
|
||||
|
// 配置漏了导致鉴权静默失效，比工具不可用严重得多。
|
||||
|
func TestResolveUID_MissingTokenKeyDeniesAll(t *testing.T) { |
||||
|
m := mcpWith("") |
||||
|
ctx := withAuthKey(context.Background(), mkToken(t, "u123", testKey, false)) |
||||
|
if _, err := m.ResolveUID(ctx, ""); err == nil { |
||||
|
t.Fatal("未配置 TokenKey 时必须拒绝所有用户数据类工具") |
||||
|
} |
||||
|
} |
||||
|
|
||||
|
// token 里没有 uid（ID 为空）时不能返回空串当成合法身份 ——
|
||||
|
// 空 uid 会让下游的 ApplyMemoryQuery 走 "1 = 0" 分支返回空，
|
||||
|
// 但那是最后一道防线，不该指望它。
|
||||
|
func TestResolveUID_RejectsEmptyIDClaim(t *testing.T) { |
||||
|
m := mcpWith(testKey) |
||||
|
ctx := withAuthKey(context.Background(), mkToken(t, "", testKey, false)) |
||||
|
if uid, err := m.ResolveUID(ctx, ""); err == nil { |
||||
|
t.Fatalf("token 里没有 uid 时必须拒绝，却返回 %q", uid) |
||||
|
} |
||||
|
} |
||||
@ -0,0 +1,367 @@ |
|||||
|
package mcp |
||||
|
|
||||
|
import ( |
||||
|
"context" |
||||
|
"strings" |
||||
|
"time" |
||||
|
|
||||
|
"yunyan/comm" |
||||
|
"yunyan/lego/core" |
||||
|
"yunyan/lego/core/cbase" |
||||
|
"yunyan/lego/sys/mysql" |
||||
|
"yunyan/pb" |
||||
|
"yunyan/utils" |
||||
|
|
||||
|
"github.com/mark3labs/mcp-go/mcp" |
||||
|
) |
||||
|
|
||||
|
/* |
||||
|
记忆中心（拾忆）的 MCP 工具：让 EMAI 能回答「我明天有什么事」「这个月花了多少」。 |
||||
|
|
||||
|
## 只读 |
||||
|
|
||||
|
新增/修改/删除**不放在这里**。写路径的校验、幂等（client_key）、提醒时刻重算 |
||||
|
都在 home 的 memory 模块内，在 MCP 里复制一份必然漂移。 |
||||
|
录入走端侧指令（tool_calls），查询走 MCP —— 这个分工是刻意的。 |
||||
|
|
||||
|
## 身份 |
||||
|
|
||||
|
uid 一律从 Authorization 的 JWT 解（见 auth.go），**不信任模型传的参数**。 |
||||
|
*/ |
||||
|
|
||||
|
// ---------- get_memory_items ----------
|
||||
|
|
||||
|
type tool_get_memory_items struct { |
||||
|
cbase.ModuleCompBase |
||||
|
module *Mcp |
||||
|
} |
||||
|
|
||||
|
func (this *tool_get_memory_items) Init(service core.IService, module core.IModule, comp core.IModuleComp, opt core.IModuleOptions) (err error) { |
||||
|
this.ModuleCompBase.Init(service, module, comp, opt) |
||||
|
this.module = module.(*Mcp) |
||||
|
return |
||||
|
} |
||||
|
|
||||
|
func (this *tool_get_memory_items) Start() (err error) { |
||||
|
err = this.ModuleCompBase.Start() |
||||
|
this.module.AddTool(ToolGroup_GLOBAL, this.Tool(), this.Handl) |
||||
|
return |
||||
|
} |
||||
|
|
||||
|
func (this *tool_get_memory_items) Tool() mcp.Tool { |
||||
|
return mcp.NewTool("get_memory_items", |
||||
|
mcp.WithDescription("查询用户记录的事项（待办/闹钟/灵感/花销），可按时间范围与分类筛选"), |
||||
|
mcp.WithString("range", |
||||
|
mcp.Description("相对时间范围，取值："+strings.Join(comm.MemoryRelativeRanges, "/")+ |
||||
|
"。与 start_date/end_date 二选一，优先用本字段"), |
||||
|
), |
||||
|
mcp.WithString("start_date", mcp.Description("起始日期 YYYY-MM-DD")), |
||||
|
mcp.WithString("end_date", mcp.Description("结束日期 YYYY-MM-DD")), |
||||
|
mcp.WithString("category", |
||||
|
mcp.Description("分类筛选：todo=待办 alarm=闹钟 idea=灵感 expense=花销，留空=全部"), |
||||
|
), |
||||
|
mcp.WithNumber("limit", mcp.Description("返回条数上限（默认50，最大200）"), mcp.DefaultNumber(50)), |
||||
|
) |
||||
|
} |
||||
|
|
||||
|
func (this *tool_get_memory_items) Handl(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) { |
||||
|
uid, err := this.module.ResolveUID(ctx, request.GetString("uid", "")) |
||||
|
if err != nil { |
||||
|
return mcp.NewToolResultError(err.Error()), nil |
||||
|
} |
||||
|
start, end, errRange := resolveRange(request) |
||||
|
if errRange != nil { |
||||
|
return mcp.NewToolResultError(errRange.Error()), nil |
||||
|
} |
||||
|
|
||||
|
var cats []string |
||||
|
if c := strings.TrimSpace(request.GetString("category", "")); c != "" { |
||||
|
// 模型填错分类不该让整次查询失败——忽略非法值当作"不限分类"，
|
||||
|
// 比返回一个错误让它重试要好（它多半会再填错一次）。
|
||||
|
cats, _ = comm.NormalizeMemoryCategories([]string{c}) |
||||
|
} |
||||
|
|
||||
|
limit := int(request.GetInt("limit", 50)) |
||||
|
if limit <= 0 || limit > 200 { |
||||
|
limit = 50 |
||||
|
} |
||||
|
|
||||
|
items := make([]*pb.DBMemoryItem, 0) |
||||
|
tx := comm.ApplyMemoryQuery(mysql.Table(comm.TableMemoryItem), comm.MemoryQuery{ |
||||
|
Uid: uid, StartDate: start, EndDate: end, Categories: cats, |
||||
|
}) |
||||
|
if e := comm.OrderMemoryItems(tx).Limit(limit).Find(&items).Error; e != nil { |
||||
|
this.module.Errorln("get_memory_items error:", e) |
||||
|
return mcp.NewToolResultError("查询失败"), nil |
||||
|
} |
||||
|
|
||||
|
return mcp.NewToolResultText(utils.ToString(map[string]interface{}{ |
||||
|
"result": "success", |
||||
|
"start_date": start, |
||||
|
"end_date": end, |
||||
|
"total": len(items), |
||||
|
"items": briefItems(items), |
||||
|
})), nil |
||||
|
} |
||||
|
|
||||
|
// ---------- get_memory_stats ----------
|
||||
|
|
||||
|
type tool_get_memory_stats struct { |
||||
|
cbase.ModuleCompBase |
||||
|
module *Mcp |
||||
|
} |
||||
|
|
||||
|
func (this *tool_get_memory_stats) Init(service core.IService, module core.IModule, comp core.IModuleComp, opt core.IModuleOptions) (err error) { |
||||
|
this.ModuleCompBase.Init(service, module, comp, opt) |
||||
|
this.module = module.(*Mcp) |
||||
|
return |
||||
|
} |
||||
|
|
||||
|
func (this *tool_get_memory_stats) Start() (err error) { |
||||
|
err = this.ModuleCompBase.Start() |
||||
|
this.module.AddTool(ToolGroup_GLOBAL, this.Tool(), this.Handl) |
||||
|
return |
||||
|
} |
||||
|
|
||||
|
func (this *tool_get_memory_stats) Tool() mcp.Tool { |
||||
|
return mcp.NewTool("get_memory_stats", |
||||
|
mcp.WithDescription("统计用户某段时间的待办完成情况、灵感数量与花销合计"), |
||||
|
mcp.WithString("range", |
||||
|
mcp.Description("相对时间范围，取值："+strings.Join(comm.MemoryRelativeRanges, "/")), |
||||
|
), |
||||
|
mcp.WithString("start_date", mcp.Description("起始日期 YYYY-MM-DD")), |
||||
|
mcp.WithString("end_date", mcp.Description("结束日期 YYYY-MM-DD")), |
||||
|
) |
||||
|
} |
||||
|
|
||||
|
func (this *tool_get_memory_stats) Handl(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) { |
||||
|
uid, err := this.module.ResolveUID(ctx, request.GetString("uid", "")) |
||||
|
if err != nil { |
||||
|
return mcp.NewToolResultError(err.Error()), nil |
||||
|
} |
||||
|
start, end, errRange := resolveRange(request) |
||||
|
if errRange != nil { |
||||
|
return mcp.NewToolResultError(errRange.Error()), nil |
||||
|
} |
||||
|
|
||||
|
type catCount struct { |
||||
|
Category string |
||||
|
State int32 |
||||
|
N int64 |
||||
|
} |
||||
|
rows := make([]catCount, 0) |
||||
|
if e := comm.ApplyMemoryQuery(mysql.Table(comm.TableMemoryItem), |
||||
|
comm.MemoryQuery{Uid: uid, StartDate: start, EndDate: end}). |
||||
|
Select("category, state, count(*) as n").Group("category, state"). |
||||
|
Find(&rows).Error; e != nil { |
||||
|
this.module.Errorln("get_memory_stats error:", e) |
||||
|
return mcp.NewToolResultError("统计失败"), nil |
||||
|
} |
||||
|
|
||||
|
var todoTotal, todoDone, alarmTotal, ideaCount int64 |
||||
|
for _, r := range rows { |
||||
|
switch r.Category { |
||||
|
case comm.MemoryCatTodo: |
||||
|
todoTotal += r.N |
||||
|
if r.State == int32(pb.MemoryState_MemoryState_Done) { |
||||
|
todoDone += r.N |
||||
|
} |
||||
|
case comm.MemoryCatAlarm: |
||||
|
alarmTotal += r.N |
||||
|
case comm.MemoryCatIdea: |
||||
|
ideaCount += r.N |
||||
|
} |
||||
|
} |
||||
|
|
||||
|
// 花销**按币种分行**，不做汇率换算：把 CNY 和 JPY 加成一个数是错的。
|
||||
|
type expRow struct { |
||||
|
Currency string |
||||
|
TotalCents int64 |
||||
|
Count int64 |
||||
|
} |
||||
|
exps := make([]expRow, 0) |
||||
|
if e := comm.ApplyMemoryQuery(mysql.Table(comm.TableMemoryItem), |
||||
|
comm.MemoryQuery{Uid: uid, StartDate: start, EndDate: end, |
||||
|
Categories: []string{comm.MemoryCatExpense}}). |
||||
|
Select("currency, sum(amount_cents) as total_cents, count(*) as count"). |
||||
|
Group("currency").Find(&exps).Error; e != nil { |
||||
|
this.module.Errorln("get_memory_stats expense error:", e) |
||||
|
return mcp.NewToolResultError("统计失败"), nil |
||||
|
} |
||||
|
expOut := make([]map[string]interface{}, 0, len(exps)) |
||||
|
for _, e := range exps { |
||||
|
expOut = append(expOut, map[string]interface{}{ |
||||
|
"currency": e.Currency, |
||||
|
// 同时给「分」和「元」：模型直接念「元」不容易错，
|
||||
|
// 但保留「分」以免它想自己做加减时用浮点数
|
||||
|
"total_cents": e.TotalCents, |
||||
|
"total": float64(e.TotalCents) / 100.0, |
||||
|
"count": e.Count, |
||||
|
}) |
||||
|
} |
||||
|
|
||||
|
return mcp.NewToolResultText(utils.ToString(map[string]interface{}{ |
||||
|
"result": "success", |
||||
|
"start_date": start, |
||||
|
"end_date": end, |
||||
|
"todo_total": todoTotal, |
||||
|
"todo_done": todoDone, |
||||
|
"todo_undone": todoTotal - todoDone, |
||||
|
"alarm_total": alarmTotal, |
||||
|
"idea_count": ideaCount, |
||||
|
"expenses": expOut, |
||||
|
})), nil |
||||
|
} |
||||
|
|
||||
|
// ---------- get_memory_report ----------
|
||||
|
|
||||
|
type tool_get_memory_report struct { |
||||
|
cbase.ModuleCompBase |
||||
|
module *Mcp |
||||
|
} |
||||
|
|
||||
|
func (this *tool_get_memory_report) Init(service core.IService, module core.IModule, comp core.IModuleComp, opt core.IModuleOptions) (err error) { |
||||
|
this.ModuleCompBase.Init(service, module, comp, opt) |
||||
|
this.module = module.(*Mcp) |
||||
|
return |
||||
|
} |
||||
|
|
||||
|
func (this *tool_get_memory_report) Start() (err error) { |
||||
|
err = this.ModuleCompBase.Start() |
||||
|
this.module.AddTool(ToolGroup_GLOBAL, this.Tool(), this.Handl) |
||||
|
return |
||||
|
} |
||||
|
|
||||
|
func (this *tool_get_memory_report) Tool() mcp.Tool { |
||||
|
return mcp.NewTool("get_memory_report", |
||||
|
mcp.WithDescription("获取用户的周报或月报（系统定期生成的总结）"), |
||||
|
mcp.WithString("period_type", |
||||
|
mcp.Description("week=周报 month=月报"), |
||||
|
mcp.DefaultString("week"), |
||||
|
), |
||||
|
mcp.WithString("period_key", |
||||
|
mcp.Description("周期标识，周报形如 2026-W35，月报形如 2026-08；留空=最近一份"), |
||||
|
), |
||||
|
) |
||||
|
} |
||||
|
|
||||
|
func (this *tool_get_memory_report) Handl(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) { |
||||
|
uid, err := this.module.ResolveUID(ctx, request.GetString("uid", "")) |
||||
|
if err != nil { |
||||
|
return mcp.NewToolResultError(err.Error()), nil |
||||
|
} |
||||
|
ptype := strings.ToLower(strings.TrimSpace(request.GetString("period_type", comm.MemoryPeriodWeek))) |
||||
|
if ptype != comm.MemoryPeriodWeek && ptype != comm.MemoryPeriodMonth { |
||||
|
ptype = comm.MemoryPeriodWeek |
||||
|
} |
||||
|
pkey := strings.TrimSpace(request.GetString("period_key", "")) |
||||
|
|
||||
|
recs := make([]*pb.DBMemoryReport, 0) |
||||
|
tx := mysql.Table(comm.TableMemoryReport). |
||||
|
Where("uid = ? and period_type = ?", uid, ptype). |
||||
|
// 只给已生成完的：还在生成中的报告 summary 是空的，
|
||||
|
// 给出去只会让模型编一段话填空。
|
||||
|
Where("state = ?", int32(pb.MemoryReportState_MemoryReportState_Done)) |
||||
|
if pkey != "" { |
||||
|
tx = tx.Where("period_key = ?", pkey) |
||||
|
} |
||||
|
if e := tx.Order("period_end desc").Limit(1).Find(&recs).Error; e != nil { |
||||
|
this.module.Errorln("get_memory_report error:", e) |
||||
|
return mcp.NewToolResultError("查询失败"), nil |
||||
|
} |
||||
|
if len(recs) == 0 { |
||||
|
return mcp.NewToolResultText(utils.ToString(map[string]interface{}{ |
||||
|
"result": "empty", "message": "还没有生成过对应的报告", |
||||
|
})), nil |
||||
|
} |
||||
|
r := recs[0] |
||||
|
return mcp.NewToolResultText(utils.ToString(map[string]interface{}{ |
||||
|
"result": "success", |
||||
|
"period_type": r.PeriodType, |
||||
|
"period_key": r.PeriodKey, |
||||
|
"period_start": r.PeriodStart, |
||||
|
"period_end": r.PeriodEnd, |
||||
|
"summary": r.Summary, |
||||
|
"stat": r.StatJson, |
||||
|
})), nil |
||||
|
} |
||||
|
|
||||
|
// ---------- 公共辅助 ----------
|
||||
|
|
||||
|
// resolveRange 解析时间范围：range 优先，其次 start_date/end_date，都没有则默认最近 7 天。
|
||||
|
//
|
||||
|
// 提供 range 这种相对说法是刻意的：模型对「今天几号」的认知来自对话上下文，
|
||||
|
// 很不可靠，让它自己算 start_date 经常算错（SET_clock 把「下午 3:30」填成
|
||||
|
// 03:30 是同一类问题）。
|
||||
|
func resolveRange(request mcp.CallToolRequest) (start, end string, err error) { |
||||
|
now := time.Now() |
||||
|
if r := strings.TrimSpace(request.GetString("range", "")); r != "" { |
||||
|
return comm.MemoryRelativeRange(r, now) |
||||
|
} |
||||
|
start = strings.TrimSpace(request.GetString("start_date", "")) |
||||
|
end = strings.TrimSpace(request.GetString("end_date", "")) |
||||
|
if start == "" && end == "" { |
||||
|
return comm.FormatMemoryDate(now.AddDate(0, 0, -7)), |
||||
|
comm.FormatMemoryDate(now.AddDate(0, 0, 7)), nil |
||||
|
} |
||||
|
if start != "" { |
||||
|
if _, ok := comm.ParseMemoryDate(start); !ok { |
||||
|
return "", "", errBadDate(start) |
||||
|
} |
||||
|
} |
||||
|
if end != "" { |
||||
|
if _, ok := comm.ParseMemoryDate(end); !ok { |
||||
|
return "", "", errBadDate(end) |
||||
|
} |
||||
|
} |
||||
|
return start, end, nil |
||||
|
} |
||||
|
|
||||
|
func errBadDate(s string) error { |
||||
|
return &badDateError{s} |
||||
|
} |
||||
|
|
||||
|
type badDateError struct{ v string } |
||||
|
|
||||
|
func (e *badDateError) Error() string { |
||||
|
return "日期格式应为 YYYY-MM-DD：" + e.v |
||||
|
} |
||||
|
|
||||
|
// briefItems 只给模型必要字段。
|
||||
|
//
|
||||
|
// 不把整行 DBMemoryItem 丢过去：那里面有 client_key、gen_round、extra 这类
|
||||
|
// 内部字段，对回答问题毫无用处，只会占 token 并诱导模型去解释它们。
|
||||
|
func briefItems(items []*pb.DBMemoryItem) []map[string]interface{} { |
||||
|
out := make([]map[string]interface{}, 0, len(items)) |
||||
|
for _, it := range items { |
||||
|
m := map[string]interface{}{ |
||||
|
"id": it.Id, |
||||
|
"category": it.Category, |
||||
|
"title": it.Title, |
||||
|
"date": it.HappenDate, |
||||
|
"done": it.State == pb.MemoryState_MemoryState_Done, |
||||
|
} |
||||
|
if it.HappenTime != "" { |
||||
|
m["time"] = it.HappenTime |
||||
|
} |
||||
|
if it.Detail != "" { |
||||
|
m["detail"] = it.Detail |
||||
|
} |
||||
|
if it.Owner != "" { |
||||
|
m["owner"] = it.Owner |
||||
|
} |
||||
|
if it.DueRaw != "" { |
||||
|
m["due"] = it.DueRaw |
||||
|
} |
||||
|
if it.Category == comm.MemoryCatExpense { |
||||
|
m["amount"] = float64(it.AmountCents) / 100.0 |
||||
|
m["currency"] = it.Currency |
||||
|
} |
||||
|
if !it.DateCertain { |
||||
|
// 让模型知道这个日期是推断的，别把它当成用户确认过的安排来陈述
|
||||
|
m["date_is_inferred"] = true |
||||
|
} |
||||
|
out = append(out, m) |
||||
|
} |
||||
|
return out |
||||
|
} |
||||
Loading…
Reference in new issue