Files
ai_site/platform/internal/agentstore/store.go
2026-07-31 10:31:17 +08:00

689 lines
20 KiB
Go

package agentstore
import (
"context"
"crypto/rand"
"database/sql"
"encoding/hex"
"errors"
"fmt"
"strings"
"sync"
"time"
"golang.org/x/crypto/bcrypt"
"aijianzhan/platform/internal/authx"
)
const (
StatusActive = "active"
StatusPending = "pending"
StatusDisabled = "disabled"
)
type Account struct {
AgentID int64 `json:"agent_id"`
TenantID int64 `json:"tenant_id"`
Name string `json:"name"`
ClientID string `json:"client_id"`
HostKey string `json:"host_key,omitempty"`
RoleID int64 `json:"role_id,omitempty"`
RoleCode string `json:"role_code,omitempty"`
RoleName string `json:"role_name,omitempty"`
Status string `json:"status"`
Perms []string `json:"permissions"`
AppSlugs []string `json:"app_slugs"`
CreatedBy int64 `json:"created_by"`
CreatedAt time.Time `json:"created_at"`
LastTokenAt *time.Time `json:"last_token_at,omitempty"`
}
type CreateInput struct {
Name string
Perms []string
AppSlugs []string
Status string // 空则 active
HostKey string
RoleID int64
}
type UpdateInput struct {
Name *string
Status *string
Perms *[]string
AppSlugs *[]string
RoleID *int64
}
type Store interface {
List(ctx context.Context, tenantID int64) ([]Account, error)
Get(ctx context.Context, tenantID, agentID int64) (*Account, error)
Create(ctx context.Context, tenantID, createdBy int64, in CreateInput) (acc *Account, clientSecret string, err error)
Register(ctx context.Context, tenantID int64, name, hostKey string) (acc *Account, clientSecret string, reused bool, err error)
Update(ctx context.Context, tenantID, agentID int64, in UpdateInput) (*Account, error)
RotateSecret(ctx context.Context, tenantID, agentID int64) (clientSecret string, err error)
Delete(ctx context.Context, tenantID, agentID int64) error
Authenticate(ctx context.Context, clientID, clientSecret string) (*Account, error)
TouchToken(ctx context.Context, agentID int64) error
HasAppAccess(ctx context.Context, agentID int64, slug string) (bool, error)
// GrantAppSlug 将 slug 写入智能体可访问模块(幂等)。新建发布时自动授权用。
GrantAppSlug(ctx context.Context, agentID int64, slug string) error
}
type memAcc struct {
Account
SecretHash string
}
type MemoryStore struct {
mu sync.Mutex
seq int64
byID map[int64]*memAcc
byCID map[string]*memAcc
}
func NewMemoryStore() *MemoryStore {
return &MemoryStore{byID: map[int64]*memAcc{}, byCID: map[string]*memAcc{}}
}
func (s *MemoryStore) List(_ context.Context, tenantID int64) ([]Account, error) {
s.mu.Lock()
defer s.mu.Unlock()
out := []Account{}
for _, a := range s.byID {
if a.TenantID == tenantID {
out = append(out, cloneAcc(&a.Account))
}
}
return out, nil
}
func (s *MemoryStore) Get(_ context.Context, tenantID, agentID int64) (*Account, error) {
s.mu.Lock()
defer s.mu.Unlock()
a, ok := s.byID[agentID]
if !ok || a.TenantID != tenantID {
return nil, fmt.Errorf("agent not found")
}
cp := cloneAcc(&a.Account)
return &cp, nil
}
func (s *MemoryStore) Create(_ context.Context, tenantID, createdBy int64, in CreateInput) (*Account, string, error) {
s.mu.Lock()
defer s.mu.Unlock()
return s.createLocked(tenantID, createdBy, in)
}
func (s *MemoryStore) Register(_ context.Context, tenantID int64, name, hostKey string) (*Account, string, bool, error) {
s.mu.Lock()
defer s.mu.Unlock()
hostKey = strings.TrimSpace(hostKey)
if hostKey != "" {
for _, a := range s.byID {
if a.TenantID == tenantID && a.HostKey == hostKey {
if a.Status == StatusPending {
secret := randomHex(24)
hash, err := bcrypt.GenerateFromPassword([]byte(secret), bcrypt.DefaultCost)
if err != nil {
return nil, "", false, err
}
a.SecretHash = string(hash)
if n := strings.TrimSpace(name); n != "" {
a.Name = n
}
cp := cloneAcc(&a.Account)
return &cp, secret, true, nil
}
return nil, "", false, fmt.Errorf("host already registered as %s (status=%s)", a.ClientID, a.Status)
}
}
}
acc, secret, err := s.createLocked(tenantID, 0, CreateInput{
Name: name,
Status: StatusPending,
HostKey: hostKey,
})
return acc, secret, false, err
}
func (s *MemoryStore) createLocked(tenantID, createdBy int64, in CreateInput) (*Account, string, error) {
cid := "agt_" + randomHex(8)
secret := randomHex(24)
hash, err := bcrypt.GenerateFromPassword([]byte(secret), bcrypt.DefaultCost)
if err != nil {
return nil, "", err
}
st := strings.TrimSpace(in.Status)
if st == "" {
st = StatusActive
}
if !validStatus(st) {
return nil, "", fmt.Errorf("invalid status")
}
s.seq++
a := &memAcc{
Account: Account{
AgentID: s.seq,
TenantID: tenantID,
Name: strings.TrimSpace(in.Name),
ClientID: cid,
HostKey: strings.TrimSpace(in.HostKey),
RoleID: in.RoleID,
Status: st,
Perms: authx.NormalizePerms(uniq(in.Perms)),
AppSlugs: uniq(in.AppSlugs),
CreatedBy: createdBy,
CreatedAt: time.Now().UTC(),
},
SecretHash: string(hash),
}
if a.Name == "" {
a.Name = cid
}
s.byID[a.AgentID] = a
s.byCID[a.ClientID] = a
cp := cloneAcc(&a.Account)
return &cp, secret, nil
}
func (s *MemoryStore) Update(_ context.Context, tenantID, agentID int64, in UpdateInput) (*Account, error) {
s.mu.Lock()
defer s.mu.Unlock()
a, ok := s.byID[agentID]
if !ok || a.TenantID != tenantID {
return nil, fmt.Errorf("agent not found")
}
if in.Name != nil {
a.Name = strings.TrimSpace(*in.Name)
}
if in.Status != nil {
st := strings.TrimSpace(*in.Status)
if !validStatus(st) {
return nil, fmt.Errorf("invalid status")
}
a.Status = st
}
if in.Perms != nil {
a.Perms = uniq(*in.Perms)
}
if in.AppSlugs != nil {
a.AppSlugs = uniq(*in.AppSlugs)
}
if in.RoleID != nil {
a.RoleID = *in.RoleID
}
cp := cloneAcc(&a.Account)
return &cp, nil
}
func (s *MemoryStore) RotateSecret(_ context.Context, tenantID, agentID int64) (string, error) {
s.mu.Lock()
defer s.mu.Unlock()
a, ok := s.byID[agentID]
if !ok || a.TenantID != tenantID {
return "", fmt.Errorf("agent not found")
}
secret := randomHex(24)
hash, err := bcrypt.GenerateFromPassword([]byte(secret), bcrypt.DefaultCost)
if err != nil {
return "", err
}
a.SecretHash = string(hash)
return secret, nil
}
func (s *MemoryStore) Delete(_ context.Context, tenantID, agentID int64) error {
s.mu.Lock()
defer s.mu.Unlock()
a, ok := s.byID[agentID]
if !ok || a.TenantID != tenantID {
return fmt.Errorf("agent not found")
}
delete(s.byCID, a.ClientID)
delete(s.byID, agentID)
return nil
}
func (s *MemoryStore) Authenticate(_ context.Context, clientID, clientSecret string) (*Account, error) {
s.mu.Lock()
defer s.mu.Unlock()
a, ok := s.byCID[clientID]
if !ok {
return nil, fmt.Errorf("invalid client credentials")
}
if bcrypt.CompareHashAndPassword([]byte(a.SecretHash), []byte(clientSecret)) != nil {
return nil, fmt.Errorf("invalid client credentials")
}
switch a.Status {
case StatusPending:
return nil, fmt.Errorf("agent pending approval")
case StatusDisabled:
return nil, fmt.Errorf("agent disabled")
case StatusActive:
// ok
default:
return nil, fmt.Errorf("agent not active")
}
cp := cloneAcc(&a.Account)
return &cp, nil
}
func (s *MemoryStore) TouchToken(_ context.Context, agentID int64) error {
s.mu.Lock()
defer s.mu.Unlock()
if a, ok := s.byID[agentID]; ok {
now := time.Now().UTC()
a.LastTokenAt = &now
}
return nil
}
func (s *MemoryStore) HasAppAccess(_ context.Context, agentID int64, slug string) (bool, error) {
s.mu.Lock()
defer s.mu.Unlock()
a, ok := s.byID[agentID]
if !ok {
return false, fmt.Errorf("agent not found")
}
// 未配置可访问模块 = 不限制,智能体可自由发布/访问自建模块
if len(a.AppSlugs) == 0 {
return true, nil
}
for _, s0 := range a.AppSlugs {
if s0 == "*" || s0 == slug {
return true, nil
}
}
return false, nil
}
func (s *MemoryStore) GrantAppSlug(_ context.Context, agentID int64, slug string) error {
slug = strings.TrimSpace(slug)
if slug == "" {
return fmt.Errorf("empty slug")
}
s.mu.Lock()
defer s.mu.Unlock()
a, ok := s.byID[agentID]
if !ok {
return fmt.Errorf("agent not found")
}
for _, s0 := range a.AppSlugs {
if s0 == slug {
return nil
}
}
a.AppSlugs = append(a.AppSlugs, slug)
return nil
}
type PostgresStore struct {
DB *sql.DB
}
func NewPostgresStore(db *sql.DB) *PostgresStore { return &PostgresStore{DB: db} }
func (s *PostgresStore) List(ctx context.Context, tenantID int64) ([]Account, error) {
rows, err := s.DB.QueryContext(ctx, `
SELECT agent_id, tenant_id, name, client_id, COALESCE(host_key,''), COALESCE(role_id,0), status, created_by, created_at, last_token_at
FROM platform_meta.agent_accounts WHERE tenant_id=$1
ORDER BY CASE status WHEN 'pending' THEN 0 WHEN 'active' THEN 1 ELSE 2 END, agent_id`, tenantID)
if err != nil {
return nil, err
}
defer rows.Close()
var out []Account
for rows.Next() {
var a Account
var last sql.NullTime
if err := rows.Scan(&a.AgentID, &a.TenantID, &a.Name, &a.ClientID, &a.HostKey, &a.RoleID, &a.Status, &a.CreatedBy, &a.CreatedAt, &last); err != nil {
return nil, err
}
if last.Valid {
t := last.Time.UTC()
a.LastTokenAt = &t
}
perms, slugs, err := s.loadKids(ctx, a.AgentID)
if err != nil {
return nil, err
}
a.Perms, a.AppSlugs = perms, slugs
out = append(out, a)
}
return out, rows.Err()
}
func (s *PostgresStore) Get(ctx context.Context, tenantID, agentID int64) (*Account, error) {
var a Account
var last sql.NullTime
err := s.DB.QueryRowContext(ctx, `
SELECT agent_id, tenant_id, name, client_id, COALESCE(host_key,''), COALESCE(role_id,0), status, created_by, created_at, last_token_at
FROM platform_meta.agent_accounts WHERE agent_id=$1 AND tenant_id=$2`, agentID, tenantID,
).Scan(&a.AgentID, &a.TenantID, &a.Name, &a.ClientID, &a.HostKey, &a.RoleID, &a.Status, &a.CreatedBy, &a.CreatedAt, &last)
if errors.Is(err, sql.ErrNoRows) {
return nil, fmt.Errorf("agent not found")
}
if err != nil {
return nil, err
}
if last.Valid {
t := last.Time.UTC()
a.LastTokenAt = &t
}
perms, slugs, err := s.loadKids(ctx, a.AgentID)
if err != nil {
return nil, err
}
a.Perms, a.AppSlugs = perms, slugs
return &a, nil
}
func (s *PostgresStore) Create(ctx context.Context, tenantID, createdBy int64, in CreateInput) (*Account, string, error) {
return s.insert(ctx, tenantID, createdBy, in)
}
func (s *PostgresStore) Register(ctx context.Context, tenantID int64, name, hostKey string) (*Account, string, bool, error) {
hostKey = strings.TrimSpace(hostKey)
if hostKey != "" {
var a Account
var last sql.NullTime
err := s.DB.QueryRowContext(ctx, `
SELECT agent_id, tenant_id, name, client_id, COALESCE(host_key,''), COALESCE(role_id,0), status, created_by, created_at, last_token_at
FROM platform_meta.agent_accounts WHERE tenant_id=$1 AND host_key=$2`, tenantID, hostKey,
).Scan(&a.AgentID, &a.TenantID, &a.Name, &a.ClientID, &a.HostKey, &a.RoleID, &a.Status, &a.CreatedBy, &a.CreatedAt, &last)
if err == nil {
if a.Status == StatusPending {
secret, err := s.RotateSecret(ctx, tenantID, a.AgentID)
if err != nil {
return nil, "", false, err
}
if n := strings.TrimSpace(name); n != "" {
_, _ = s.Update(ctx, tenantID, a.AgentID, UpdateInput{Name: &n})
}
acc, err := s.Get(ctx, tenantID, a.AgentID)
return acc, secret, true, err
}
return nil, "", false, fmt.Errorf("host already registered as %s (status=%s)", a.ClientID, a.Status)
}
if !errors.Is(err, sql.ErrNoRows) {
return nil, "", false, err
}
}
acc, secret, err := s.insert(ctx, tenantID, 0, CreateInput{
Name: name,
Status: StatusPending,
HostKey: hostKey,
})
return acc, secret, false, err
}
func (s *PostgresStore) insert(ctx context.Context, tenantID, createdBy int64, in CreateInput) (*Account, string, error) {
cid := "agt_" + randomHex(8)
secret := randomHex(24)
hash, err := bcrypt.GenerateFromPassword([]byte(secret), bcrypt.DefaultCost)
if err != nil {
return nil, "", err
}
name := strings.TrimSpace(in.Name)
if name == "" {
name = cid
}
st := strings.TrimSpace(in.Status)
if st == "" {
st = StatusActive
}
if !validStatus(st) {
return nil, "", fmt.Errorf("invalid status")
}
var a Account
err = s.DB.QueryRowContext(ctx, `
INSERT INTO platform_meta.agent_accounts(tenant_id, name, client_id, client_secret_hash, status, created_by, host_key, role_id)
VALUES($1,$2,$3,$4,$5,$6,$7,NULLIF($8,0))
RETURNING agent_id, tenant_id, name, client_id, COALESCE(host_key,''), COALESCE(role_id,0), status, created_by, created_at`,
tenantID, name, cid, string(hash), st, createdBy, strings.TrimSpace(in.HostKey), in.RoleID,
).Scan(&a.AgentID, &a.TenantID, &a.Name, &a.ClientID, &a.HostKey, &a.RoleID, &a.Status, &a.CreatedBy, &a.CreatedAt)
if err != nil {
return nil, "", err
}
perms, slugs := uniq(in.Perms), uniq(in.AppSlugs)
if err := s.replaceKids(ctx, a.AgentID, perms, slugs); err != nil {
return nil, "", err
}
a.Perms, a.AppSlugs = perms, slugs
return &a, secret, nil
}
func (s *PostgresStore) Update(ctx context.Context, tenantID, agentID int64, in UpdateInput) (*Account, error) {
a, err := s.Get(ctx, tenantID, agentID)
if err != nil {
return nil, err
}
name, status := a.Name, a.Status
roleID := a.RoleID
if in.Name != nil {
name = strings.TrimSpace(*in.Name)
}
if in.Status != nil {
status = strings.TrimSpace(*in.Status)
if !validStatus(status) {
return nil, fmt.Errorf("invalid status")
}
}
if in.RoleID != nil {
roleID = *in.RoleID
}
if _, err := s.DB.ExecContext(ctx, `
UPDATE platform_meta.agent_accounts SET name=$1, status=$2, role_id=NULLIF($3,0) WHERE agent_id=$4 AND tenant_id=$5`,
name, status, roleID, agentID, tenantID); err != nil {
return nil, err
}
perms, slugs := a.Perms, a.AppSlugs
if in.Perms != nil {
perms = uniq(*in.Perms)
}
if in.AppSlugs != nil {
slugs = uniq(*in.AppSlugs)
}
if err := s.replaceKids(ctx, agentID, perms, slugs); err != nil {
return nil, err
}
return s.Get(ctx, tenantID, agentID)
}
func (s *PostgresStore) RotateSecret(ctx context.Context, tenantID, agentID int64) (string, error) {
if _, err := s.Get(ctx, tenantID, agentID); err != nil {
return "", err
}
secret := randomHex(24)
hash, err := bcrypt.GenerateFromPassword([]byte(secret), bcrypt.DefaultCost)
if err != nil {
return "", err
}
_, err = s.DB.ExecContext(ctx, `
UPDATE platform_meta.agent_accounts SET client_secret_hash=$1 WHERE agent_id=$2 AND tenant_id=$3`,
string(hash), agentID, tenantID)
return secret, err
}
func (s *PostgresStore) Delete(ctx context.Context, tenantID, agentID int64) error {
res, err := s.DB.ExecContext(ctx, `
DELETE FROM platform_meta.agent_accounts WHERE agent_id=$1 AND tenant_id=$2`, agentID, tenantID)
if err != nil {
return err
}
n, _ := res.RowsAffected()
if n == 0 {
return fmt.Errorf("agent not found")
}
return nil
}
func (s *PostgresStore) Authenticate(ctx context.Context, clientID, clientSecret string) (*Account, error) {
var a Account
var hash string
var last sql.NullTime
err := s.DB.QueryRowContext(ctx, `
SELECT agent_id, tenant_id, name, client_id, COALESCE(host_key,''), COALESCE(role_id,0), client_secret_hash, status, created_by, created_at, last_token_at
FROM platform_meta.agent_accounts WHERE client_id=$1`, clientID,
).Scan(&a.AgentID, &a.TenantID, &a.Name, &a.ClientID, &a.HostKey, &a.RoleID, &hash, &a.Status, &a.CreatedBy, &a.CreatedAt, &last)
if errors.Is(err, sql.ErrNoRows) {
return nil, fmt.Errorf("invalid client credentials")
}
if err != nil {
return nil, err
}
if bcrypt.CompareHashAndPassword([]byte(hash), []byte(clientSecret)) != nil {
return nil, fmt.Errorf("invalid client credentials")
}
switch a.Status {
case StatusPending:
return nil, fmt.Errorf("agent pending approval")
case StatusDisabled:
return nil, fmt.Errorf("agent disabled")
case StatusActive:
default:
return nil, fmt.Errorf("agent not active")
}
if last.Valid {
t := last.Time.UTC()
a.LastTokenAt = &t
}
perms, slugs, err := s.loadKids(ctx, a.AgentID)
if err != nil {
return nil, err
}
a.Perms, a.AppSlugs = perms, slugs
return &a, nil
}
func (s *PostgresStore) TouchToken(ctx context.Context, agentID int64) error {
_, err := s.DB.ExecContext(ctx, `
UPDATE platform_meta.agent_accounts SET last_token_at=now() WHERE agent_id=$1`, agentID)
return err
}
func (s *PostgresStore) HasAppAccess(ctx context.Context, agentID int64, slug string) (bool, error) {
var total int
if err := s.DB.QueryRowContext(ctx, `
SELECT COUNT(1) FROM platform_meta.agent_app_grants WHERE agent_id=$1`, agentID).Scan(&total); err != nil {
return false, err
}
// 未配置可访问模块 = 不限制,可自由发布自己的模块
if total == 0 {
return true, nil
}
var n int
err := s.DB.QueryRowContext(ctx, `
SELECT COUNT(1) FROM platform_meta.agent_app_grants
WHERE agent_id=$1 AND (slug=$2 OR slug='*')`, agentID, slug).Scan(&n)
return n > 0, err
}
func (s *PostgresStore) GrantAppSlug(ctx context.Context, agentID int64, slug string) error {
slug = strings.TrimSpace(slug)
if slug == "" {
return fmt.Errorf("empty slug")
}
ok, err := s.HasAppAccess(ctx, agentID, slug)
if err != nil {
return err
}
if ok {
return nil
}
_, err = s.DB.ExecContext(ctx, `
INSERT INTO platform_meta.agent_app_grants(agent_id, slug) VALUES($1,$2)`, agentID, slug)
return err
}
func (s *PostgresStore) loadKids(ctx context.Context, agentID int64) ([]string, []string, error) {
prows, err := s.DB.QueryContext(ctx, `SELECT perm FROM platform_meta.agent_permissions WHERE agent_id=$1`, agentID)
if err != nil {
return nil, nil, err
}
defer prows.Close()
var perms []string
for prows.Next() {
var p string
if err := prows.Scan(&p); err != nil {
return nil, nil, err
}
perms = append(perms, p)
}
srows, err := s.DB.QueryContext(ctx, `SELECT slug FROM platform_meta.agent_app_grants WHERE agent_id=$1`, agentID)
if err != nil {
return nil, nil, err
}
defer srows.Close()
var slugs []string
for srows.Next() {
var slug string
if err := srows.Scan(&slug); err != nil {
return nil, nil, err
}
slugs = append(slugs, slug)
}
return authx.NormalizePerms(perms), slugs, nil
}
func (s *PostgresStore) replaceKids(ctx context.Context, agentID int64, perms, slugs []string) error {
tx, err := s.DB.BeginTx(ctx, nil)
if err != nil {
return err
}
defer func() { _ = tx.Rollback() }()
if _, err := tx.ExecContext(ctx, `DELETE FROM platform_meta.agent_permissions WHERE agent_id=$1`, agentID); err != nil {
return err
}
if _, err := tx.ExecContext(ctx, `DELETE FROM platform_meta.agent_app_grants WHERE agent_id=$1`, agentID); err != nil {
return err
}
for _, p := range authx.NormalizePerms(perms) {
if _, err := tx.ExecContext(ctx, `INSERT INTO platform_meta.agent_permissions(agent_id, perm) VALUES($1,$2)`, agentID, p); err != nil {
return err
}
}
for _, slug := range slugs {
if _, err := tx.ExecContext(ctx, `INSERT INTO platform_meta.agent_app_grants(agent_id, slug) VALUES($1,$2)`, agentID, slug); err != nil {
return err
}
}
return tx.Commit()
}
func validStatus(st string) bool {
return st == StatusActive || st == StatusPending || st == StatusDisabled
}
func cloneAcc(a *Account) Account {
cp := *a
cp.Perms = append([]string{}, a.Perms...)
cp.AppSlugs = append([]string{}, a.AppSlugs...)
return cp
}
func uniq(in []string) []string {
seen := map[string]struct{}{}
out := make([]string, 0, len(in))
for _, s := range in {
s = strings.TrimSpace(s)
if s == "" {
continue
}
if _, ok := seen[s]; ok {
continue
}
seen[s] = struct{}{}
out = append(out, s)
}
return out
}
func randomHex(n int) string {
b := make([]byte, n)
_, _ = rand.Read(b)
return hex.EncodeToString(b)
}