root/geo
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: add tenant and user management with migrations, handlers, and tests

Browse Source 
- Implemented tenant and user management features including:
  - Tenant creation and management with associated migrations.
  - User creation and management with associated migrations.
  - Tenant membership management with associated migrations.
  - Platform user roles management with associated migrations.
  - Quota management with associated migrations.
  - Article and template management with associated migrations.
- Added HTTP handlers for templates and workspaces.
- Created tests for protected and public routes.
- Introduced a script to check tenant scope in SQL queries.
- Documented task plan for backend completion and frontend foundation.
This commit is contained in:
liangxu
2026-04-01 00:58:42 +08:00
commit de30497f59
210 changed files with 23733 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
server/internal/shared/auth/claims.go
+12
View File
@@ -0,0 +1,12 @@
package auth
import (
"github.com/golang-jwt/jwt/v5"
)
type Claims struct {
UserID int64 `json:"user_id"`
TenantID int64 `json:"tenant_id"`
Role string `json:"role"`
jwt.RegisteredClaims
}
server/internal/shared/auth/context.go
+28
View File
@@ -0,0 +1,28 @@
package auth
import "context"
type Actor struct {
UserID int64
TenantID int64
Role string
}
type actorKey struct{}
func WithActor(ctx context.Context, actor Actor) context.Context {
return context.WithValue(ctx, actorKey{}, actor)
}
func ActorFromCtx(ctx context.Context) (Actor, bool) {
actor, ok := ctx.Value(actorKey{}).(Actor)
return actor, ok
}
func MustActor(ctx context.Context) Actor {
actor, ok := ActorFromCtx(ctx)
if !ok {
panic("auth: actor not in context")
}
return actor
}
server/internal/shared/auth/context_test.go
+36
View File
@@ -0,0 +1,36 @@
package auth
import (
"context"
"testing"
"github.com/stretchr/testify/assert"
)
func TestWithActorAndFromCtx(t *testing.T) {
actor := Actor{UserID: 1, TenantID: 10, Role: "editor"}
ctx := WithActor(context.Background(), actor)
got, ok := ActorFromCtx(ctx)
assert.True(t, ok)
assert.Equal(t, actor, got)
}
func TestActorFromCtxMissing(t *testing.T) {
_, ok := ActorFromCtx(context.Background())
assert.False(t, ok)
}
func TestMustActorPanics(t *testing.T) {
assert.Panics(t, func() {
MustActor(context.Background())
})
}
func TestMustActorReturns(t *testing.T) {
actor := Actor{UserID: 5, TenantID: 20, Role: "tenant_admin"}
ctx := WithActor(context.Background(), actor)
got := MustActor(ctx)
assert.Equal(t, actor, got)
}
server/internal/shared/auth/jwt.go
+103
View File
@@ -0,0 +1,103 @@
package auth
import (
"crypto/sha256"
"fmt"
"time"
"github.com/golang-jwt/jwt/v5"
"github.com/google/uuid"
)
type Manager struct {
secret []byte
accessTTL time.Duration
refreshTTL time.Duration
}
func NewManager(secret string, accessTTL, refreshTTL time.Duration) *Manager {
return &Manager{
secret: []byte(secret),
accessTTL: accessTTL,
refreshTTL: refreshTTL,
}
}
type TokenPair struct {
AccessToken string `json:"access_token"`
RefreshToken string `json:"refresh_token"`
AccessJTI string `json:"-"`
RefreshJTI string `json:"-"`
ExpiresAt int64 `json:"expires_at"`
}
func (m *Manager) Issue(userID, tenantID int64, role string) (*TokenPair, error) {
now := time.Now()
accessJTI := uuid.New().String()
refreshJTI := uuid.New().String()
accessClaims := Claims{
UserID: userID,
TenantID: tenantID,
Role: role,
RegisteredClaims: jwt.RegisteredClaims{
ID: accessJTI,
Subject: "access",
IssuedAt: jwt.NewNumericDate(now),
ExpiresAt: jwt.NewNumericDate(now.Add(m.accessTTL)),
},
}
accessToken, err := jwt.NewWithClaims(jwt.SigningMethodHS256, accessClaims).SignedString(m.secret)
if err != nil {
return nil, fmt.Errorf("sign access token: %w", err)
}
refreshClaims := Claims{
UserID: userID,
TenantID: tenantID,
Role: role,
RegisteredClaims: jwt.RegisteredClaims{
ID: refreshJTI,
Subject: "refresh",
IssuedAt: jwt.NewNumericDate(now),
ExpiresAt: jwt.NewNumericDate(now.Add(m.refreshTTL)),
},
}
refreshToken, err := jwt.NewWithClaims(jwt.SigningMethodHS256, refreshClaims).SignedString(m.secret)
if err != nil {
return nil, fmt.Errorf("sign refresh token: %w", err)
}
return &TokenPair{
AccessToken: accessToken,
RefreshToken: refreshToken,
AccessJTI: accessJTI,
RefreshJTI: refreshJTI,
ExpiresAt: now.Add(m.accessTTL).Unix(),
}, nil
}
func (m *Manager) Parse(tokenStr string) (*Claims, error) {
token, err := jwt.ParseWithClaims(tokenStr, &Claims{}, func(t *jwt.Token) (interface{}, error) {
if _, ok := t.Method.(*jwt.SigningMethodHMAC); !ok {
return nil, fmt.Errorf("unexpected signing method: %v", t.Header["alg"])
}
return m.secret, nil
})
if err != nil {
return nil, err
}
claims, ok := token.Claims.(*Claims)
if !ok || !token.Valid {
return nil, fmt.Errorf("invalid token claims")
}
return claims, nil
}
func (m *Manager) AccessTTL() time.Duration { return m.accessTTL }
func (m *Manager) RefreshTTL() time.Duration { return m.refreshTTL }
func HashToken(raw string) string {
h := sha256.Sum256([]byte(raw))
return fmt.Sprintf("%x", h)
}
server/internal/shared/auth/jwt_test.go
+85
View File
@@ -0,0 +1,85 @@
package auth
import (
"testing"
"time"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
func TestIssueAndParse(t *testing.T) {
mgr := NewManager("test-secret-key", 15*time.Minute, 7*24*time.Hour)
pair, err := mgr.Issue(42, 10, "tenant_admin")
require.NoError(t, err)
assert.NotEmpty(t, pair.AccessToken)
assert.NotEmpty(t, pair.RefreshToken)
assert.NotEmpty(t, pair.AccessJTI)
assert.NotEmpty(t, pair.RefreshJTI)
assert.NotZero(t, pair.ExpiresAt)
claims, err := mgr.Parse(pair.AccessToken)
require.NoError(t, err)
assert.Equal(t, int64(42), claims.UserID)
assert.Equal(t, int64(10), claims.TenantID)
assert.Equal(t, "tenant_admin", claims.Role)
assert.Equal(t, "access", claims.Subject)
assert.Equal(t, pair.AccessJTI, claims.ID)
}
func TestParseRefreshToken(t *testing.T) {
mgr := NewManager("test-secret-key", 15*time.Minute, 7*24*time.Hour)
pair, err := mgr.Issue(1, 1, "viewer")
require.NoError(t, err)
claims, err := mgr.Parse(pair.RefreshToken)
require.NoError(t, err)
assert.Equal(t, "refresh", claims.Subject)
assert.Equal(t, pair.RefreshJTI, claims.ID)
}
func TestParseExpired(t *testing.T) {
mgr := NewManager("test-secret-key", -1*time.Second, -1*time.Second)
pair, err := mgr.Issue(1, 1, "viewer")
require.NoError(t, err)
_, err = mgr.Parse(pair.AccessToken)
assert.Error(t, err)
}
func TestParseWrongSecret(t *testing.T) {
mgr1 := NewManager("secret-one", 15*time.Minute, 7*24*time.Hour)
mgr2 := NewManager("secret-two", 15*time.Minute, 7*24*time.Hour)
pair, err := mgr1.Issue(1, 1, "viewer")
require.NoError(t, err)
_, err = mgr2.Parse(pair.AccessToken)
assert.Error(t, err)
}
func TestParseGarbage(t *testing.T) {
mgr := NewManager("test-secret-key", 15*time.Minute, 7*24*time.Hour)
_, err := mgr.Parse("not-a-real-token")
assert.Error(t, err)
}
func TestHashToken(t *testing.T) {
h1 := HashToken("some-token")
h2 := HashToken("some-token")
h3 := HashToken("different-token")
assert.Equal(t, h1, h2)
assert.NotEqual(t, h1, h3)
assert.Len(t, h1, 64) // SHA-256 hex
}
func TestAccessAndRefreshTTL(t *testing.T) {
mgr := NewManager("s", 15*time.Minute, 24*time.Hour)
assert.Equal(t, 15*time.Minute, mgr.AccessTTL())
assert.Equal(t, 24*time.Hour, mgr.RefreshTTL())
}
server/internal/shared/auth/middleware.go
+50
View File
@@ -0,0 +1,50 @@
package auth
import (
"strings"
"github.com/gin-gonic/gin"
"github.com/geo-platform/tenant-api/internal/shared/response"
)
func Middleware(jwtMgr *Manager, sessions *SessionStore) gin.HandlerFunc {
return func(c *gin.Context) {
raw, err := bearerToken(c.GetHeader("Authorization"))
if err != nil {
response.Error(c, response.ErrUnauthorized(40101, "missing_bearer_token", "Authorization header is required"))
c.Abort()
return
}
claims, err := jwtMgr.Parse(raw)
if err != nil || claims.Subject != "access" {
response.Error(c, response.ErrUnauthorized(40102, "invalid_access_token", "token is invalid or expired"))
c.Abort()
return
}
revoked, _ := sessions.IsBlacklisted(c.Request.Context(), claims.ID)
if revoked {
response.Error(c, response.ErrUnauthorized(40103, "token_revoked", "token has been logged out"))
c.Abort()
return
}
actor := Actor{
UserID: claims.UserID,
TenantID: claims.TenantID,
Role: claims.Role,
}
c.Request = c.Request.WithContext(WithActor(c.Request.Context(), actor))
c.Next()
}
}
func bearerToken(header string) (string, error) {
parts := strings.SplitN(header, " ", 2)
if len(parts) != 2 || !strings.EqualFold(parts[0], "bearer") {
return "", response.ErrUnauthorized(40101, "missing_bearer_token", "Authorization header is required")
}
return parts[1], nil
}
server/internal/shared/auth/middleware_test.go
+124
View File
@@ -0,0 +1,124 @@
package auth
import (
"fmt"
"net/http"
"net/http/httptest"
"testing"
"time"
"github.com/alicebob/miniredis/v2"
"github.com/gin-gonic/gin"
goredis "github.com/redis/go-redis/v9"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
func init() {
gin.SetMode(gin.TestMode)
}
func setupTestMiddleware(t *testing.T) (*Manager, *SessionStore, *miniredis.Miniredis) {
t.Helper()
mr := miniredis.RunT(t)
rdb := goredis.NewClient(&goredis.Options{Addr: mr.Addr()})
mgr := NewManager("test-middleware-secret", 15*time.Minute, 7*24*time.Hour)
sessions := NewSessionStore(rdb)
return mgr, sessions, mr
}
func TestAuthMiddleware_ValidToken(t *testing.T) {
mgr, sessions, _ := setupTestMiddleware(t)
pair, err := mgr.Issue(1, 10, "editor")
require.NoError(t, err)
r := gin.New()
r.Use(Middleware(mgr, sessions))
var gotActor Actor
r.GET("/test", func(c *gin.Context) {
a, ok := ActorFromCtx(c.Request.Context())
assert.True(t, ok)
gotActor = a
c.Status(http.StatusOK)
})
w := httptest.NewRecorder()
req := httptest.NewRequest(http.MethodGet, "/test", nil)
req.Header.Set("Authorization", "Bearer "+pair.AccessToken)
r.ServeHTTP(w, req)
assert.Equal(t, http.StatusOK, w.Code)
assert.Equal(t, int64(1), gotActor.UserID)
assert.Equal(t, int64(10), gotActor.TenantID)
assert.Equal(t, "editor", gotActor.Role)
}
func TestAuthMiddleware_MissingHeader(t *testing.T) {
mgr, sessions, _ := setupTestMiddleware(t)
r := gin.New()
r.Use(Middleware(mgr, sessions))
r.GET("/test", func(c *gin.Context) { c.Status(http.StatusOK) })
w := httptest.NewRecorder()
req := httptest.NewRequest(http.MethodGet, "/test", nil)
r.ServeHTTP(w, req)
assert.Equal(t, http.StatusUnauthorized, w.Code)
}
func TestAuthMiddleware_InvalidToken(t *testing.T) {
mgr, sessions, _ := setupTestMiddleware(t)
r := gin.New()
r.Use(Middleware(mgr, sessions))
r.GET("/test", func(c *gin.Context) { c.Status(http.StatusOK) })
w := httptest.NewRecorder()
req := httptest.NewRequest(http.MethodGet, "/test", nil)
req.Header.Set("Authorization", "Bearer bad-token-value")
r.ServeHTTP(w, req)
assert.Equal(t, http.StatusUnauthorized, w.Code)
}
func TestAuthMiddleware_RefreshTokenRejected(t *testing.T) {
mgr, sessions, _ := setupTestMiddleware(t)
pair, err := mgr.Issue(1, 10, "editor")
require.NoError(t, err)
r := gin.New()
r.Use(Middleware(mgr, sessions))
r.GET("/test", func(c *gin.Context) { c.Status(http.StatusOK) })
w := httptest.NewRecorder()
req := httptest.NewRequest(http.MethodGet, "/test", nil)
req.Header.Set("Authorization", "Bearer "+pair.RefreshToken)
r.ServeHTTP(w, req)
assert.Equal(t, http.StatusUnauthorized, w.Code)
}
func TestAuthMiddleware_BlacklistedToken(t *testing.T) {
mgr, sessions, mr := setupTestMiddleware(t)
pair, err := mgr.Issue(1, 10, "editor")
require.NoError(t, err)
// Blacklist the access token's JTI
require.NoError(t, mr.Set(fmt.Sprintf("blacklist:%s", pair.AccessJTI), "1"))
r := gin.New()
r.Use(Middleware(mgr, sessions))
r.GET("/test", func(c *gin.Context) { c.Status(http.StatusOK) })
w := httptest.NewRecorder()
req := httptest.NewRequest(http.MethodGet, "/test", nil)
req.Header.Set("Authorization", "Bearer "+pair.AccessToken)
r.ServeHTTP(w, req)
assert.Equal(t, http.StatusUnauthorized, w.Code)
}
server/internal/shared/auth/permissions.go
+30
View File
@@ -0,0 +1,30 @@
package auth
var rolePermissions = map[string][]string{
"tenant_admin": {
"brand:read", "brand:write", "brand:delete",
"article:read", "article:write", "article:delete",
"template:read", "template:generate",
"workspace:read",
"member:read", "member:write",
},
"editor": {
"brand:read", "brand:write",
"article:read", "article:write",
"template:read", "template:generate",
"workspace:read",
},
"viewer": {
"brand:read",
"article:read",
"template:read",
"workspace:read",
},
}
func PermissionsForRole(role string) []string {
if perms, ok := rolePermissions[role]; ok {
return perms
}
return nil
}
server/internal/shared/auth/permissions_test.go
+35
View File
@@ -0,0 +1,35 @@
package auth
import (
"testing"
"github.com/stretchr/testify/assert"
)
func TestPermissionsForRole_Admin(t *testing.T) {
perms := PermissionsForRole("tenant_admin")
assert.Contains(t, perms, "brand:write")
assert.Contains(t, perms, "brand:delete")
assert.Contains(t, perms, "member:write")
}
func TestPermissionsForRole_Editor(t *testing.T) {
perms := PermissionsForRole("editor")
assert.Contains(t, perms, "article:write")
assert.Contains(t, perms, "template:generate")
assert.NotContains(t, perms, "brand:delete")
assert.NotContains(t, perms, "member:write")
}
func TestPermissionsForRole_Viewer(t *testing.T) {
perms := PermissionsForRole("viewer")
assert.Contains(t, perms, "brand:read")
assert.Contains(t, perms, "article:read")
assert.NotContains(t, perms, "brand:write")
assert.NotContains(t, perms, "article:write")
}
func TestPermissionsForRole_Unknown(t *testing.T) {
perms := PermissionsForRole("nonexistent")
assert.Nil(t, perms)
}
server/internal/shared/auth/session_store.go
+93
View File
@@ -0,0 +1,93 @@
package auth
import (
"context"
"errors"
"fmt"
"time"
"github.com/redis/go-redis/v9"
)
var (
ErrRefreshSessionNotFound = errors.New("refresh session not found")
ErrRefreshTokenMismatch = errors.New("refresh token mismatch")
)
var rotateRefreshScript = redis.NewScript(`
local current = redis.call("GET", KEYS[1])
if not current then
return 0
end
if current ~= ARGV[1] then
return -1
end
redis.call("DEL", KEYS[1])
redis.call("PSETEX", KEYS[2], ARGV[3], ARGV[2])
return 1
`)
type SessionStore struct {
rdb *redis.Client
}
func NewSessionStore(rdb *redis.Client) *SessionStore {
return &SessionStore{rdb: rdb}
}
func (s *SessionStore) SaveRefresh(ctx context.Context, jti string, tokenHash string, ttl time.Duration) error {
key := fmt.Sprintf("refresh:%s", jti)
return s.rdb.Set(ctx, key, tokenHash, ttl).Err()
}
func (s *SessionStore) GetRefresh(ctx context.Context, jti string) (string, error) {
key := fmt.Sprintf("refresh:%s", jti)
return s.rdb.Get(ctx, key).Result()
}
func (s *SessionStore) DeleteRefresh(ctx context.Context, jti string) error {
key := fmt.Sprintf("refresh:%s", jti)
return s.rdb.Del(ctx, key).Err()
}
func (s *SessionStore) RotateRefresh(ctx context.Context, oldJTI, oldHash, newJTI, newHash string, ttl time.Duration) error {
oldKey := fmt.Sprintf("refresh:%s", oldJTI)
newKey := fmt.Sprintf("refresh:%s", newJTI)
result, err := rotateRefreshScript.Run(
ctx,
s.rdb,
[]string{oldKey, newKey},
oldHash,
newHash,
ttl.Milliseconds(),
).Int()
if err != nil {
return err
}
switch result {
case 1:
return nil
case 0:
return ErrRefreshSessionNotFound
case -1:
return ErrRefreshTokenMismatch
default:
return fmt.Errorf("unexpected rotate refresh result: %d", result)
}
}
func (s *SessionStore) Blacklist(ctx context.Context, accessJTI string, ttl time.Duration) error {
key := fmt.Sprintf("blacklist:%s", accessJTI)
return s.rdb.Set(ctx, key, "1", ttl).Err()
}
func (s *SessionStore) IsBlacklisted(ctx context.Context, accessJTI string) (bool, error) {
key := fmt.Sprintf("blacklist:%s", accessJTI)
n, err := s.rdb.Exists(ctx, key).Result()
if err != nil {
return false, err
}
return n > 0, nil
}
server/internal/shared/auth/session_store_test.go
+52
View File
@@ -0,0 +1,52 @@
package auth
import (
"context"
"testing"
"time"
"github.com/alicebob/miniredis/v2"
goredis "github.com/redis/go-redis/v9"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
func TestSessionStoreRotateRefreshSuccess(t *testing.T) {
mr := miniredis.RunT(t)
client := goredis.NewClient(&goredis.Options{Addr: mr.Addr()})
t.Cleanup(func() { _ = client.Close() })
store := NewSessionStore(client)
ctx := context.Background()
require.NoError(t, store.SaveRefresh(ctx, "old-jti", "old-hash", time.Minute))
require.NoError(t, store.RotateRefresh(ctx, "old-jti", "old-hash", "new-jti", "new-hash", 2*time.Minute))
_, err := store.GetRefresh(ctx, "old-jti")
assert.Error(t, err)
value, err := store.GetRefresh(ctx, "new-jti")
require.NoError(t, err)
assert.Equal(t, "new-hash", value)
}
func TestSessionStoreRotateRefreshRejectsMismatch(t *testing.T) {
mr := miniredis.RunT(t)
client := goredis.NewClient(&goredis.Options{Addr: mr.Addr()})
t.Cleanup(func() { _ = client.Close() })
store := NewSessionStore(client)
ctx := context.Background()
require.NoError(t, store.SaveRefresh(ctx, "old-jti", "old-hash", time.Minute))
err := store.RotateRefresh(ctx, "old-jti", "wrong-hash", "new-jti", "new-hash", 2*time.Minute)
require.ErrorIs(t, err, ErrRefreshTokenMismatch)
value, getErr := store.GetRefresh(ctx, "old-jti")
require.NoError(t, getErr)
assert.Equal(t, "old-hash", value)
_, getErr = store.GetRefresh(ctx, "new-jti")
assert.Error(t, getErr)
}
server/internal/shared/config/config.go
+121
View File
@@ -0,0 +1,121 @@
package config
import (
"fmt"
"os"
"strings"
"time"
"github.com/spf13/viper"
)
type Config struct {
Server ServerConfig `mapstructure:"server"`
Database DatabaseConfig `mapstructure:"database"`
Redis RedisConfig `mapstructure:"redis"`
JWT JWTConfig `mapstructure:"jwt"`
Log LogConfig `mapstructure:"log"`
LLM LLMConfig `mapstructure:"llm"`
Generation GenerationConfig `mapstructure:"generation"`
}
type ServerConfig struct {
Port int `mapstructure:"port"`
Mode string `mapstructure:"mode"`
}
type DatabaseConfig struct {
Host string `mapstructure:"host"`
Port int `mapstructure:"port"`
User string `mapstructure:"user"`
Password string `mapstructure:"password"`
DBName string `mapstructure:"dbname"`
SSLMode string `mapstructure:"sslmode"`
MaxOpenConns int `mapstructure:"max_open_conns"`
MaxIdleConns int `mapstructure:"max_idle_conns"`
}
func (d DatabaseConfig) DSN() string {
return fmt.Sprintf("postgres://%s:%s@%s:%d/%s?sslmode=%s",
d.User, d.Password, d.Host, d.Port, d.DBName, d.SSLMode)
}
type RedisConfig struct {
Addr string `mapstructure:"addr"`
DB int `mapstructure:"db"`
}
type JWTConfig struct {
Secret string `mapstructure:"secret"`
AccessTTL time.Duration `mapstructure:"access_ttl"`
RefreshTTL time.Duration `mapstructure:"refresh_ttl"`
}
type LogConfig struct {
Level string `mapstructure:"level"`
Format string `mapstructure:"format"`
}
type LLMConfig struct {
Provider string `mapstructure:"provider"`
APIKey string `mapstructure:"api_key"`
BaseURL string `mapstructure:"base_url"`
Model string `mapstructure:"model"`
Timeout time.Duration `mapstructure:"timeout"`
MaxOutputTokens int64 `mapstructure:"max_output_tokens"`
Temperature float64 `mapstructure:"temperature"`
}
type GenerationConfig struct {
QueueSize int `mapstructure:"queue_size"`
WorkerConcurrency int `mapstructure:"worker_concurrency"`
StreamEnabled bool `mapstructure:"stream_enabled"`
}
func Load(configPath string) (*Config, error) {
v := viper.New()
v.SetConfigFile(configPath)
v.AutomaticEnv()
v.SetEnvKeyReplacer(strings.NewReplacer(".", "_"))
if err := v.ReadInConfig(); err != nil {
return nil, fmt.Errorf("read config: %w", err)
}
// Allow local overrides
v.SetConfigFile(strings.Replace(configPath, ".yaml", ".local.yaml", 1))
_ = v.MergeInConfig()
var cfg Config
if err := v.Unmarshal(&cfg); err != nil {
return nil, fmt.Errorf("unmarshal config: %w", err)
}
applyEnvOverrides(&cfg)
return &cfg, nil
}
func applyEnvOverrides(cfg *Config) {
if cfg == nil {
return
}
if apiKey, ok := lookupNonEmptyEnv("LLM_API_KEY"); ok {
cfg.LLM.APIKey = apiKey
}
}
func lookupNonEmptyEnv(key string) (string, bool) {
value, ok := os.LookupEnv(key)
if !ok {
return "", false
}
value = strings.TrimSpace(value)
if value == "" {
return "", false
}
return value, true
}
server/internal/shared/config/config_test.go
+57
View File
@@ -0,0 +1,57 @@
package config
import (
"os"
"path/filepath"
"testing"
)
func TestLoadPrefersEnvLLMAPIKey(t *testing.T) {
t.Setenv("LLM_API_KEY", "env-priority-key")
configPath := writeTestConfig(t, `
llm:
provider: ark
api_key: "config-file-key"
`)
cfg, err := Load(configPath)
if err != nil {
t.Fatalf("Load() error = %v", err)
}
if cfg.LLM.APIKey != "env-priority-key" {
t.Fatalf("expected env api key, got %q", cfg.LLM.APIKey)
}
}
func TestLoadFallsBackToConfigLLMAPIKey(t *testing.T) {
t.Setenv("LLM_API_KEY", "")
configPath := writeTestConfig(t, `
llm:
provider: ark
api_key: "config-file-key"
`)
cfg, err := Load(configPath)
if err != nil {
t.Fatalf("Load() error = %v", err)
}
if cfg.LLM.APIKey != "config-file-key" {
t.Fatalf("expected config api key, got %q", cfg.LLM.APIKey)
}
}
func writeTestConfig(t *testing.T, body string) string {
t.Helper()
dir := t.TempDir()
path := filepath.Join(dir, "config.yaml")
if err := os.WriteFile(path, []byte(body), 0o600); err != nil {
t.Fatalf("write config: %v", err)
}
return path
}
server/internal/shared/llm/ark.go
+178
View File
@@ -0,0 +1,178 @@
package llm
import (
"context"
"errors"
"fmt"
"io"
"strings"
"time"
"github.com/volcengine/volcengine-go-sdk/service/arkruntime"
"github.com/volcengine/volcengine-go-sdk/service/arkruntime/model/responses"
"github.com/geo-platform/tenant-api/internal/shared/config"
)
const (
defaultArkBaseURL = "https://ark.cn-beijing.volces.com/api/v3"
defaultArkModel = "doubao-seed-2-0-lite-260215"
defaultTimeout = 2 * time.Minute
)
type arkClient struct {
client *arkruntime.Client
model string
timeout time.Duration
maxOutputTokens int64
temperature float64
}
func NewArkClient(cfg config.LLMConfig) Client {
if strings.TrimSpace(cfg.APIKey) == "" {
return disabledClient{reason: "missing llm.api_key or LLM_API_KEY for Ark provider"}
}
baseURL := strings.TrimSpace(cfg.BaseURL)
if baseURL == "" {
baseURL = defaultArkBaseURL
}
model := strings.TrimSpace(cfg.Model)
if model == "" {
model = defaultArkModel
}
timeout := cfg.Timeout
if timeout <= 0 {
timeout = defaultTimeout
}
maxOutputTokens := cfg.MaxOutputTokens
if maxOutputTokens <= 0 {
maxOutputTokens = 4000
}
temperature := cfg.Temperature
if temperature <= 0 {
temperature = 0.7
}
return &arkClient{
client: arkruntime.NewClientWithApiKey(cfg.APIKey, arkruntime.WithBaseUrl(baseURL)),
model: model,
timeout: timeout,
maxOutputTokens: maxOutputTokens,
temperature: temperature,
}
}
func (c *arkClient) Validate() error {
if c == nil || c.client == nil {
return fmt.Errorf("%w: ark client is not initialized", ErrNotConfigured)
}
if strings.TrimSpace(c.model) == "" {
return fmt.Errorf("%w: ark model is empty", ErrNotConfigured)
}
return nil
}
func (c *arkClient) Generate(ctx context.Context, req GenerateRequest, onDelta func(string)) (*GenerateResult, error) {
if err := c.Validate(); err != nil {
return nil, err
}
prompt := strings.TrimSpace(req.Prompt)
if prompt == "" {
return nil, errors.New("prompt is required")
}
callCtx, cancel := context.WithTimeout(ctx, c.timeout)
defer cancel()
store := false
inputMessage := &responses.ItemInputMessage{
Role: responses.MessageRole_user,
Content: []*responses.ContentItem{
{
Union: &responses.ContentItem_Text{
Text: &responses.ContentItemText{
Type: responses.ContentItemType_input_text,
Text: prompt,
},
},
},
},
}
stream, err := c.client.CreateResponsesStream(callCtx, &responses.ResponsesRequest{
Model: c.model,
MaxOutputTokens: &c.maxOutputTokens,
Store: &store,
Temperature: &c.temperature,
Input: &responses.ResponsesInput{
Union: &responses.ResponsesInput_ListValue{
ListValue: &responses.InputItemList{
ListValue: []*responses.InputItem{
{
Union: &responses.InputItem_InputMessage{
InputMessage: inputMessage,
},
},
},
},
},
},
})
if err != nil {
return nil, fmt.Errorf("create ark response stream: %w", err)
}
defer stream.Close()
var content strings.Builder
for {
event, err := stream.Recv()
if err != nil {
if errors.Is(err, io.EOF) {
break
}
return nil, fmt.Errorf("receive ark stream event: %w", err)
}
if delta := event.GetText(); delta != nil {
chunk := delta.GetDelta()
if chunk != "" {
content.WriteString(chunk)
if onDelta != nil {
onDelta(chunk)
}
}
continue
}
if failed := event.GetResponseFailed(); failed != nil {
if respErr := failed.GetResponse().GetError(); respErr != nil {
return nil, fmt.Errorf("ark response failed: %s", respErr.GetMessage())
}
return nil, errors.New("ark response failed")
}
if incomplete := event.GetResponseIncomplete(); incomplete != nil {
if respErr := incomplete.GetResponse().GetError(); respErr != nil {
return nil, fmt.Errorf("ark response incomplete: %s", respErr.GetMessage())
}
return nil, errors.New("ark response incomplete")
}
}
result := strings.TrimSpace(content.String())
if result == "" {
return nil, errors.New("ark returned empty content")
}
return &GenerateResult{
Content: result,
Model: c.model,
}, nil
}
server/internal/shared/llm/client.go
+53
View File
@@ -0,0 +1,53 @@
package llm
import (
"context"
"errors"
"fmt"
"strings"
"github.com/geo-platform/tenant-api/internal/shared/config"
)
var ErrNotConfigured = errors.New("llm provider is not configured")
type GenerateRequest struct {
Prompt string
}
type GenerateResult struct {
Content string
Model string
}
type Client interface {
Validate() error
Generate(ctx context.Context, req GenerateRequest, onDelta func(string)) (*GenerateResult, error)
}
func New(cfg config.LLMConfig) Client {
provider := strings.ToLower(strings.TrimSpace(cfg.Provider))
switch provider {
case "", "disabled":
return disabledClient{reason: "llm provider is disabled"}
case "ark":
return NewArkClient(cfg)
default:
return disabledClient{reason: fmt.Sprintf("unsupported llm provider %q", cfg.Provider)}
}
}
type disabledClient struct {
reason string
}
func (c disabledClient) Validate() error {
if c.reason == "" {
c.reason = "missing LLM configuration"
}
return fmt.Errorf("%w: %s", ErrNotConfigured, c.reason)
}
func (c disabledClient) Generate(context.Context, GenerateRequest, func(string)) (*GenerateResult, error) {
return nil, c.Validate()
}
server/internal/shared/middleware/cors.go
+21
View File
@@ -0,0 +1,21 @@
package middleware
import (
"github.com/gin-gonic/gin"
)
func CORS() gin.HandlerFunc {
return func(c *gin.Context) {
c.Header("Access-Control-Allow-Origin", "*")
c.Header("Access-Control-Allow-Methods", "GET, POST, PUT, DELETE, OPTIONS")
c.Header("Access-Control-Allow-Headers", "Origin, Content-Type, Authorization, X-Request-ID")
c.Header("Access-Control-Expose-Headers", "X-Request-ID")
c.Header("Access-Control-Max-Age", "86400")
if c.Request.Method == "OPTIONS" {
c.AbortWithStatus(204)
return
}
c.Next()
}
}
server/internal/shared/middleware/logger.go
+23
View File
@@ -0,0 +1,23 @@
package middleware
import (
"time"
"github.com/gin-gonic/gin"
"go.uber.org/zap"
)
func Logger(logger *zap.Logger) gin.HandlerFunc {
return func(c *gin.Context) {
start := time.Now()
c.Next()
logger.Info("request",
zap.String("method", c.Request.Method),
zap.String("path", c.Request.URL.Path),
zap.Int("status", c.Writer.Status()),
zap.Duration("latency", time.Since(start)),
zap.String("request_id", RequestIDFromGin(c)),
zap.String("client_ip", c.ClientIP()),
)
}
}
server/internal/shared/middleware/recovery.go
+30
View File
@@ -0,0 +1,30 @@
package middleware
import (
"net/http"
"runtime/debug"
"github.com/gin-gonic/gin"
"go.uber.org/zap"
)
func Recovery(logger *zap.Logger) gin.HandlerFunc {
return func(c *gin.Context) {
defer func() {
if r := recover(); r != nil {
logger.Error("panic recovered",
zap.Any("error", r),
zap.String("stack", string(debug.Stack())),
zap.String("request_id", RequestIDFromGin(c)),
)
c.AbortWithStatusJSON(http.StatusInternalServerError, gin.H{
"code": 50000,
"message": "internal_error",
"detail": "An unexpected error occurred",
"request_id": RequestIDFromGin(c),
})
}
}()
c.Next()
}
}
server/internal/shared/middleware/request_id.go
+28
View File
@@ -0,0 +1,28 @@
package middleware
import (
"github.com/gin-gonic/gin"
"github.com/google/uuid"
)
const RequestIDHeader = "X-Request-ID"
const requestIDKey = "request_id"
func RequestID() gin.HandlerFunc {
return func(c *gin.Context) {
rid := c.GetHeader(RequestIDHeader)
if rid == "" {
rid = uuid.New().String()
}
c.Set(requestIDKey, rid)
c.Header(RequestIDHeader, rid)
c.Next()
}
}
func RequestIDFromGin(c *gin.Context) string {
if rid, ok := c.Get(requestIDKey); ok {
return rid.(string)
}
return ""
}
server/internal/shared/middleware/request_id_test.go
+43
View File
@@ -0,0 +1,43 @@
package middleware
import (
"net/http"
"net/http/httptest"
"testing"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/assert"
)
func TestRequestID_GeneratesNew(t *testing.T) {
r := gin.New()
r.Use(RequestID())
r.GET("/test", func(c *gin.Context) {
rid := RequestIDFromGin(c)
assert.NotEmpty(t, rid)
c.Status(http.StatusOK)
})
w := httptest.NewRecorder()
req := httptest.NewRequest(http.MethodGet, "/test", nil)
r.ServeHTTP(w, req)
assert.Equal(t, http.StatusOK, w.Code)
assert.NotEmpty(t, w.Header().Get(RequestIDHeader))
}
func TestRequestID_UsesExisting(t *testing.T) {
r := gin.New()
r.Use(RequestID())
r.GET("/test", func(c *gin.Context) {
c.Status(http.StatusOK)
})
w := httptest.NewRecorder()
req := httptest.NewRequest(http.MethodGet, "/test", nil)
req.Header.Set(RequestIDHeader, "my-custom-rid")
r.ServeHTTP(w, req)
assert.Equal(t, http.StatusOK, w.Code)
assert.Equal(t, "my-custom-rid", w.Header().Get(RequestIDHeader))
}
server/internal/shared/middleware/tenant_scope.go
+37
View File
@@ -0,0 +1,37 @@
package middleware
import (
"context"
"github.com/gin-gonic/gin"
"github.com/geo-platform/tenant-api/internal/shared/auth"
"github.com/geo-platform/tenant-api/internal/shared/response"
)
type tenantIDKey struct{}
func TenantScope() gin.HandlerFunc {
return func(c *gin.Context) {
actor, ok := auth.ActorFromCtx(c.Request.Context())
if !ok || actor.TenantID == 0 {
response.Error(c, response.ErrUnauthorized(40104, "tenant_scope_missing", "tenant context is required"))
c.Abort()
return
}
ctx := WithTenantID(c.Request.Context(), actor.TenantID)
c.Request = c.Request.WithContext(ctx)
c.Next()
}
}
func WithTenantID(ctx context.Context, tenantID int64) context.Context {
return context.WithValue(ctx, tenantIDKey{}, tenantID)
}
func TenantIDFromCtx(ctx context.Context) int64 {
if v, ok := ctx.Value(tenantIDKey{}).(int64); ok {
return v
}
return 0
}
server/internal/shared/middleware/tenant_scope_test.go
+78
View File
@@ -0,0 +1,78 @@
package middleware
import (
"context"
"net/http"
"net/http/httptest"
"testing"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/assert"
"github.com/geo-platform/tenant-api/internal/shared/auth"
)
func init() {
gin.SetMode(gin.TestMode)
}
func TestTenantScope_ValidActor(t *testing.T) {
r := gin.New()
r.Use(func(c *gin.Context) {
actor := auth.Actor{UserID: 1, TenantID: 10, Role: "editor"}
c.Request = c.Request.WithContext(auth.WithActor(c.Request.Context(), actor))
c.Next()
})
r.Use(TenantScope())
var gotTenantID int64
r.GET("/test", func(c *gin.Context) {
gotTenantID = TenantIDFromCtx(c.Request.Context())
c.Status(http.StatusOK)
})
w := httptest.NewRecorder()
req := httptest.NewRequest(http.MethodGet, "/test", nil)
r.ServeHTTP(w, req)
assert.Equal(t, http.StatusOK, w.Code)
assert.Equal(t, int64(10), gotTenantID)
}
func TestTenantScope_MissingActor(t *testing.T) {
r := gin.New()
r.Use(TenantScope())
r.GET("/test", func(c *gin.Context) { c.Status(http.StatusOK) })
w := httptest.NewRecorder()
req := httptest.NewRequest(http.MethodGet, "/test", nil)
r.ServeHTTP(w, req)
assert.Equal(t, http.StatusUnauthorized, w.Code)
}
func TestTenantScope_ZeroTenantID(t *testing.T) {
r := gin.New()
r.Use(func(c *gin.Context) {
actor := auth.Actor{UserID: 1, TenantID: 0, Role: "editor"}
c.Request = c.Request.WithContext(auth.WithActor(c.Request.Context(), actor))
c.Next()
})
r.Use(TenantScope())
r.GET("/test", func(c *gin.Context) { c.Status(http.StatusOK) })
w := httptest.NewRecorder()
req := httptest.NewRequest(http.MethodGet, "/test", nil)
r.ServeHTTP(w, req)
assert.Equal(t, http.StatusUnauthorized, w.Code)
}
func TestTenantIDFromCtx_RoundTrip(t *testing.T) {
ctx := WithTenantID(context.Background(), 42)
assert.Equal(t, int64(42), TenantIDFromCtx(ctx))
}
func TestTenantIDFromCtx_Missing(t *testing.T) {
assert.Equal(t, int64(0), TenantIDFromCtx(context.Background()))
}
server/internal/shared/observability/logger.go
+23
View File
@@ -0,0 +1,23 @@
package observability
import (
"go.uber.org/zap"
"go.uber.org/zap/zapcore"
)
func NewLogger(level, format string) (*zap.Logger, error) {
var cfg zap.Config
if format == "json" {
cfg = zap.NewProductionConfig()
} else {
cfg = zap.NewDevelopmentConfig()
}
lvl, err := zapcore.ParseLevel(level)
if err != nil {
lvl = zapcore.InfoLevel
}
cfg.Level = zap.NewAtomicLevelAt(lvl)
return cfg.Build()
}
server/internal/shared/repository/postgres/postgres.go
+32
View File
@@ -0,0 +1,32 @@
package postgres
import (
"context"
"fmt"
"github.com/jackc/pgx/v5/pgxpool"
"github.com/geo-platform/tenant-api/internal/shared/config"
)
func NewPool(ctx context.Context, cfg config.DatabaseConfig) (*pgxpool.Pool, error) {
poolCfg, err := pgxpool.ParseConfig(cfg.DSN())
if err != nil {
return nil, fmt.Errorf("parse pool config: %w", err)
}
poolCfg.MaxConns = int32(cfg.MaxOpenConns)
poolCfg.MinConns = int32(cfg.MaxIdleConns)
pool, err := pgxpool.NewWithConfig(ctx, poolCfg)
if err != nil {
return nil, fmt.Errorf("create pool: %w", err)
}
if err := pool.Ping(ctx); err != nil {
pool.Close()
return nil, fmt.Errorf("ping postgres: %w", err)
}
return pool, nil
}
server/internal/shared/repository/redis/redis.go
+23
View File
@@ -0,0 +1,23 @@
package redis
import (
"context"
"fmt"
goredis "github.com/redis/go-redis/v9"
"github.com/geo-platform/tenant-api/internal/shared/config"
)
func NewClient(ctx context.Context, cfg config.RedisConfig) (*goredis.Client, error) {
client := goredis.NewClient(&goredis.Options{
Addr: cfg.Addr,
DB: cfg.DB,
})
if err := client.Ping(ctx).Err(); err != nil {
return nil, fmt.Errorf("ping redis: %w", err)
}
return client, nil
}
server/internal/shared/response/errors.go
+35
View File
@@ -0,0 +1,35 @@
package response
import "net/http"
func ErrBadRequest(code int, message, detail string) *AppError {
return &AppError{HTTPStatus: http.StatusBadRequest, Code: code, Message: message, Detail: detail}
}
func ErrUnauthorized(code int, message, detail string) *AppError {
return &AppError{HTTPStatus: http.StatusUnauthorized, Code: code, Message: message, Detail: detail}
}
func ErrForbidden(code int, message, detail string) *AppError {
return &AppError{HTTPStatus: http.StatusForbidden, Code: code, Message: message, Detail: detail}
}
func ErrNotFound(code int, message, detail string) *AppError {
return &AppError{HTTPStatus: http.StatusNotFound, Code: code, Message: message, Detail: detail}
}
func ErrConflict(code int, message, detail string) *AppError {
return &AppError{HTTPStatus: http.StatusConflict, Code: code, Message: message, Detail: detail}
}
func ErrTooManyRequests(code int, message, detail string) *AppError {
return &AppError{HTTPStatus: http.StatusTooManyRequests, Code: code, Message: message, Detail: detail}
}
func ErrInternal(code int, message, detail string) *AppError {
return &AppError{HTTPStatus: http.StatusInternalServerError, Code: code, Message: message, Detail: detail}
}
func ErrServiceUnavailable(code int, message, detail string) *AppError {
return &AppError{HTTPStatus: http.StatusServiceUnavailable, Code: code, Message: message, Detail: detail}
}
server/internal/shared/response/response.go
+70
View File
@@ -0,0 +1,70 @@
package response
import (
"net/http"
"github.com/gin-gonic/gin"
)
type AppError struct {
HTTPStatus int `json:"-"`
Code int `json:"code"`
Message string `json:"message"`
Detail string `json:"detail,omitempty"`
Cause error `json:"-"`
}
func (e *AppError) Error() string {
if e.Cause != nil {
return e.Cause.Error()
}
return e.Message
}
func Success(c *gin.Context, data interface{}) {
c.JSON(http.StatusOK, gin.H{
"code": 0,
"message": "ok",
"data": data,
"request_id": getRequestID(c),
})
}
func SuccessWithStatus(c *gin.Context, status int, data interface{}) {
c.JSON(status, gin.H{
"code": 0,
"message": "ok",
"data": data,
"request_id": getRequestID(c),
})
}
func Error(c *gin.Context, err error) {
appErr := Normalize(err)
c.JSON(appErr.HTTPStatus, gin.H{
"code": appErr.Code,
"message": appErr.Message,
"detail": appErr.Detail,
"request_id": getRequestID(c),
})
}
func Normalize(err error) *AppError {
if appErr, ok := err.(*AppError); ok {
return appErr
}
return &AppError{
HTTPStatus: http.StatusInternalServerError,
Code: 50000,
Message: "internal_error",
Detail: err.Error(),
Cause: err,
}
}
func getRequestID(c *gin.Context) string {
if rid, ok := c.Get("request_id"); ok {
return rid.(string)
}
return ""
}
server/internal/shared/stream/generation.go
+221
View File
@@ -0,0 +1,221 @@
package stream
import (
"sync"
"time"
)
type GenerationEvent struct {
Type string `json:"type"`
ArticleID int64 `json:"article_id"`
TaskID int64 `json:"task_id"`
Title string `json:"title,omitempty"`
Status string `json:"status,omitempty"`
Delta string `json:"delta,omitempty"`
Content string `json:"content,omitempty"`
Error string `json:"error,omitempty"`
Done bool `json:"done,omitempty"`
UpdatedAt time.Time `json:"updated_at"`
}
type generationStream struct {
snapshot GenerationEvent
subscribers map[int]chan GenerationEvent
nextSubscriber int
completed bool
}
type GenerationHub struct {
mu sync.Mutex
streams map[int64]*generationStream
}
func NewGenerationHub() *GenerationHub {
return &GenerationHub{
streams: make(map[int64]*generationStream),
}
}
func (h *GenerationHub) Start(articleID, taskID int64, title string) {
h.publish(articleID, func(state *generationStream) GenerationEvent {
now := time.Now()
state.completed = false
state.snapshot = GenerationEvent{
Type: "snapshot",
ArticleID: articleID,
TaskID: taskID,
Title: title,
Status: "generating",
UpdatedAt: now,
}
return GenerationEvent{
Type: "status",
ArticleID: articleID,
TaskID: taskID,
Title: title,
Status: "generating",
UpdatedAt: now,
}
})
}
func (h *GenerationHub) AppendDelta(articleID, taskID int64, title, delta string) {
h.publish(articleID, func(state *generationStream) GenerationEvent {
now := time.Now()
state.snapshot.Type = "snapshot"
state.snapshot.ArticleID = articleID
state.snapshot.TaskID = taskID
state.snapshot.Title = title
state.snapshot.Status = "generating"
state.snapshot.Content += delta
state.snapshot.UpdatedAt = now
return GenerationEvent{
Type: "delta",
ArticleID: articleID,
TaskID: taskID,
Title: title,
Status: "generating",
Delta: delta,
Content: state.snapshot.Content,
UpdatedAt: now,
}
})
}
func (h *GenerationHub) Complete(articleID, taskID int64, title, content string) {
h.publish(articleID, func(state *generationStream) GenerationEvent {
now := time.Now()
state.completed = true
state.snapshot = GenerationEvent{
Type: "snapshot",
ArticleID: articleID,
TaskID: taskID,
Title: title,
Status: "completed",
Content: content,
Done: true,
UpdatedAt: now,
}
return GenerationEvent{
Type: "completed",
ArticleID: articleID,
TaskID: taskID,
Title: title,
Status: "completed",
Content: content,
Done: true,
UpdatedAt: now,
}
})
}
func (h *GenerationHub) Fail(articleID, taskID int64, title, errMsg string) {
h.publish(articleID, func(state *generationStream) GenerationEvent {
now := time.Now()
state.completed = true
state.snapshot.Type = "snapshot"
state.snapshot.ArticleID = articleID
state.snapshot.TaskID = taskID
state.snapshot.Title = title
state.snapshot.Status = "failed"
state.snapshot.Error = errMsg
state.snapshot.Done = true
state.snapshot.UpdatedAt = now
return GenerationEvent{
Type: "error",
ArticleID: articleID,
TaskID: taskID,
Title: title,
Status: "failed",
Error: errMsg,
Done: true,
UpdatedAt: now,
}
})
}
func (h *GenerationHub) Subscribe(articleID int64) (<-chan GenerationEvent, *GenerationEvent, func()) {
h.mu.Lock()
state := h.streams[articleID]
if state == nil {
h.mu.Unlock()
return nil, nil, func() {}
}
id := state.nextSubscriber
state.nextSubscriber++
ch := make(chan GenerationEvent, 64)
state.subscribers[id] = ch
snapshot := state.snapshot
h.mu.Unlock()
cancel := func() {
h.mu.Lock()
defer h.mu.Unlock()
current := h.streams[articleID]
if current == nil {
close(ch)
return
}
delete(current.subscribers, id)
shouldDelete := current.completed && len(current.subscribers) == 0
if shouldDelete {
delete(h.streams, articleID)
}
close(ch)
}
return ch, &snapshot, cancel
}
func (h *GenerationHub) Snapshot(articleID int64) (*GenerationEvent, bool) {
h.mu.Lock()
defer h.mu.Unlock()
state := h.streams[articleID]
if state == nil {
return nil, false
}
snapshot := state.snapshot
return &snapshot, true
}
func (h *GenerationHub) publish(articleID int64, mutate func(*generationStream) GenerationEvent) {
h.mu.Lock()
state := h.ensureLocked(articleID)
event := mutate(state)
subscribers := make([]chan GenerationEvent, 0, len(state.subscribers))
for _, ch := range state.subscribers {
subscribers = append(subscribers, ch)
}
h.mu.Unlock()
for _, ch := range subscribers {
select {
case ch <- event:
default:
}
}
}
func (h *GenerationHub) ensureLocked(articleID int64) *generationStream {
state := h.streams[articleID]
if state != nil {
return state
}
state = &generationStream{
subscribers: make(map[int]chan GenerationEvent),
}
h.streams[articleID] = state
return state
}