diff --git a/go.mod b/go.mod index f01728e..9bf29e4 100644 --- a/go.mod +++ b/go.mod @@ -13,6 +13,8 @@ require ( require ( agent-squad v0.0.0-00010101000000-000000000000 // indirect error-code-system v0.0.0 // indirect + github.com/cespare/xxhash/v2 v2.3.0 // indirect + github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f // indirect github.com/dustin/go-humanize v1.0.1 // indirect github.com/google/uuid v1.6.0 // indirect github.com/jackc/pgpassfile v1.0.0 // indirect @@ -21,6 +23,7 @@ require ( github.com/jackc/puddle/v2 v2.2.2 // indirect github.com/mattn/go-isatty v0.0.24 // indirect github.com/ncruces/go-strftime v1.0.0 // indirect + github.com/redis/go-redis/v9 v9.14.0 // indirect github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec // indirect golang.org/x/crypto v0.37.0 // indirect golang.org/x/sync v0.23.0 // indirect diff --git a/go.sum b/go.sum index 1db8903..535e370 100644 --- a/go.sum +++ b/go.sum @@ -1,4 +1,8 @@ +github.com/cespare/xxhash/v2 v2.3.0 h1:UL815xU9SqsFlibzuggzjXhog7bL6oX9BbNZnL2UFvs= +github.com/cespare/xxhash/v2 v2.3.0/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs= github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= +github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f h1:lO4WD4F/rVNCu3HqELle0jiPLLBs70cWOduZpkS1E78= +github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f/go.mod h1:cuUVRXasLTGF7a8hSLbxyZXjz+1KgoB3wDUb6vlszIc= github.com/dustin/go-humanize v1.0.1 h1:GzkhY7T5VNhEkwH0PVJgjz+fX1rhBrR7pRT3mDkpeCY= github.com/dustin/go-humanize v1.0.1/go.mod h1:Mu1zIs6XwVuF/gI1OepvI0qD18qycQx+mFykh5fBlto= github.com/google/pprof v0.0.0-20260802141513-ef3492d7dac3 h1:LMLX+LgTNWpfvCBdFebv6EsYotImrt/Ppc5cXIriCSo= @@ -22,6 +26,8 @@ github.com/mattn/go-isatty v0.0.24/go.mod h1:nMCL3Zebbrt45jsMDgnfIwz6ydEQApk5oEI github.com/ncruces/go-strftime v1.0.0 h1:HMFp8mLCTPp341M/ZnA4qaf7ZlsbTc+miZjCLOFAw7w= github.com/ncruces/go-strftime v1.0.0/go.mod h1:Fwc5htZGVVkseilnfgOVb9mKy6w1naJmn9CehxcKcls= github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= +github.com/redis/go-redis/v9 v9.14.0 h1:u4tNCjXOyzfgeLN+vAZaW1xUooqWDqVEsZN0U01jfAE= +github.com/redis/go-redis/v9 v9.14.0/go.mod h1:huWgSWd8mW6+m0VPhJjSSQ+d6Nh1VICQ6Q5lHuCH/Iw= github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec h1:W09IVJc94icq4NjY3clb7Lk8O1qJ8BdBEF8z0ibU0rE= github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec/go.mod h1:qqbHyh8v60DhA7CoWK5oRCqLrMHRGoxYCSS9EjAz6Eo= github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME= diff --git a/internal/app/app.go b/internal/app/app.go index 1d07dcf..4d2d246 100644 --- a/internal/app/app.go +++ b/internal/app/app.go @@ -11,28 +11,31 @@ import ( "os" "os/signal" "path/filepath" + "sync" "syscall" "time" - core "agent-system-core" "agent-system-backend/internal/bus" "agent-system-backend/internal/intentgate" + "agent-system-backend/internal/pgbus" "agent-system-backend/internal/plugins" - "agent-system-backend/internal/rag" "agent-system-backend/internal/profiles" + "agent-system-backend/internal/rag" + "agent-system-backend/internal/ratelimit" "agent-system-backend/internal/runmgr" "agent-system-backend/internal/runtime" "agent-system-backend/internal/server" "agent-system-backend/internal/settings" "agent-system-backend/internal/store" "agent-system-backend/internal/tools" + core "agent-system-core" ) type Options struct { - Addr string - DataDir string - PluginsDir string - ForceMock bool + Addr string + DataDir string + PluginsDir string + ForceMock bool StorageDSN string // 空 = 嵌入式 SQLite 档；非空 = PostgreSQL 档（多用户/多实例） ArtifactsDir string // PG 档工件对象目录（空 = /artifacts） AuthMode string // auto（默认：PG 档开、嵌入式关）| on | off @@ -101,6 +104,22 @@ func Run(opt Options) error { srv := server.New(st, mgr, reg, profSvc, setSvc, pluginHost) srv.Retriever = pipe srv.Indexer = pipe + + // 限流（阶段 3）：默认进程内档；设置 runtime.redisAddr 切原生 Redis 共享档。 + var limiter ratelimit.Limiter = ratelimit.NewMemory() + if addr := setSvc.Effective(store.OwnerDefault, "").GetString("runtime.redisAddr", ""); addr != "" { + rl := ratelimit.NewRedis(addr, "", 0) + if err := rl.Ping(context.Background()); err != nil { + return fmt.Errorf("redis %s: %w", addr, err) + } + defer rl.Close() + limiter = rl + log.Printf("rate limit: shared via redis %s", addr) + } + srv.Limiter = limiter + srv.LimitLookup = cachedInt(5*time.Second, func() int { + return int(setSvc.Effective(store.OwnerDefault, "").GetNumber("runtime.rateLimitPerMin", 0)) + }) // 认证策略（auto）：PG 档默认开启（多用户面不能裸奔）；嵌入式档默认关闭（本地单用户）。 authOn := opt.AuthMode == "on" || (opt.AuthMode != "off" && opt.StorageDSN != "") srv.Auth = server.NewAuthenticator(st, authOn) @@ -120,6 +139,17 @@ func Run(opt Options) error { } }() + // PG 档：事件跨实例扇出（NOTIFY/LISTEN；SQLite 档不需要）。 + if st.Dialect() == "postgres" { + mgr.SetEventNotifier(pgbus.NewNotifier(st).Notify) + listener, err := pgbus.Start(context.Background(), opt.StorageDSN, st, b) + if err != nil { + return fmt.Errorf("pgbus listen: %w", err) + } + defer listener.Close(context.Background()) + log.Printf("pgbus: LISTEN agent_events（跨实例事件扇出）") + } + stop := make(chan os.Signal, 1) signal.Notify(stop, os.Interrupt, syscall.SIGTERM) <-stop @@ -184,3 +214,19 @@ func RevokeAPIKey(storageDSN, dataDir, owner, name string) (bool, error) { defer st.Close() return st.RevokeAPIKeyByName(owner, name) } + +// cachedInt 带短 TTL 的整型缓存读（限流值热读取免打库）。 +func cachedInt(ttl time.Duration, fn func() int) func() int { + var mu sync.Mutex + var val int + var at time.Time + return func() int { + mu.Lock() + defer mu.Unlock() + if time.Since(at) > ttl { + val = fn() + at = time.Now() + } + return val + } +} diff --git a/internal/blobstore/blobstore.go b/internal/blobstore/blobstore.go new file mode 100644 index 0000000..5c18263 --- /dev/null +++ b/internal/blobstore/blobstore.go @@ -0,0 +1,64 @@ +// Package blobstore 是工件对象存储抽象（阶段 3：多实例前置）。 +// 默认本地目录档（单实例，零外部服务）；多实例部署切共享卷路径或按本接口补 +// 对象存储实现（S3 兼容等）——上层（store 工件方法）只依赖本接口。 +package blobstore + +import ( + "fmt" + "os" + "path/filepath" + "sync" +) + +// Store 是对象存储接口：Put 原子落盘并返回记录进库的引用（本地档=绝对路径， +// 对象存储档=对象键）；Get/Delete 按该引用操作。 +type Store interface { + Put(id string, content []byte) (ref string, err error) + Get(ref string) ([]byte, error) + Delete(ref string) error +} + +// Local 是本地目录实现（data/artifacts；多实例时把目录指向共享卷即可）。 +type Local struct { + dir string + mu sync.Mutex +} + +func NewLocal(dir string) (*Local, error) { + if err := os.MkdirAll(dir, 0o755); err != nil { + return nil, err + } + return &Local{dir: dir}, nil +} + +// Put 临时文件 → rename 原子写：崩溃只可能残留 .tmp 孤立文件，不产生半写对象。 +func (l *Local) Put(id string, content []byte) (string, error) { + final := filepath.Join(l.dir, id) + tmp := final + ".tmp-" + id + if err := os.WriteFile(tmp, content, 0o644); err != nil { + return "", err + } + l.mu.Lock() + defer l.mu.Unlock() + if err := os.Rename(tmp, final); err != nil { + os.Remove(tmp) + return "", err + } + return final, nil +} + +func (l *Local) Get(ref string) ([]byte, error) { + b, err := os.ReadFile(ref) + if os.IsNotExist(err) { + return nil, fmt.Errorf("blobstore: object %s not found", filepath.Base(ref)) + } + return b, err +} + +func (l *Local) Delete(ref string) error { + err := os.Remove(ref) + if os.IsNotExist(err) { + return nil + } + return err +} diff --git a/internal/pgbus/pgbus.go b/internal/pgbus/pgbus.go new file mode 100644 index 0000000..171be69 --- /dev/null +++ b/internal/pgbus/pgbus.go @@ -0,0 +1,108 @@ +// Package pgbus 是 PG 档的事件跨实例扇出（阶段 3）： +// 落盘事务提交后 NOTIFY agent_events '|'；每实例一个 LISTEN 连接， +// 收到通知按 seq 从库取回事件并转发本地 bus（SSE 以 seq 去重，双投无害）。 +// 多实例部署契约：run 在其归属实例内存中执行，会话流量需按会话亲和路由 +//（网关/客户端择点）；本包只解决“其他实例的订阅者”可见性。 +package pgbus + +import ( + "context" + "log" + "strconv" + "strings" + "time" + + core "agent-system-core" + "agent-system-backend/internal/bus" + "agent-system-backend/internal/store" + + "github.com/jackc/pgx/v5" +) + +const channel = "agent_events" + +// Notifier 落盘后通知（任意连接池连接可用）。 +type Notifier struct{ st *store.Store } + +func NewNotifier(st *store.Store) *Notifier { return &Notifier{st: st} } + +// Notify 在 AppendEvents 成功后调用：一批事件一条通知（载荷=会话|末 seq）。 +func (n *Notifier) Notify(conversationID string, evs []core.StreamEvent) { + if len(evs) == 0 { + return + } + ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second) + defer cancel() + payload := conversationID + "|" + strconv.FormatInt(evs[len(evs)-1].Seq, 10) + if _, err := n.st.NotifyPG(ctx, channel, payload); err != nil { + log.Printf("[pgbus] notify %s: %v", conversationID, err) + } +} + +// Listener 每实例一个：LISTEN → 按 seq 取回 → 本地 bus 转发。 +type Listener struct { + conn *pgx.Conn + st *store.Store + bus *bus.Bus + done chan struct{} +} + +// Start 建立监听连接并开始转发循环；返回后即可用，随 ctx 取消停机。 +func Start(ctx context.Context, dsn string, st *store.Store, b *bus.Bus) (*Listener, error) { + conn, err := pgx.Connect(ctx, dsn) + if err != nil { + return nil, err + } + if _, err := conn.Exec(ctx, "LISTEN "+channel); err != nil { + conn.Close(ctx) + return nil, err + } + l := &Listener{conn: conn, st: st, bus: b, done: make(chan struct{})} + go l.loop(ctx) + return l, nil +} + +func (l *Listener) loop(ctx context.Context) { + defer close(l.done) + for { + n, err := l.conn.WaitForNotification(ctx) + if err != nil { + if ctx.Err() != nil { + return // 停机 + } + log.Printf("[pgbus] wait: %v（1s 后重连）", err) + time.Sleep(time.Second) + continue + } + conv, seq, ok := parsePayload(n.Payload) + if !ok { + continue + } + ev, err := l.st.EventAt(conv, seq) + if err != nil { + continue // 事件不可见（极早期窗口）——SSE 重放兜底 + } + l.bus.Publish(ev) + } +} + +func parsePayload(p string) (string, int64, bool) { + i := strings.LastIndex(p, "|") + if i < 0 { + return "", 0, false + } + seq, err := strconv.ParseInt(p[i+1:], 10, 64) + if err != nil { + return "", 0, false + } + return p[:i], seq, true +} + +// Close 关闭监听连接。 +func (l *Listener) Close(ctx context.Context) { + l.conn.Close(ctx) + select { + case <-l.done: + case <-time.After(2 * time.Second): + } +} diff --git a/internal/pgbus/pgbus_test.go b/internal/pgbus/pgbus_test.go new file mode 100644 index 0000000..4fe778f --- /dev/null +++ b/internal/pgbus/pgbus_test.go @@ -0,0 +1,82 @@ +package pgbus + +// PG 档事件扇出集成测试（环境门控：AGENT_PG_TEST_DSN，原生 PostgreSQL 实例）。 +// 验证：实例 A 落盘+NOTIFY → 实例 B 的 LISTEN 取回并转发到其本地 bus。 + +import ( + "context" + "database/sql" + "fmt" + "os" + "path/filepath" + "strconv" + "testing" + "time" + + core "agent-system-core" + "agent-system-backend/internal/bus" + "agent-system-backend/internal/store" +) + +func pgTestStore(t *testing.T) (*store.Store, string) { + t.Helper() + dsn := os.Getenv("AGENT_PG_TEST_DSN") + if dsn == "" { + t.Skip("AGENT_PG_TEST_DSN not set") + } + schema := "pgbus_" + strconv.Itoa(os.Getpid()) + "_" + fmt.Sprint(time.Now().UnixNano()%1e6) + admin, err := sql.Open("pgx", dsn) + if err != nil { + t.Fatal(err) + } + if _, err := admin.Exec("CREATE SCHEMA " + schema); err != nil { + admin.Close() + t.Fatal(err) + } + t.Cleanup(func() { + admin.Exec("DROP SCHEMA IF EXISTS " + schema + " CASCADE") + admin.Close() + }) + full := dsn + "&search_path=" + schema + st, err := store.OpenPostgres(full, filepath.Join(t.TempDir(), "artifacts")) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { st.Close() }) + return st, full +} + +func TestNotifyListenerFanout(t *testing.T) { + st, dsn := pgTestStore(t) + + // “实例 B”：独立 bus + 监听连接。 + b2 := bus.New() + listener, err := Start(context.Background(), dsn, st, b2) + if err != nil { + t.Fatalf("listen: %v", err) + } + defer listener.Close(context.Background()) + + ch, stop := b2.Subscribe("conv_x") + defer stop() + + // “实例 A”：批落盘 + NOTIFY（Notifier 与 listener 共库不同连接）。 + notifier := NewNotifier(st) + evs := []core.StreamEvent{ + {ConversationID: "conv_x", RunID: "r1", TS: time.Now().UTC().Format(time.RFC3339Nano), V: 1, Type: core.EvRunPhase, Payload: []byte(`{"phase":"thinking"}`)}, + {ConversationID: "conv_x", RunID: "r1", TS: time.Now().UTC().Format(time.RFC3339Nano), V: 1, Type: core.EvRunCompleted, Payload: []byte(`{}`)}, + } + if _, err := st.AppendEvents("conv_x", evs); err != nil { + t.Fatal(err) + } + notifier.Notify("conv_x", evs) + + select { + case e := <-ch: + if e.ConversationID != "conv_x" || e.Seq == 0 { + t.Fatalf("fanout event malformed: %+v", e) + } + case <-time.After(3 * time.Second): + t.Fatal("no fanout within 3s") + } +} diff --git a/internal/ratelimit/ratelimit.go b/internal/ratelimit/ratelimit.go new file mode 100644 index 0000000..a22331a --- /dev/null +++ b/internal/ratelimit/ratelimit.go @@ -0,0 +1,82 @@ +// Package ratelimit 提供按租户的限流（阶段 3）： +// - Memory：进程内固定窗口（默认，单实例档，零外部服务）； +// - Redis：多实例共享固定窗口（原生 Redis，设置 runtime.redisAddr 启用）。 +package ratelimit + +import ( + "context" + "sync" + "time" + + "github.com/redis/go-redis/v9" +) + +type Limiter interface { + // Allow 判定 key（租户 owner）在窗口内第 N 次请求是否放行。 + Allow(ctx context.Context, key string, limit int, window time.Duration) bool +} + +// ---- 进程内固定窗口 ---- + +type memWindow struct { + n int + reset time.Time +} + +type Memory struct { + mu sync.Mutex + windows map[string]*memWindow + now func() time.Time +} + +func NewMemory() *Memory { + return &Memory{windows: map[string]*memWindow{}, now: time.Now} +} + +func (m *Memory) Allow(_ context.Context, key string, limit int, window time.Duration) bool { + if limit <= 0 { + return true + } + m.mu.Lock() + defer m.mu.Unlock() + now := m.now() + w, ok := m.windows[key] + if !ok || !now.Before(w.reset) { + m.windows[key] = &memWindow{n: 1, reset: now.Add(window)} + return true + } + if w.n >= limit { + return false + } + w.n++ + return true +} + +// ---- Redis 固定窗口（INCR + EXPIRE；首击设窗）---- + +type Redis struct { + rdb *redis.Client +} + +func NewRedis(addr, password string, db int) *Redis { + return &Redis{rdb: redis.NewClient(&redis.Options{Addr: addr, Password: password, DB: db})} +} + +func (r *Redis) Allow(ctx context.Context, key string, limit int, window time.Duration) bool { + if limit <= 0 { + return true + } + rkey := "rl:" + key + ":" + time.Now().UTC().Truncate(window).Format("150405") + n, err := r.rdb.Incr(ctx, rkey).Result() + if err != nil { + return true // Redis 故障不阻断请求（限流是柔性行为，放行并记日志于调用方） + } + if n == 1 { + r.rdb.Expire(ctx, rkey, window+time.Second) + } + return n <= int64(limit) +} + +func (r *Redis) Ping(ctx context.Context) error { return r.rdb.Ping(ctx).Err() } + +func (r *Redis) Close() error { return r.rdb.Close() } diff --git a/internal/ratelimit/ratelimit_test.go b/internal/ratelimit/ratelimit_test.go new file mode 100644 index 0000000..94d6a95 --- /dev/null +++ b/internal/ratelimit/ratelimit_test.go @@ -0,0 +1,57 @@ +package ratelimit + +import ( + "context" + "os" + "testing" + "time" +) + +func TestMemoryWindow(t *testing.T) { + m := NewMemory() + ctx := context.Background() + for i := 1; i <= 3; i++ { + if !m.Allow(ctx, "o", 3, time.Minute) { + t.Fatalf("request %d should pass", i) + } + } + if m.Allow(ctx, "o", 3, time.Minute) { + t.Fatal("4th request must be limited") + } + // 其他租户互不影响。 + if !m.Allow(ctx, "p", 3, time.Minute) { + t.Fatal("other tenant must pass") + } + // 窗口滚动后恢复。 + m.now = func() time.Time { return time.Now().Add(2 * time.Minute) } + if !m.Allow(ctx, "o", 3, time.Minute) { + t.Fatal("new window must reset") + } + // limit<=0 = 不限。 + if !m.Allow(ctx, "x", 0, time.Minute) { + t.Fatal("limit 0 means unlimited") + } +} + +// Redis 共享档（环境门控：AGENT_REDIS_TEST_ADDR，本机原生实例）。 +func TestRedisWindow(t *testing.T) { + addr := os.Getenv("AGENT_REDIS_TEST_ADDR") + if addr == "" { + t.Skip("AGENT_REDIS_TEST_ADDR not set") + } + r := NewRedis(addr, "", 15) // 独立 db 避免污染 + defer r.Close() + if err := r.Ping(context.Background()); err != nil { + t.Fatalf("ping: %v", err) + } + ctx := context.Background() + key := "test-" + time.Now().Format("150405.000000000") + for i := 1; i <= 2; i++ { + if !r.Allow(ctx, key, 2, time.Minute) { + t.Fatalf("request %d should pass", i) + } + } + if r.Allow(ctx, key, 2, time.Minute) { + t.Fatal("3rd request must be limited") + } +} diff --git a/internal/runmgr/events.go b/internal/runmgr/events.go index 221edbf..e1cedcc 100644 --- a/internal/runmgr/events.go +++ b/internal/runmgr/events.go @@ -24,6 +24,8 @@ type eventBuffer struct { bus *bus.Bus interval time.Duration // 定时冲盘窗口 maxBuf int // 单会话缓冲条数上限（触发立即冲盘） + // onPersisted 落盘+本地发布后的回调（PG 档 NOTIFY 跨实例扇出；nil = 无）。 + onPersisted func(conversationID string, evs []core.StreamEvent) mu sync.Mutex buf map[string][]core.StreamEvent @@ -119,6 +121,9 @@ func (eb *eventBuffer) append(conversationID string, evs []core.StreamEvent) err for i := range evs { eb.bus.Publish(evs[i]) } + if eb.onPersisted != nil { + eb.onPersisted(conversationID, evs) + } return nil } diff --git a/internal/runmgr/runmgr.go b/internal/runmgr/runmgr.go index cf2f028..d95d04e 100644 --- a/internal/runmgr/runmgr.go +++ b/internal/runmgr/runmgr.go @@ -70,6 +70,11 @@ func New(st *store.Store, b *bus.Bus, reg *tools.Registry, p *profiles.Service, // Shutdown 停机清空事件缓冲（在 store.Close 之前调用）。 func (m *Manager) Shutdown() { m.events.stop() } +// SetEventNotifier 注入落盘后回调（PG 档 pgbus.Notifier；SQLite 档不注入）。 +func (m *Manager) SetEventNotifier(fn func(conversationID string, evs []core.StreamEvent)) { + m.events.onPersisted = fn +} + // IntentGate 是可选意图闸门端口（intentgate.Gate 实现；nil = 未启用，行为回退）。 type IntentGate interface { Match(ctx context.Context, input string) (branch, topID string) diff --git a/internal/server/ratelimit.go b/internal/server/ratelimit.go new file mode 100644 index 0000000..9fcc9ac --- /dev/null +++ b/internal/server/ratelimit.go @@ -0,0 +1,26 @@ +package server + +import ( + "net/http" + "time" + + "agent-system-backend/internal/ratelimit" +) + +// 限流中间件（阶段 3）：按租户 owner 固定窗口；limit<=0（默认）直接放行。 +// 位于认证之后（key=owner）；Limiter/LimitLookup 由组合根注入。 +func (s *Server) withRateLimit(next http.Handler) http.Handler { + return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if s.Limiter != nil && s.LimitLookup != nil { + if limit := s.LimitLookup(); limit > 0 { + if !s.Limiter.Allow(r.Context(), "owner:"+s.owner(r), limit, time.Minute) { + http.Error(w, "rate limit exceeded（runtime.rateLimitPerMin）", http.StatusTooManyRequests) + return + } + } + } + next.ServeHTTP(w, r) + }) +} + +var _ = ratelimit.Memory{} // 接口形状锚点（Limiter 注入自组合根） diff --git a/internal/server/ratelimit_e2e_test.go b/internal/server/ratelimit_e2e_test.go new file mode 100644 index 0000000..d74f47e --- /dev/null +++ b/internal/server/ratelimit_e2e_test.go @@ -0,0 +1,40 @@ +package server + +// 限流 e2e：limit=3 时第 4 个请求 429；其他租户不受影响。 + +import ( + "net/http" + "net/http/httptest" + "testing" + + "agent-system-backend/internal/ratelimit" +) + +func TestRateLimitMiddleware(t *testing.T) { + srv, _ := newTestServer(t) + srv.Limiter = ratelimit.NewMemory() + n := 3 + srv.LimitLookup = func() int { return n } + h := srv.Handler() + + get := func() int { + req := httptest.NewRequest("GET", "/api/projects", nil) + rec := httptest.NewRecorder() + h.ServeHTTP(rec, req) + return rec.Code + } + for i := 1; i <= n; i++ { + if c := get(); c != http.StatusOK { + t.Fatalf("request %d: %d", i, c) + } + } + if c := get(); c != http.StatusTooManyRequests { + t.Fatalf("4th request: %d (want 429)", c) + } + // lookup<=0 时放行（默认不限）。 + n = 0 + srv.Limiter = ratelimit.NewMemory() + if c := get(); c != http.StatusOK { + t.Fatalf("unlimited: %d", c) + } +} diff --git a/internal/server/server.go b/internal/server/server.go index 74a2ad9..01f2eb0 100644 --- a/internal/server/server.go +++ b/internal/server/server.go @@ -19,6 +19,7 @@ import ( "agent-system-backend/internal/ingest" "agent-system-backend/internal/plugins" "agent-system-backend/internal/profiles" + "agent-system-backend/internal/ratelimit" "agent-system-backend/internal/runmgr" "agent-system-backend/internal/settings" "agent-system-backend/internal/store" @@ -31,6 +32,8 @@ type Server struct { Auth *Authenticator // API key 认证（nil/未启用 = 默认租户；组合根注入） Retriever core.Retriever // 混合检索（nil = store 全文；组合根注入） Indexer core.DocumentIndexer // 检索索引入口（nil = 整文档单块；组合根注入） + Limiter ratelimit.Limiter // 租户限流（nil = 不限；组合根注入） + LimitLookup func() int // 当前限流值（次/分；<=0 不限；组合根缓存注入） mgr *runmgr.Manager reg *tools.Registry prof *profiles.Service @@ -113,10 +116,11 @@ func (s *Server) Handler() http.Handler { // API 未匹配路径一律 404（不回落 SPA，避免 /api/* 返回 HTML 200）。 mux.HandleFunc("/api/", func(w http.ResponseWriter, r *http.Request) { http.NotFound(w, r) }) - // 静态前端（embed）；API 全链路经认证中间件（未启用 = 默认租户）。 - var h http.Handler = mux + // 静态前端（embed）；API 全链路：认证（未启用 = 默认租户）→ 租户限流。 + inner := s.withRateLimit(mux) + var h http.Handler = inner if s.Auth != nil { - h = s.Auth.Middleware(mux) + h = s.Auth.Middleware(inner) } mux.HandleFunc("/", s.serveStatic) return h diff --git a/internal/settings/settings.go b/internal/settings/settings.go index 2306e4e..ed2d655 100644 --- a/internal/settings/settings.go +++ b/internal/settings/settings.go @@ -29,6 +29,8 @@ func SystemDefaults() map[string]any { "runtime": map[string]any{ "mock": false, // true = 全部会话走 mock runner（演示/测试） "squad": true, // 小队编排（Agent-Squad 引擎）可用 + "rateLimitPerMin": 0, // 租户限流（次/分；0 = 不限） + "redisAddr": "", // 限流共享后端（原生 Redis addr；空 = 进程内） }, "storage": map[string]any{ "dataDir": "./data", // 启动注入实际值（--data） @@ -60,6 +62,8 @@ func Schema() []core.SettingDesc { {Key: "models.rerank.model", Group: "models", Label: "Reranker 模型", Type: core.SettingString, Default: "", Effect: core.EffectHot}, {Key: "runtime.mock", Group: "runtime", Label: "演示模式（mock runner）", Type: core.SettingBool, Default: false, Effect: core.EffectNewRun, Description: "不调用真实模型，按脚本回放事件流"}, {Key: "runtime.squad", Group: "runtime", Label: "小队编排（Agent-Squad）", Type: core.SettingBool, Default: true, Effect: core.EffectNewRun, Description: "启用 Agent-Squad 引擎编排（#squad 指令与 squad_orchestrate 工具）"}, + {Key: "runtime.rateLimitPerMin", Group: "runtime", Label: "租户限流（次/分）", Type: core.SettingNumber, Default: 0, Effect: core.EffectHot, Description: "0 = 不限；多用户建议 300–1200"}, + {Key: "runtime.redisAddr", Group: "runtime", Label: "Redis 地址（限流共享）", Type: core.SettingString, Default: "", Effect: core.EffectRestart, Description: "原生 Redis（如 127.0.0.1:6379）；空 = 进程内限流（单实例）"}, {Key: "context.maxHistoryMessages", Group: "context", Label: "历史窗口条数", Type: core.SettingNumber, Default: 40, Effect: core.EffectNewRun}, {Key: "context.maxToolResultChars", Group: "context", Label: "工具结果上限（字符）", Type: core.SettingNumber, Default: 20000, Effect: core.EffectNewRun}, {Key: "retrieval.topK", Group: "retrieval", Label: "检索命中数", Type: core.SettingNumber, Default: 5, Effect: core.EffectHot}, diff --git a/internal/store/content.go b/internal/store/content.go index 9e14d45..4c0b34b 100644 --- a/internal/store/content.go +++ b/internal/store/content.go @@ -5,8 +5,6 @@ import ( "database/sql" "errors" "fmt" - "os" - "path/filepath" "time" core "agent-system-core" @@ -16,18 +14,6 @@ import ( // 文件写走 临时文件 → rename → 落库：崩溃残留只可能产生无主孤立文件， // 不产生库里有行、对象缺失的悬空引用。 -func (s *Store) writeBlobAtomic(path string, content []byte) error { - tmp := path + ".tmp-" + newID("x") - if err := os.WriteFile(tmp, content, 0o644); err != nil { - return err - } - if err := os.Rename(tmp, path); err != nil { - os.Remove(tmp) - return err - } - return nil -} - func (s *Store) CreateArtifact(a *core.Artifact, content []byte) error { if a.ID == "" { a.ID = newID("art") @@ -35,14 +21,14 @@ func (s *Store) CreateArtifact(a *core.Artifact, content []byte) error { now := time.Now().UTC() a.CreatedAt, a.UpdatedAt = now, now a.Size = int64(len(content)) - path := filepath.Join(s.dataDir, a.ID) - if err := s.writeBlobAtomic(path, content); err != nil { + ref, err := s.blobs.Put(a.ID, content) + if err != nil { return err } - _, err := s.db.Exec(s.q(`INSERT INTO artifacts(id,project_id,owner,name,kind,mime,size,source_run,content_path,created_at,updated_at) VALUES(?,?,?,?,?,?,?,?,?,?,?)`), - a.ID, a.ProjectID, a.Owner, a.Name, a.Kind, a.MIME, a.Size, a.SourceRun, path, nowISO(), nowISO()) + _, err = s.db.Exec(s.q(`INSERT INTO artifacts(id,project_id,owner,name,kind,mime,size,source_run,content_path,created_at,updated_at) VALUES(?,?,?,?,?,?,?,?,?,?,?)`), + a.ID, a.ProjectID, a.Owner, a.Name, a.Kind, a.MIME, a.Size, a.SourceRun, ref, nowISO(), nowISO()) if err != nil { - os.Remove(path) // 落库失败：回收对象，不留悬空 + _ = s.blobs.Delete(ref) // 落库失败：回收对象，不留悬空 return err } return nil @@ -57,7 +43,8 @@ func (s *Store) UpdateArtifactContent(owner, id string, content []byte) error { if err != nil { return err } - if err := s.writeBlobAtomic(path, content); err != nil { + // 同 ID 原子覆写（临时文件+rename）；行内引用不变，避免行/对象漂移。 + if _, err := s.blobs.Put(id, content); err != nil { return err } res, err := s.db.Exec(s.q(`UPDATE artifacts SET size=?, updated_at=? WHERE id=? AND owner=?`), len(content), nowISO(), id, owner) @@ -96,15 +83,15 @@ func (s *Store) GetArtifact(owner, id string) (core.Artifact, error) { } func (s *Store) ArtifactContent(owner, id string) ([]byte, error) { - var path string - err := s.db.QueryRow(s.q(`SELECT content_path FROM artifacts WHERE id=? AND owner=?`), id, owner).Scan(&path) + var ref string + err := s.db.QueryRow(s.q(`SELECT content_path FROM artifacts WHERE id=? AND owner=?`), id, owner).Scan(&ref) if errors.Is(err, sql.ErrNoRows) { return nil, ErrNotFound } if err != nil { return nil, err } - return os.ReadFile(path) + return s.blobs.Get(ref) } func (s *Store) ListArtifacts(owner, projectID string) ([]core.Artifact, error) { @@ -140,10 +127,7 @@ func (s *Store) DeleteArtifact(owner, id string) error { // 联动清理检索索引（知识行 + 向量；失败不阻断删除，残留可由重建索引收敛）。 _ = s.DeleteKnowledgeByRef(owner, id) _ = s.DeleteKnowledgeVectors(context.Background(), owner, id) - if err := os.Remove(path); err != nil && !os.IsNotExist(err) { - return err - } - return nil + return s.blobs.Delete(path) } // ---- HITL 卡片（resolve 原子化：条件更新消灭先查后改窗口）---- diff --git a/internal/store/store.go b/internal/store/store.go index 2f4c0fb..594c8fe 100644 --- a/internal/store/store.go +++ b/internal/store/store.go @@ -9,6 +9,7 @@ package store import ( + "context" "database/sql" "encoding/json" "errors" @@ -17,6 +18,8 @@ import ( "path/filepath" "time" + "agent-system-backend/internal/blobstore" + _ "github.com/jackc/pgx/v5/stdlib" // PG 档驱动（注册名 "pgx"） _ "modernc.org/sqlite" // 嵌入式档驱动（注册名 "sqlite"，CGO-free） ) @@ -29,7 +32,8 @@ const OwnerDefault = "local" type Store struct { db *sql.DB dialect dialect - dataDir string // 工件对象目录（/artifacts 或配置的共享目录） + dataDir string // 工件对象目录（备份路径等诊断用） + blobs blobstore.Store // 工件对象存储（本地档；多实例切共享卷/对象存储） } // OpenSQLite 打开（必要时创建）嵌入式存储并执行迁移。 @@ -44,7 +48,7 @@ func OpenSQLite(path string) (*Store, error) { } // SQLite 单写者：串行化连接，语句级原子 + 版本化迁移事务。 db.SetMaxOpenConns(1) - return finishOpen(db, sqliteDialect{}, filepath.Join(filepath.Dir(path), "artifacts")) + return finishOpen(db, sqliteDialect{}, filepath.Join(filepath.Dir(path), "artifacts"), nil) } // OpenPostgres 打开 PostgreSQL 档存储并执行迁移。 @@ -58,11 +62,31 @@ func OpenPostgres(dsn, artifactsDir string) (*Store, error) { db.SetMaxOpenConns(16) db.SetMaxIdleConns(8) db.SetConnMaxIdleTime(5 * time.Minute) - return finishOpen(db, postgresDialect{}, artifactsDir) + return finishOpen(db, postgresDialect{}, artifactsDir, nil) } -func finishOpen(db *sql.DB, d dialect, artifactsDir string) (*Store, error) { - s := &Store{db: db, dialect: d, dataDir: artifactsDir} +// OpenPostgresWithBlobs 是注入自定义对象存储的 PG 档入口（多实例共享卷/对象存储）。 +func OpenPostgresWithBlobs(dsn, artifactsDir string, blobs blobstore.Store) (*Store, error) { + db, err := sql.Open("pgx", dsn) + if err != nil { + return nil, err + } + db.SetMaxOpenConns(16) + db.SetMaxIdleConns(8) + db.SetConnMaxIdleTime(5 * time.Minute) + return finishOpen(db, postgresDialect{}, artifactsDir, blobs) +} + +func finishOpen(db *sql.DB, d dialect, artifactsDir string, blobs blobstore.Store) (*Store, error) { + if blobs == nil { + var err error + blobs, err = blobstore.NewLocal(artifactsDir) + if err != nil { + db.Close() + return nil, err + } + } + s := &Store{db: db, dialect: d, dataDir: artifactsDir, blobs: blobs} if err := s.migrate(); err != nil { db.Close() return nil, err @@ -82,6 +106,18 @@ func (s *Store) Dialect() string { return s.dialect.name() } // DataDir 返回工件对象目录。 func (s *Store) DataDir() string { return s.dataDir } +// BlobStore 返回工件对象存储（诊断/运维）。 +func (s *Store) BlobStore() blobstore.Store { return s.blobs } + +// NotifyPG 发送 PostgreSQL NOTIFY（PG 档事件扇出用；SQLite 档返回未执行）。 +func (s *Store) NotifyPG(ctx context.Context, channel, payload string) (bool, error) { + if s.dialect.name() != "postgres" { + return false, nil + } + _, err := s.db.ExecContext(ctx, `SELECT pg_notify($1, $2)`, channel, payload) + return true, err +} + // q 按方言重绑占位符（查询统一用 ? 书写，PG 执行前转 $n）。 // 本包 SQL 不含字符串字面量 "?"，重绑安全。 func (s *Store) q(query string) string { return s.dialect.rebind(query) }