weichen
29f14ac5f2
Details (original message): Add multi-instance groundwork: blobstore seam for artifacts, PG LISTEN/NOTIFY event fanout, and per-tenant rate limiting with native Redis backend
233 lines
7.5 KiB
Go
233 lines
7.5 KiB
Go
// Package app 组合根：装配存储（嵌入式/PG 档）、总线、注册表、服务、runmgr、
|
||
// 认证、插件宿主与 HTTP 服务。
|
||
package app
|
||
|
||
import (
|
||
"context"
|
||
"encoding/json"
|
||
"fmt"
|
||
"log"
|
||
"net/http"
|
||
"os"
|
||
"os/signal"
|
||
"path/filepath"
|
||
"sync"
|
||
"syscall"
|
||
"time"
|
||
|
||
"agent-system-backend/internal/bus"
|
||
"agent-system-backend/internal/intentgate"
|
||
"agent-system-backend/internal/pgbus"
|
||
"agent-system-backend/internal/plugins"
|
||
"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
|
||
StorageDSN string // 空 = 嵌入式 SQLite 档；非空 = PostgreSQL 档（多用户/多实例）
|
||
ArtifactsDir string // PG 档工件对象目录（空 = <DataDir>/artifacts）
|
||
AuthMode string // auto（默认：PG 档开、嵌入式关）| on | off
|
||
}
|
||
|
||
// Run 装配并阻塞服务直至 SIGINT/SIGTERM（优雅停机：中断在途 run、冲净事件缓冲）。
|
||
func Run(opt Options) error {
|
||
var st *store.Store
|
||
var err error
|
||
if opt.StorageDSN != "" {
|
||
artifacts := opt.ArtifactsDir
|
||
if artifacts == "" {
|
||
artifacts = filepath.Join(opt.DataDir, "artifacts")
|
||
}
|
||
st, err = store.OpenPostgres(opt.StorageDSN, artifacts)
|
||
} else {
|
||
st, err = store.OpenSQLite(filepath.Join(opt.DataDir, "agent-system.db"))
|
||
}
|
||
if err != nil {
|
||
return fmt.Errorf("open store: %w", err)
|
||
}
|
||
defer st.Close()
|
||
log.Printf("storage tier: %s", st.Dialect())
|
||
|
||
b := bus.New()
|
||
reg := tools.NewRegistry()
|
||
if err := reg.RegisterBuiltin(); err != nil {
|
||
return fmt.Errorf("builtin tools: %w", err)
|
||
}
|
||
setSvc := settings.NewService(st)
|
||
setSvc.SetSystemOverride(map[string]any{"storage": map[string]any{"dataDir": opt.DataDir}})
|
||
profSvc := profiles.New(st)
|
||
if err := profSvc.EnsureDefault(); err != nil {
|
||
return fmt.Errorf("default profile: %w", err)
|
||
}
|
||
if opt.ForceMock {
|
||
if err := setSvc.MergeUser(store.OwnerDefault, map[string]any{"runtime": map[string]any{"mock": true}}); err != nil {
|
||
return err
|
||
}
|
||
}
|
||
|
||
convRunner := &runtime.LLMRunner{}
|
||
subRunner := &runtime.SubLLMRunner{}
|
||
mockRunner := &runtime.MockRunner{}
|
||
mgr := runmgr.New(st, b, reg, profSvc, setSvc, convRunner, subRunner, mockRunner, &runtime.MockSubagent{})
|
||
|
||
// 意图闸门：插件意图 overlay → agent-intent 注册表（未装载意图 = 行为回退）。
|
||
gate := intentgate.New()
|
||
mgr.SetGate(gate)
|
||
|
||
pluginHost := plugins.New(st, reg, setSvc, opt.PluginsDir)
|
||
pluginHost.OnIntents = func(overlay []json.RawMessage) error { return gate.Rebuild(context.Background(), overlay) }
|
||
for id, err := range pluginHost.LoadDir() {
|
||
if err != nil {
|
||
log.Printf("[plugins] %s: %v", id, err) // fail-closed：单个失败不阻断启动
|
||
}
|
||
}
|
||
|
||
// 检索管线（4.12）：分块+全文+（配置时）向量混合与 rerank；未配置自动退化为全文。
|
||
pipe := rag.NewPipeline(st, setSvc, nil, nil)
|
||
pipe.SetEmbedder(rag.NewHTTPEmbedder(pipe.EmbeddingRoute))
|
||
pipe.SetReranker(rag.NewHTTPReranker(pipe.RerankRoute))
|
||
mgr.SetRetriever(pipe)
|
||
mgr.SetIndexer(pipe)
|
||
|
||
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)
|
||
if authOn {
|
||
log.Printf("auth: API key required (create with: agent-system apikey create --owner <name>)")
|
||
}
|
||
srv.Subscribe = func(conversationID string) (<-chan core.StreamEvent, func()) {
|
||
return b.Subscribe(conversationID)
|
||
}
|
||
|
||
httpSrv := &http.Server{Addr: opt.Addr, Handler: srv.Handler(), ReadHeaderTimeout: 10 * time.Second}
|
||
|
||
go func() {
|
||
log.Printf("agent-system serving on http://%s (data: %s, plugins: %s)", opt.Addr, opt.DataDir, opt.PluginsDir)
|
||
if err := httpSrv.ListenAndServe(); err != nil && err != http.ErrServerClosed {
|
||
log.Fatalf("http: %v", err)
|
||
}
|
||
}()
|
||
|
||
// 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
|
||
log.Printf("shutting down…")
|
||
mgr.InterruptRunning()
|
||
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
||
defer cancel()
|
||
if err := httpSrv.Shutdown(ctx); err != nil {
|
||
return err
|
||
}
|
||
mgr.Shutdown() // 冲净事件缓冲后再关存储
|
||
return nil
|
||
}
|
||
|
||
// CreateAPIKey 是 apikey 子命令的实现：在指定存储档上为 owner 创建密钥（明文仅打印一次）。
|
||
func CreateAPIKey(storageDSN, dataDir, owner, name string) (string, error) {
|
||
var st *store.Store
|
||
var err error
|
||
if storageDSN != "" {
|
||
st, err = store.OpenPostgres(storageDSN, filepath.Join(dataDir, "artifacts"))
|
||
} else {
|
||
st, err = store.OpenSQLite(filepath.Join(dataDir, "agent-system.db"))
|
||
}
|
||
if err != nil {
|
||
return "", err
|
||
}
|
||
defer st.Close()
|
||
return st.CreateAPIKey(owner, name)
|
||
}
|
||
|
||
// ListAPIKeys 列出租户密钥（摘要视图）。
|
||
func ListAPIKeys(storageDSN, dataDir, owner string) ([]store.APIKeyInfo, error) {
|
||
var st *store.Store
|
||
var err error
|
||
if storageDSN != "" {
|
||
st, err = store.OpenPostgres(storageDSN, filepath.Join(dataDir, "artifacts"))
|
||
} else {
|
||
st, err = store.OpenSQLite(filepath.Join(dataDir, "agent-system.db"))
|
||
}
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
defer st.Close()
|
||
if _, err := st.GetUser(owner); err != nil {
|
||
return nil, fmt.Errorf("owner %q 不存在（先 apikey create 创建）", owner)
|
||
}
|
||
return st.ListAPIKeys(owner)
|
||
}
|
||
|
||
// RevokeAPIKey 撤销指定名称密钥。
|
||
func RevokeAPIKey(storageDSN, dataDir, owner, name string) (bool, error) {
|
||
var st *store.Store
|
||
var err error
|
||
if storageDSN != "" {
|
||
st, err = store.OpenPostgres(storageDSN, filepath.Join(dataDir, "artifacts"))
|
||
} else {
|
||
st, err = store.OpenSQLite(filepath.Join(dataDir, "agent-system.db"))
|
||
}
|
||
if err != nil {
|
||
return false, err
|
||
}
|
||
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
|
||
}
|
||
}
|