chore: import upstream snapshot with attribution

This commit is contained in:
wehub-resource-sync
2026-07-13 13:28:13 +08:00
commit 0878425be3
1160 changed files with 491311 additions and 0 deletions
+121
View File
@@ -0,0 +1,121 @@
package auth
import (
"sync"
"time"
"pentagi/pkg/server/models"
"github.com/jinzhu/gorm"
)
// tokenCacheEntry represents a cached token status entry
type tokenCacheEntry struct {
status models.TokenStatus
privileges []string
notFound bool // negative caching
expiresAt time.Time
}
// TokenCache provides caching for token status lookups
type TokenCache struct {
cache sync.Map
ttl time.Duration
db *gorm.DB
}
// NewTokenCache creates a new token cache instance
func NewTokenCache(db *gorm.DB) *TokenCache {
return &TokenCache{
ttl: 5 * time.Minute,
db: db,
}
}
// SetTTL sets the TTL for the token cache
func (tc *TokenCache) SetTTL(ttl time.Duration) {
tc.ttl = ttl
}
// GetStatus retrieves token status and privileges from cache or database
func (tc *TokenCache) GetStatus(tokenID string) (models.TokenStatus, []string, error) {
// check cache first
if entry, ok := tc.cache.Load(tokenID); ok {
cached := entry.(tokenCacheEntry)
if time.Now().Before(cached.expiresAt) {
// return cached "not found" error
if cached.notFound {
return "", nil, gorm.ErrRecordNotFound
}
return cached.status, cached.privileges, nil
}
// cache entry expired, remove it
tc.cache.Delete(tokenID)
}
// load from database with role privileges
var token models.APIToken
if err := tc.db.Where("token_id = ? AND deleted_at IS NULL", tokenID).First(&token).Error; err != nil {
if gorm.IsRecordNotFoundError(err) {
// cache negative result (token not found)
tc.cache.Store(tokenID, tokenCacheEntry{
notFound: true,
expiresAt: time.Now().Add(tc.ttl),
})
return "", nil, gorm.ErrRecordNotFound
}
return "", nil, err
}
// load privileges for the token's role
var privileges []models.Privilege
if err := tc.db.Where("role_id = ?", token.RoleID).Find(&privileges).Error; err != nil {
return "", nil, err
}
// extract privilege names
privNames := make([]string, len(privileges))
for i, priv := range privileges {
privNames[i] = priv.Name
}
// always add automation privilege for API tokens
privNames = append(privNames, PrivilegeAutomation)
// update cache with positive result
tc.cache.Store(tokenID, tokenCacheEntry{
status: token.Status,
privileges: privNames,
notFound: false,
expiresAt: time.Now().Add(tc.ttl),
})
return token.Status, privNames, nil
}
// Invalidate removes a specific token from cache
func (tc *TokenCache) Invalidate(tokenID string) {
tc.cache.Delete(tokenID)
}
// InvalidateUser removes all tokens for a specific user from cache
func (tc *TokenCache) InvalidateUser(userID uint64) {
// load all tokens for this user
var tokens []models.APIToken
if err := tc.db.Where("user_id = ? AND deleted_at IS NULL", userID).Find(&tokens).Error; err != nil {
return
}
// invalidate each token in cache
for _, token := range tokens {
tc.cache.Delete(token.TokenID)
}
}
// InvalidateAll clears the entire cache
func (tc *TokenCache) InvalidateAll() {
tc.cache.Range(func(key, value any) bool {
tc.cache.Delete(key)
return true
})
}
@@ -0,0 +1,322 @@
package auth_test
import (
"testing"
"time"
"pentagi/pkg/server/auth"
"pentagi/pkg/server/models"
"github.com/jinzhu/gorm"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
func TestTokenCache_GetStatus(t *testing.T) {
db := setupTestDB(t)
defer db.Close()
cache := auth.NewTokenCache(db)
tokenID := "testtoken1"
// Insert test token
token := models.APIToken{
TokenID: tokenID,
UserID: 1,
RoleID: 2,
TTL: 3600,
Status: models.TokenStatusActive,
}
err := db.Create(&token).Error
require.NoError(t, err)
// Test: Get status (should hit database)
status, privileges, err := cache.GetStatus(tokenID)
require.NoError(t, err)
assert.Equal(t, models.TokenStatusActive, status)
assert.NotEmpty(t, privileges)
assert.Contains(t, privileges, auth.PrivilegeAutomation)
assert.Contains(t, privileges, "flows.create")
assert.Contains(t, privileges, "settings.tokens.view")
// Test: Get status again (should hit cache)
status, privileges, err = cache.GetStatus(tokenID)
require.NoError(t, err)
assert.Equal(t, models.TokenStatusActive, status)
assert.NotEmpty(t, privileges)
assert.Contains(t, privileges, auth.PrivilegeAutomation)
// Test: Non-existent token
_, _, err = cache.GetStatus("nonexistent")
assert.Error(t, err)
assert.Equal(t, gorm.ErrRecordNotFound, err)
}
func TestTokenCache_Invalidate(t *testing.T) {
db := setupTestDB(t)
defer db.Close()
cache := auth.NewTokenCache(db)
tokenID := "testtoken2"
// Insert test token
token := models.APIToken{
TokenID: tokenID,
UserID: 1,
RoleID: 2,
TTL: 3600,
Status: models.TokenStatusActive,
}
err := db.Create(&token).Error
require.NoError(t, err)
// Get status to populate cache
status, privileges, err := cache.GetStatus(tokenID)
require.NoError(t, err)
assert.Equal(t, models.TokenStatusActive, status)
assert.NotEmpty(t, privileges)
// Update token in database
db.Model(&token).Update("status", models.TokenStatusRevoked)
// Status should still be active (from cache)
status, privileges, err = cache.GetStatus(tokenID)
require.NoError(t, err)
assert.Equal(t, models.TokenStatusActive, status)
assert.NotEmpty(t, privileges)
// Invalidate cache
cache.Invalidate(tokenID)
// Status should now be revoked (from database)
status, privileges, err = cache.GetStatus(tokenID)
require.NoError(t, err)
assert.Equal(t, models.TokenStatusRevoked, status)
assert.NotEmpty(t, privileges)
}
func TestTokenCache_InvalidateUser(t *testing.T) {
db := setupTestDB(t)
defer db.Close()
cache := auth.NewTokenCache(db)
userID := uint64(1)
// Insert multiple tokens for user
tokens := []models.APIToken{
{
TokenID: "token1",
UserID: userID,
RoleID: 2,
TTL: 3600,
Status: models.TokenStatusActive,
},
{
TokenID: "token2",
UserID: userID,
RoleID: 2,
TTL: 3600,
Status: models.TokenStatusActive,
},
}
for _, token := range tokens {
err := db.Create(&token).Error
require.NoError(t, err)
}
// Populate cache
for _, token := range tokens {
_, _, err := cache.GetStatus(token.TokenID)
require.NoError(t, err)
}
// Update tokens in database
db.Model(&models.APIToken{}).Where("user_id = ?", userID).Update("status", models.TokenStatusRevoked)
// Invalidate all user tokens
cache.InvalidateUser(userID)
// All tokens should now show revoked status
for _, token := range tokens {
status, privileges, err := cache.GetStatus(token.TokenID)
require.NoError(t, err)
assert.Equal(t, models.TokenStatusRevoked, status)
assert.NotEmpty(t, privileges)
}
}
func TestTokenCache_Expiration(t *testing.T) {
db := setupTestDB(t)
defer db.Close()
// Create cache with very short TTL for testing
cache := auth.NewTokenCache(db)
cache.SetTTL(300 * time.Millisecond)
tokenID := "testtoken3"
// Insert test token
token := models.APIToken{
TokenID: tokenID,
UserID: 1,
RoleID: 2,
TTL: 3600,
Status: models.TokenStatusActive,
}
err := db.Create(&token).Error
require.NoError(t, err)
// Get status to populate cache
status, privileges, err := cache.GetStatus(tokenID)
require.NoError(t, err)
assert.Equal(t, models.TokenStatusActive, status)
assert.NotEmpty(t, privileges)
// Update token in database
db.Model(&token).Update("status", models.TokenStatusRevoked)
// Wait for cache to expire
time.Sleep(500 * time.Millisecond)
// Status should now be revoked (cache expired, reading from DB)
status, privileges, err = cache.GetStatus(tokenID)
require.NoError(t, err)
assert.Equal(t, models.TokenStatusRevoked, status)
assert.NotEmpty(t, privileges)
}
func TestTokenCache_PrivilegesByRole(t *testing.T) {
db := setupTestDB(t)
defer db.Close()
cache := auth.NewTokenCache(db)
// Test Admin token (role_id = 1)
adminTokenID := "admin_token"
adminToken := models.APIToken{
TokenID: adminTokenID,
UserID: 1,
RoleID: 1,
TTL: 3600,
Status: models.TokenStatusActive,
}
err := db.Create(&adminToken).Error
require.NoError(t, err)
status, adminPrivs, err := cache.GetStatus(adminTokenID)
require.NoError(t, err)
assert.Equal(t, models.TokenStatusActive, status)
assert.NotEmpty(t, adminPrivs)
assert.Contains(t, adminPrivs, auth.PrivilegeAutomation)
assert.Contains(t, adminPrivs, "users.create")
assert.Contains(t, adminPrivs, "users.delete")
assert.Contains(t, adminPrivs, "settings.tokens.admin")
// Test User token (role_id = 2)
userTokenID := "user_token"
userToken := models.APIToken{
TokenID: userTokenID,
UserID: 2,
RoleID: 2,
TTL: 3600,
Status: models.TokenStatusActive,
}
err = db.Create(&userToken).Error
require.NoError(t, err)
status, userPrivs, err := cache.GetStatus(userTokenID)
require.NoError(t, err)
assert.Equal(t, models.TokenStatusActive, status)
assert.NotEmpty(t, userPrivs)
assert.Contains(t, userPrivs, auth.PrivilegeAutomation)
assert.Contains(t, userPrivs, "flows.create")
assert.Contains(t, userPrivs, "settings.tokens.view")
// User should NOT have admin privileges
assert.NotContains(t, userPrivs, "users.create")
assert.NotContains(t, userPrivs, "users.delete")
assert.NotContains(t, userPrivs, "settings.tokens.admin")
// Admin should have more privileges than User
assert.Greater(t, len(adminPrivs), len(userPrivs))
}
func TestTokenCache_NegativeCaching(t *testing.T) {
db := setupTestDB(t)
defer db.Close()
cache := auth.NewTokenCache(db)
nonExistentTokenID := "nonexistent"
// First call - should hit database and cache the "not found"
_, _, err := cache.GetStatus(nonExistentTokenID)
require.Error(t, err)
assert.Equal(t, gorm.ErrRecordNotFound, err)
// Second call - should return from cache without hitting DB
// We can verify this by checking error is still the same
_, _, err = cache.GetStatus(nonExistentTokenID)
require.Error(t, err)
assert.Equal(t, gorm.ErrRecordNotFound, err, "Should return cached not found error")
// Now create the token in DB
token := models.APIToken{
TokenID: nonExistentTokenID,
UserID: 1,
RoleID: 2,
TTL: 3600,
Status: models.TokenStatusActive,
}
err = db.Create(&token).Error
require.NoError(t, err)
// Should still return cached "not found" until invalidated
_, _, err = cache.GetStatus(nonExistentTokenID)
require.Error(t, err)
assert.Equal(t, gorm.ErrRecordNotFound, err, "Should still return cached not found")
// Invalidate cache
cache.Invalidate(nonExistentTokenID)
// Now should find the token
status, privileges, err := cache.GetStatus(nonExistentTokenID)
require.NoError(t, err)
assert.Equal(t, models.TokenStatusActive, status)
assert.NotEmpty(t, privileges)
}
func TestTokenCache_NegativeCachingExpiration(t *testing.T) {
db := setupTestDB(t)
defer db.Close()
cache := auth.NewTokenCache(db)
cache.SetTTL(300 * time.Millisecond)
nonExistentTokenID := "temp_nonexistent"
// First call - cache the "not found"
_, _, err := cache.GetStatus(nonExistentTokenID)
require.Error(t, err)
assert.Equal(t, gorm.ErrRecordNotFound, err)
// Create token in DB
token := models.APIToken{
TokenID: nonExistentTokenID,
UserID: 1,
RoleID: 2,
TTL: 3600,
Status: models.TokenStatusActive,
}
err = db.Create(&token).Error
require.NoError(t, err)
// Wait for cache to expire
time.Sleep(500 * time.Millisecond)
// Now should find the token (cache expired)
status, privileges, err := cache.GetStatus(nonExistentTokenID)
require.NoError(t, err)
assert.Equal(t, models.TokenStatusActive, status)
assert.NotEmpty(t, privileges)
}
+28
View File
@@ -0,0 +1,28 @@
package auth
import (
"crypto/rand"
"fmt"
"math/big"
)
const (
Base62Chars = "0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz"
TokenIDLength = 10
)
// GenerateTokenID generates a random base62 string of specified length
func GenerateTokenID() (string, error) {
b := make([]byte, TokenIDLength)
maxIdx := big.NewInt(int64(len(Base62Chars)))
for i := range b {
idx, err := rand.Int(rand.Reader, maxIdx)
if err != nil {
return "", fmt.Errorf("error generating token ID: %w", err)
}
b[i] = Base62Chars[idx.Int64()]
}
return string(b), nil
}
@@ -0,0 +1,49 @@
package auth_test
import (
"testing"
"pentagi/pkg/server/auth"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
func TestGenerateTokenID(t *testing.T) {
// Test basic generation
tokenID, err := auth.GenerateTokenID()
require.NoError(t, err)
assert.Len(t, tokenID, auth.TokenIDLength, "Token ID should have correct length")
// Test that all characters are from base62 charset
for _, char := range tokenID {
assert.Contains(t, auth.Base62Chars, string(char), "Token ID should only contain base62 characters")
}
// Test uniqueness (generate multiple tokens and check they're different)
tokens := make(map[string]bool)
for i := 0; i < 100; i++ {
token, err := auth.GenerateTokenID()
require.NoError(t, err)
assert.Len(t, token, auth.TokenIDLength)
assert.False(t, tokens[token], "Generated tokens should be unique")
tokens[token] = true
}
}
func TestGenerateTokenIDFormat(t *testing.T) {
// Test that token IDs match expected format
tokenID, err := auth.GenerateTokenID()
require.NoError(t, err)
// Should be exactly 10 characters
assert.Equal(t, 10, len(tokenID))
// Should only contain alphanumeric characters
for _, char := range tokenID {
isValid := (char >= '0' && char <= '9') ||
(char >= 'A' && char <= 'Z') ||
(char >= 'a' && char <= 'z')
assert.True(t, isValid, "Character %c should be alphanumeric", char)
}
}
+62
View File
@@ -0,0 +1,62 @@
package auth
import (
"errors"
"fmt"
"time"
"pentagi/pkg/server/models"
"github.com/golang-jwt/jwt/v5"
)
func MakeAPIToken(globalSalt string, claims jwt.Claims) (string, error) {
token := jwt.NewWithClaims(jwt.SigningMethodHS256, claims)
tokenString, err := token.SignedString(MakeJWTSigningKey(globalSalt))
if err != nil {
return "", fmt.Errorf("failed to sign token: %w", err)
}
return tokenString, nil
}
func MakeAPITokenClaims(tokenID, uhash string, uid, rid, ttl uint64) jwt.Claims {
now := time.Now()
return models.APITokenClaims{
TokenID: tokenID,
RID: rid,
UID: uid,
UHASH: uhash,
RegisteredClaims: jwt.RegisteredClaims{
ExpiresAt: jwt.NewNumericDate(now.Add(time.Duration(ttl) * time.Second)),
IssuedAt: jwt.NewNumericDate(now),
Subject: "api_token",
},
}
}
func ValidateAPIToken(tokenString, globalSalt string) (*models.APITokenClaims, error) {
var claims models.APITokenClaims
token, err := jwt.ParseWithClaims(tokenString, &claims, func(token *jwt.Token) (any, error) {
// verify signing algorithm to prevent "alg: none"
if _, ok := token.Method.(*jwt.SigningMethodHMAC); !ok {
return nil, fmt.Errorf("unexpected signing method: %v", token.Header["alg"])
}
return MakeJWTSigningKey(globalSalt), nil
})
if err != nil {
if errors.Is(err, jwt.ErrTokenMalformed) {
return nil, fmt.Errorf("token is malformed")
} else if errors.Is(err, jwt.ErrTokenExpired) || errors.Is(err, jwt.ErrTokenNotValidYet) {
return nil, fmt.Errorf("token is either expired or not active yet")
} else {
return nil, fmt.Errorf("token invalid: %w", err)
}
}
if !token.Valid {
return nil, fmt.Errorf("token is invalid")
}
return &claims, nil
}
+479
View File
@@ -0,0 +1,479 @@
package auth_test
import (
"testing"
"time"
"pentagi/pkg/server/auth"
"pentagi/pkg/server/models"
"github.com/gin-gonic/gin"
"github.com/golang-jwt/jwt/v5"
"github.com/jinzhu/gorm"
_ "github.com/jinzhu/gorm/dialects/sqlite"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
// setupTestDB creates an in-memory SQLite database for testing
func setupTestDB(t *testing.T) *gorm.DB {
t.Helper()
db, err := gorm.Open("sqlite3", ":memory:")
require.NoError(t, err)
// Create roles table
result := db.Exec(`
CREATE TABLE roles (
id INTEGER PRIMARY KEY AUTOINCREMENT,
name TEXT NOT NULL UNIQUE
)
`)
require.NoError(t, result.Error, "Failed to create roles table")
// Create privileges table
result = db.Exec(`
CREATE TABLE privileges (
id INTEGER PRIMARY KEY AUTOINCREMENT,
role_id INTEGER NOT NULL,
name TEXT NOT NULL,
UNIQUE(role_id, name)
)
`)
require.NoError(t, result.Error, "Failed to create privileges table")
// Create api_tokens table for testing
result = db.Exec(`
CREATE TABLE api_tokens (
id INTEGER PRIMARY KEY AUTOINCREMENT,
token_id TEXT NOT NULL UNIQUE,
user_id INTEGER NOT NULL,
role_id INTEGER NOT NULL,
name TEXT,
ttl INTEGER NOT NULL,
status TEXT NOT NULL DEFAULT 'active',
created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
updated_at DATETIME DEFAULT CURRENT_TIMESTAMP,
deleted_at DATETIME
)
`)
require.NoError(t, result.Error, "Failed to create api_tokens table")
// Insert test roles
db.Exec("INSERT INTO roles (id, name) VALUES (1, 'Admin'), (2, 'User')")
// Insert test privileges for Admin role
db.Exec(`INSERT INTO privileges (role_id, name) VALUES
(1, 'users.create'),
(1, 'users.delete'),
(1, 'users.edit'),
(1, 'users.view'),
(1, 'roles.view'),
(1, 'flows.admin'),
(1, 'flows.create'),
(1, 'flows.delete'),
(1, 'flows.edit'),
(1, 'flows.view'),
(1, 'settings.tokens.create'),
(1, 'settings.tokens.view'),
(1, 'settings.tokens.edit'),
(1, 'settings.tokens.delete'),
(1, 'settings.tokens.admin')`)
// Insert test privileges for User role
db.Exec(`INSERT INTO privileges (role_id, name) VALUES
(2, 'roles.view'),
(2, 'flows.create'),
(2, 'flows.delete'),
(2, 'flows.edit'),
(2, 'flows.view'),
(2, 'settings.tokens.create'),
(2, 'settings.tokens.view'),
(2, 'settings.tokens.edit'),
(2, 'settings.tokens.delete')`)
// Create users table
db.Exec(`
CREATE TABLE users (
id INTEGER PRIMARY KEY AUTOINCREMENT,
hash TEXT NOT NULL UNIQUE,
type TEXT NOT NULL DEFAULT 'local',
mail TEXT NOT NULL UNIQUE,
name TEXT NOT NULL DEFAULT '',
status TEXT NOT NULL DEFAULT 'active',
role_id INTEGER NOT NULL DEFAULT 2,
password TEXT,
password_change_required BOOLEAN NOT NULL DEFAULT false,
provider TEXT,
created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
deleted_at DATETIME
)
`)
// Insert test users
db.Exec("INSERT INTO users (id, hash, mail, name, status, role_id) VALUES (1, 'testhash', 'user1@test.com', 'User 1', 'active', 2)")
db.Exec("INSERT INTO users (id, hash, mail, name, status, role_id) VALUES (2, 'testhash2', 'user2@test.com', 'User 2', 'active', 2)")
// Create user_preferences table
db.Exec(`
CREATE TABLE user_preferences (
id INTEGER PRIMARY KEY AUTOINCREMENT,
user_id INTEGER NOT NULL UNIQUE,
preferences TEXT NOT NULL DEFAULT '{"favoriteFlows": []}',
created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
updated_at DATETIME DEFAULT CURRENT_TIMESTAMP,
FOREIGN KEY (user_id) REFERENCES users(id) ON DELETE CASCADE
)
`)
// Insert preferences for test users
db.Exec("INSERT INTO user_preferences (user_id, preferences) VALUES (1, '{\"favoriteFlows\": []}')")
db.Exec("INSERT INTO user_preferences (user_id, preferences) VALUES (2, '{\"favoriteFlows\": []}')")
time.Sleep(200 * time.Millisecond) // wait for database to be ready
return db
}
func TestValidateAPIToken(t *testing.T) {
globalSalt := "test_salt"
testCases := []struct {
name string
setup func() string
expectError bool
errorMsg string
}{
{
name: "valid token",
setup: func() string {
claims := models.APITokenClaims{
TokenID: "abc123xyz9",
RID: 2,
UID: 1,
UHASH: "testhash",
RegisteredClaims: jwt.RegisteredClaims{
ExpiresAt: jwt.NewNumericDate(time.Now().Add(1 * time.Hour)),
IssuedAt: jwt.NewNumericDate(time.Now()),
Subject: "api_token",
},
}
token := jwt.NewWithClaims(jwt.SigningMethodHS256, claims)
tokenString, _ := token.SignedString(auth.MakeJWTSigningKey(globalSalt))
return tokenString
},
expectError: false,
},
{
name: "expired token",
setup: func() string {
claims := models.APITokenClaims{
TokenID: "abc123xyz9",
RID: 2,
UID: 1,
UHASH: "testhash",
RegisteredClaims: jwt.RegisteredClaims{
ExpiresAt: jwt.NewNumericDate(time.Now().Add(-1 * time.Hour)),
IssuedAt: jwt.NewNumericDate(time.Now().Add(-2 * time.Hour)),
Subject: "api_token",
},
}
token := jwt.NewWithClaims(jwt.SigningMethodHS256, claims)
tokenString, _ := token.SignedString(auth.MakeJWTSigningKey(globalSalt))
return tokenString
},
expectError: true,
errorMsg: "expired",
},
{
name: "invalid signature",
setup: func() string {
claims := models.APITokenClaims{
TokenID: "abc123xyz9",
RID: 2,
UID: 1,
UHASH: "testhash",
RegisteredClaims: jwt.RegisteredClaims{
ExpiresAt: jwt.NewNumericDate(time.Now().Add(1 * time.Hour)),
IssuedAt: jwt.NewNumericDate(time.Now()),
Subject: "api_token",
},
}
token := jwt.NewWithClaims(jwt.SigningMethodHS256, claims)
tokenString, _ := token.SignedString([]byte("wrong_key"))
return tokenString
},
expectError: true,
errorMsg: "invalid",
},
{
name: "malformed token",
setup: func() string {
return "not.a.valid.jwt.token"
},
expectError: true,
errorMsg: "malformed",
},
}
for _, tc := range testCases {
t.Run(tc.name, func(t *testing.T) {
tokenString := tc.setup()
claims, err := auth.ValidateAPIToken(tokenString, globalSalt)
if tc.expectError {
assert.Error(t, err)
if tc.errorMsg != "" {
assert.Contains(t, err.Error(), tc.errorMsg)
}
assert.Nil(t, claims)
} else {
assert.NoError(t, err)
assert.NotNil(t, claims)
assert.Equal(t, "abc123xyz9", claims.TokenID)
assert.Equal(t, uint64(1), claims.UID)
assert.Equal(t, uint64(2), claims.RID)
}
})
}
}
func TestAPITokenAuthentication_CacheExpiration(t *testing.T) {
db := setupTestDB(t)
defer db.Close()
// Create cache with short TTL for testing
tokenCache := auth.NewTokenCache(db)
tokenCache.SetTTL(100 * time.Millisecond)
userCache := auth.NewUserCache(db)
authMiddleware := auth.NewAuthMiddleware("/base/url", "test", tokenCache, userCache)
// Create active token
tokenID, err := auth.GenerateTokenID()
require.NoError(t, err)
apiToken := models.APIToken{
TokenID: tokenID,
UserID: 1,
RoleID: 2,
TTL: 3600,
Status: models.TokenStatusActive,
}
err = db.Create(&apiToken).Error
require.NoError(t, err)
// Create JWT
claims := models.APITokenClaims{
TokenID: tokenID,
RID: 2,
UID: 1,
UHASH: "testhash",
RegisteredClaims: jwt.RegisteredClaims{
ExpiresAt: jwt.NewNumericDate(time.Now().Add(1 * time.Hour)),
IssuedAt: jwt.NewNumericDate(time.Now()),
Subject: "api_token",
},
}
token := jwt.NewWithClaims(jwt.SigningMethodHS256, claims)
tokenString, err := token.SignedString(auth.MakeJWTSigningKey("test"))
require.NoError(t, err)
server := newTestServer(t, "/test", db, authMiddleware.AuthTokenRequired)
defer server.Close()
// First call: should work (status active, cached)
assert.True(t, server.CallAndGetStatus(t, "Bearer "+tokenString))
// Revoke token in DB
db.Model(&apiToken).Update("status", models.TokenStatusRevoked)
// Second call: should still work (cache not expired)
assert.True(t, server.CallAndGetStatus(t, "Bearer "+tokenString))
// Wait for cache to expire
time.Sleep(150 * time.Millisecond)
// Third call: should fail (cache expired, reads from DB)
assert.False(t, server.CallAndGetStatus(t, "Bearer "+tokenString))
}
func TestAPITokenAuthentication_DefaultSalt(t *testing.T) {
db := setupTestDB(t)
defer db.Close()
testCases := []struct {
name string
globalSalt string
shouldSkip bool
}{
{
name: "default salt 'salt'",
globalSalt: "salt",
shouldSkip: true,
},
{
name: "empty salt",
globalSalt: "",
shouldSkip: true,
},
{
name: "custom salt",
globalSalt: "custom_secure_salt",
shouldSkip: false,
},
}
for _, tc := range testCases {
t.Run(tc.name, func(t *testing.T) {
tokenCache := auth.NewTokenCache(db)
userCache := auth.NewUserCache(db)
authMiddleware := auth.NewAuthMiddleware("/base/url", tc.globalSalt, tokenCache, userCache)
// Create a token (even with default salt, for testing)
tokenID, err := auth.GenerateTokenID()
require.NoError(t, err)
claims := models.APITokenClaims{
TokenID: tokenID,
RID: 2,
UID: 1,
UHASH: "testhash",
RegisteredClaims: jwt.RegisteredClaims{
ExpiresAt: jwt.NewNumericDate(time.Now().Add(1 * time.Hour)),
IssuedAt: jwt.NewNumericDate(time.Now()),
Subject: "api_token",
},
}
token := jwt.NewWithClaims(jwt.SigningMethodHS256, claims)
tokenString, _ := token.SignedString(auth.MakeJWTSigningKey(tc.globalSalt))
server := newTestServer(t, "/test", db, authMiddleware.AuthTokenRequired)
defer server.Close()
// With default salt, token validation should be skipped
result := server.CallAndGetStatus(t, "Bearer "+tokenString)
if tc.shouldSkip {
// Should skip token auth and try cookie (which will fail)
assert.False(t, result)
} else {
// With custom salt but no DB record, should fail with "not found"
assert.False(t, result)
}
})
}
}
func TestAPITokenAuthentication_SoftDelete(t *testing.T) {
db := setupTestDB(t)
defer db.Close()
tokenCache := auth.NewTokenCache(db)
userCache := auth.NewUserCache(db)
authMiddleware := auth.NewAuthMiddleware("/base/url", "test", tokenCache, userCache)
// Create token
tokenID, err := auth.GenerateTokenID()
require.NoError(t, err)
apiToken := models.APIToken{
TokenID: tokenID,
UserID: 1,
RoleID: 2,
TTL: 3600,
Status: models.TokenStatusActive,
}
err = db.Create(&apiToken).Error
require.NoError(t, err)
// Create JWT
claims := models.APITokenClaims{
TokenID: tokenID,
RID: 2,
UID: 1,
UHASH: "testhash",
RegisteredClaims: jwt.RegisteredClaims{
ExpiresAt: jwt.NewNumericDate(time.Now().Add(1 * time.Hour)),
IssuedAt: jwt.NewNumericDate(time.Now()),
Subject: "api_token",
},
}
token := jwt.NewWithClaims(jwt.SigningMethodHS256, claims)
tokenString, err := token.SignedString(auth.MakeJWTSigningKey("test"))
require.NoError(t, err)
server := newTestServer(t, "/test", db, authMiddleware.AuthTokenRequired)
defer server.Close()
// Should work initially
assert.True(t, server.CallAndGetStatus(t, "Bearer "+tokenString))
// Soft delete
now := time.Now()
db.Model(&apiToken).Update("deleted_at", now)
tokenCache.Invalidate(tokenID)
// Should fail after soft delete
assert.False(t, server.CallAndGetStatus(t, "Bearer "+tokenString))
}
func TestAPITokenAuthentication_AlgNoneAttack(t *testing.T) {
db := setupTestDB(t)
defer db.Close()
tokenCache := auth.NewTokenCache(db)
userCache := auth.NewUserCache(db)
authMiddleware := auth.NewAuthMiddleware("/base/url", "test", tokenCache, userCache)
tokenID, err := auth.GenerateTokenID()
require.NoError(t, err)
// Create token with "none" algorithm (security attack)
claims := models.APITokenClaims{
TokenID: tokenID,
RID: 2,
UID: 1,
UHASH: "testhash",
RegisteredClaims: jwt.RegisteredClaims{
ExpiresAt: jwt.NewNumericDate(time.Now().Add(1 * time.Hour)),
IssuedAt: jwt.NewNumericDate(time.Now()),
Subject: "api_token",
},
}
// Try to use "none" algorithm
token := jwt.NewWithClaims(jwt.SigningMethodNone, claims)
tokenString, err := token.SignedString(jwt.UnsafeAllowNoneSignatureType)
require.NoError(t, err)
server := newTestServer(t, "/test", db, authMiddleware.AuthTokenRequired)
defer server.Close()
// Should reject "none" algorithm
assert.False(t, server.CallAndGetStatus(t, "Bearer "+tokenString))
}
func TestAPITokenAuthentication_LegacyProtoToken(t *testing.T) {
db := setupTestDB(t)
defer db.Close()
tokenCache := auth.NewTokenCache(db)
userCache := auth.NewUserCache(db)
authMiddleware := auth.NewAuthMiddleware("/base/url", "test", tokenCache, userCache)
server := newTestServer(t, "/test", db, authMiddleware.AuthTokenRequired)
defer server.Close()
// Authorize with cookie to get legacy proto token
server.Authorize(t, []string{auth.PrivilegeAutomation})
legacyToken := server.GetToken(t)
require.NotEmpty(t, legacyToken)
// Unauthorize cookie
server.Unauthorize(t)
// Legacy proto token should still work (fallback mechanism)
server.SetSessionCheckFunc(func(t *testing.T, c *gin.Context) {
t.Helper()
assert.Equal(t, uint64(1), c.GetUint64("uid"))
assert.Equal(t, "automation", c.GetString("cpt"))
})
assert.True(t, server.CallAndGetStatus(t, "Bearer "+legacyToken))
}
+237
View File
@@ -0,0 +1,237 @@
package auth
import (
"errors"
"fmt"
"slices"
"strings"
"time"
"pentagi/pkg/server/models"
"pentagi/pkg/server/rdb"
"pentagi/pkg/server/response"
"github.com/gin-contrib/sessions"
"github.com/gin-gonic/gin"
"github.com/jinzhu/gorm"
)
type authResult int
const (
authResultOk authResult = iota
authResultSkip
authResultFail
authResultAbort
)
type AuthMiddleware struct {
globalSalt string
tokenCache *TokenCache
userCache *UserCache
}
func NewAuthMiddleware(baseURL, globalSalt string, tokenCache *TokenCache, userCache *UserCache) *AuthMiddleware {
return &AuthMiddleware{
globalSalt: globalSalt,
tokenCache: tokenCache,
userCache: userCache,
}
}
func (p *AuthMiddleware) AuthUserRequired(c *gin.Context) {
p.tryAuth(c, true, p.tryUserCookieAuthentication)
}
func (p *AuthMiddleware) AuthTokenRequired(c *gin.Context) {
p.tryAuth(c, true, p.tryProtoTokenAuthentication, p.tryUserCookieAuthentication)
}
func (p *AuthMiddleware) TryAuth(c *gin.Context) {
p.tryAuth(c, false, p.tryProtoTokenAuthentication, p.tryUserCookieAuthentication)
}
func (p *AuthMiddleware) tryAuth(
c *gin.Context,
withFail bool,
authMethods ...func(c *gin.Context) (authResult, error),
) {
if c.IsAborted() {
return
}
result := authResultSkip
var authErr error
for _, authMethod := range authMethods {
result, authErr = authMethod(c)
if c.IsAborted() || result == authResultAbort {
return
}
if result != authResultSkip {
break
}
}
if withFail && result != authResultOk {
response.Error(c, response.ErrAuthRequired, authErr)
return
}
c.Next()
}
func (p *AuthMiddleware) tryUserCookieAuthentication(c *gin.Context) (authResult, error) {
sessionObject, exists := c.Get(sessions.DefaultKey)
if !exists {
return authResultSkip, errors.New("can't find session object")
}
session, ok := sessionObject.(sessions.Session)
if !ok {
return authResultFail, errors.New("not a session object")
}
uid := session.Get("uid")
uhash := session.Get("uhash")
rid := session.Get("rid")
prm := session.Get("prm")
exp := session.Get("exp")
gtm := session.Get("gtm")
tid := session.Get("tid")
uname := session.Get("uname")
for _, attr := range []any{uid, rid, prm, exp, gtm, uname, uhash, tid} {
if attr == nil {
return authResultFail, errors.New("cookie claim invalid")
}
}
prms, ok := prm.([]string)
if !ok {
return authResultFail, errors.New("no permissions granted")
}
// Verify session expiration
expVal, ok := exp.(int64)
if !ok {
return authResultFail, errors.New("token claim invalid")
}
if time.Now().Unix() > expVal {
return authResultFail, errors.New("session expired")
}
// Verify user hash matches database
userID := uid.(uint64)
sessionHash := uhash.(string)
dbHash, userStatus, err := p.userCache.GetUserHash(userID)
if err != nil {
if errors.Is(err, gorm.ErrRecordNotFound) {
return authResultFail, errors.New("user has been deleted")
}
return authResultFail, fmt.Errorf("error checking user status: %w", err)
}
switch userStatus {
case models.UserStatusBlocked:
return authResultFail, errors.New("user has been blocked")
case models.UserStatusCreated:
return authResultFail, errors.New("user is not ready")
case models.UserStatusActive:
}
if dbHash != sessionHash {
return authResultFail, errors.New("user hash mismatch - session invalid for this installation")
}
c.Set("prm", prms)
c.Set("uid", userID)
c.Set("uhash", sessionHash)
c.Set("rid", rid.(uint64))
c.Set("exp", exp.(int64))
c.Set("gtm", gtm.(int64))
c.Set("tid", tid.(string))
c.Set("uname", uname.(string))
if slices.Contains(prms, PrivilegeAutomation) {
c.Set("cpt", "automation")
}
return authResultOk, nil
}
const PrivilegeAutomation = "pentagi.automation"
func (p *AuthMiddleware) tryProtoTokenAuthentication(c *gin.Context) (authResult, error) {
authHeader := c.Request.Header.Get("Authorization")
if authHeader == "" {
return authResultSkip, errors.New("token required")
}
if !strings.HasPrefix(authHeader, "Bearer ") {
return authResultSkip, errors.New("bearer scheme must be used")
}
token := authHeader[7:]
if token == "" {
return authResultSkip, errors.New("token can't be empty")
}
// skip validation if using default salt (for backward compatibility)
if p.globalSalt == "" || p.globalSalt == "salt" {
return authResultSkip, errors.New("token validation disabled with default salt")
}
// try to validate as API token first (new format with JWT signing key)
apiClaims, apiErr := ValidateAPIToken(token, p.globalSalt)
if apiErr != nil {
return authResultFail, errors.New("token is invalid")
}
// check token status and get privileges through cache
status, privileges, err := p.tokenCache.GetStatus(apiClaims.TokenID)
if err != nil {
if errors.Is(err, gorm.ErrRecordNotFound) {
return authResultFail, errors.New("token not found in database")
}
return authResultFail, fmt.Errorf("error checking token status: %w", err)
}
if status != models.TokenStatusActive {
return authResultFail, errors.New("token has been revoked")
}
// Verify user hash matches database
dbHash, userStatus, err := p.userCache.GetUserHash(apiClaims.UID)
if err != nil {
if errors.Is(err, gorm.ErrRecordNotFound) {
return authResultFail, errors.New("user has been deleted")
}
return authResultFail, fmt.Errorf("error checking user status: %w", err)
}
if userStatus == models.UserStatusBlocked {
return authResultFail, errors.New("user has been blocked")
}
if dbHash != apiClaims.UHASH {
return authResultFail, errors.New("user hash mismatch - token invalid for this installation")
}
// generate UUID from user hash (fallback to empty string if hash is invalid)
uuid, err := rdb.MakeUuidStrFromHash(apiClaims.UHASH)
if err != nil {
// Use empty UUID for invalid hashes (e.g., in tests)
uuid = ""
}
// set session fields similar to regular login
c.Set("uid", apiClaims.UID)
c.Set("uhash", apiClaims.UHASH)
c.Set("rid", apiClaims.RID)
c.Set("tid", models.UserTypeAPI.String())
c.Set("prm", privileges)
c.Set("gtm", time.Now().Unix())
c.Set("exp", apiClaims.ExpiresAt.Unix())
c.Set("uuid", uuid)
c.Set("cpt", "automation")
return authResultOk, nil
}
@@ -0,0 +1,337 @@
package auth_test
import (
"bytes"
"io"
"math/rand"
"net/http"
"net/http/cookiejar"
"net/http/httptest"
"net/url"
"strconv"
"testing"
"time"
"pentagi/pkg/server/auth"
"pentagi/pkg/server/models"
"github.com/gin-contrib/sessions"
"github.com/gin-contrib/sessions/cookie"
"github.com/gin-gonic/gin"
"github.com/jinzhu/gorm"
_ "github.com/jinzhu/gorm/dialects/sqlite"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
func TestAuthTokenProtoRequiredAuthWithCookie(t *testing.T) {
db := setupTestDB(t)
defer db.Close()
tokenCache := auth.NewTokenCache(db)
userCache := auth.NewUserCache(db)
authMiddleware := auth.NewAuthMiddleware("/base/url", "test", tokenCache, userCache)
t.Run("test URL", func(t *testing.T) {
server := newTestServer(t, "/test", db, authMiddleware.AuthTokenRequired)
defer server.Close()
assert.False(t, server.CallAndGetStatus(t))
server.SetSessionCheckFunc(func(t *testing.T, c *gin.Context) {
t.Helper()
assert.Equal(t, "", c.GetString("cpt"))
})
server.Authorize(t, []string{})
assert.True(t, server.CallAndGetStatus(t))
server.Authorize(t, []string{"wrong.permission"})
assert.True(t, server.CallAndGetStatus(t))
server.SetSessionCheckFunc(func(t *testing.T, c *gin.Context) {
t.Helper()
assert.Equal(t, "automation", c.GetString("cpt"))
})
server.Authorize(t, []string{auth.PrivilegeAutomation})
assert.True(t, server.CallAndGetStatus(t))
server.Authorize(t, []string{"wrong.permission", auth.PrivilegeAutomation})
assert.True(t, server.CallAndGetStatus(t))
})
}
func TestAuthTokenProtoRequiredAuthWithToken(t *testing.T) {
db := setupTestDB(t)
defer db.Close()
tokenCache := auth.NewTokenCache(db)
userCache := auth.NewUserCache(db)
authMiddleware := auth.NewAuthMiddleware("/base/url", "test", tokenCache, userCache)
server := newTestServer(t, "/test", db, authMiddleware.AuthTokenRequired)
defer server.Close()
server.Authorize(t, []string{auth.PrivilegeAutomation})
token := server.GetToken(t)
require.NotEmpty(t, token)
server.Unauthorize(t)
assert.False(t, server.CallAndGetStatus(t))
assert.False(t, server.CallAndGetStatus(t, token))
assert.False(t, server.CallAndGetStatus(t, "not a bearer "+token))
assert.False(t, server.CallAndGetStatus(t, "Bearer"+token))
assert.False(t, server.CallAndGetStatus(t, "Bearer not_a_token"))
server.SetSessionCheckFunc(func(t *testing.T, c *gin.Context) {
t.Helper()
assert.Equal(t, uint64(1), c.GetUint64("uid"))
assert.Equal(t, uint64(2), c.GetUint64("rid"))
assert.NotNil(t, c.GetStringSlice("prm"))
// gtm and exp should now be set for API tokens
gtm := c.GetInt64("gtm")
assert.Greater(t, gtm, int64(0), "GTM should be set")
exp := c.GetInt64("exp")
assert.Greater(t, exp, gtm, "EXP should be greater than GTM")
// uuid will be empty for invalid hash (test uses "123" which is not valid MD5)
assert.NotNil(t, c.GetString("uuid"))
assert.Equal(t, "automation", c.GetString("cpt"))
assert.Empty(t, c.GetString("uname"))
})
assert.True(t, server.CallAndGetStatus(t, "Bearer "+token))
}
func TestAuthRequiredAuthWithCookie(t *testing.T) {
db := setupTestDB(t)
defer db.Close()
tokenCache := auth.NewTokenCache(db)
userCache := auth.NewUserCache(db)
authMiddleware := auth.NewAuthMiddleware("/base/url", "test", tokenCache, userCache)
server := newTestServer(t, "/test", db, authMiddleware.AuthUserRequired)
defer server.Close()
server.SetSessionCheckFunc(func(t *testing.T, c *gin.Context) {
t.Helper()
assert.Equal(t, uint64(1), c.GetUint64("uid"))
assert.Equal(t, uint64(2), c.GetUint64("rid"))
assert.NotNil(t, c.GetStringSlice("prm"))
assert.NotNil(t, c.GetInt64("gtm"))
assert.NotNil(t, c.GetInt64("exp"))
assert.Empty(t, c.GetString("uuid"))
assert.Equal(t, "User 1", c.GetString("uname"))
})
assert.False(t, server.CallAndGetStatus(t))
server.Authorize(t, []string{"some.permission"})
assert.True(t, server.CallAndGetStatus(t))
}
type testServer struct {
testEndpoint string
client *http.Client
calls map[string]struct{}
sessionCheckFunc func(t *testing.T, c *gin.Context)
db *gorm.DB
*httptest.Server
}
func newTestServer(t *testing.T, testEndpoint string, db *gorm.DB, middlewares ...gin.HandlerFunc) *testServer {
t.Helper()
server := &testServer{
db: db,
}
router := gin.New()
globalSalt := "test"
cookieStore := cookie.NewStore(auth.MakeCookieStoreKey(globalSalt)...)
router.Use(sessions.Sessions("auth", cookieStore))
server.calls = map[string]struct{}{}
if testEndpoint == "" {
testEndpoint = "/test"
}
server.testEndpoint = testEndpoint
router.GET("/auth", func(c *gin.Context) {
t.Helper()
privs, _ := c.GetQueryArray("privileges")
expString, ok := c.GetQuery("expiration")
assert.True(t, ok)
exp, err := strconv.Atoi(expString)
assert.NoError(t, err)
setTestSession(t, c, privs, exp)
})
authRoutes := router.Group("")
for _, middleware := range middlewares {
authRoutes.Use(middleware)
}
authRoutes.GET(server.testEndpoint, func(c *gin.Context) {
t.Helper()
id, _ := c.GetQuery("id")
require.NotEmpty(t, id)
if server.sessionCheckFunc != nil {
server.sessionCheckFunc(t, c)
}
server.calls[id] = struct{}{}
})
authRoutes.GET("/auth_token", func(c *gin.Context) {
t.Helper()
tokenID, err := auth.GenerateTokenID()
require.NoError(t, err)
uhash := "testhash"
uid := uint64(1)
rid := uint64(2)
ttl := uint64(3600)
claims := auth.MakeAPITokenClaims(tokenID, uhash, uid, rid, ttl)
token, err := auth.MakeAPIToken(globalSalt, claims)
require.NoError(t, err)
db.Create(&models.APIToken{
TokenID: tokenID,
UserID: uid,
RoleID: rid,
TTL: ttl,
Status: models.TokenStatusActive,
})
c.Writer.Write([]byte(token))
})
server.Server = httptest.NewServer(router)
server.client = server.Client()
jar, err := cookiejar.New(nil)
require.NoError(t, err)
server.client.Jar = jar
return server
}
func (s *testServer) Authorize(t *testing.T, privileges []string) {
t.Helper()
request, err := http.NewRequest(http.MethodGet, s.URL+"/auth", nil)
require.NoError(t, err)
query := url.Values{}
for _, p := range privileges {
query.Add("privileges", p)
}
query.Add("expiration", strconv.Itoa(5*60))
request.URL.RawQuery = query.Encode()
resp, err := s.client.Do(request)
require.NoError(t, err)
assert.Equal(t, http.StatusOK, resp.StatusCode)
}
func (s *testServer) GetToken(t *testing.T) string {
t.Helper()
request, err := http.NewRequest(http.MethodGet, s.URL+"/auth_token", nil)
require.NoError(t, err)
resp, err := s.client.Do(request)
require.NoError(t, err)
assert.Equal(t, http.StatusOK, resp.StatusCode)
token, err := io.ReadAll(resp.Body)
require.NoError(t, err)
return string(token)
}
func (s *testServer) SetSessionCheckFunc(f func(t *testing.T, c *gin.Context)) {
s.sessionCheckFunc = f
}
func (s *testServer) Unauthorize(t *testing.T) {
t.Helper()
request, err := http.NewRequest(http.MethodGet, s.URL+"/auth", nil)
require.NoError(t, err)
query := url.Values{}
query.Add("expiration", strconv.Itoa(-1))
request.URL.RawQuery = query.Encode()
resp, err := s.client.Do(request)
require.NoError(t, err)
assert.Equal(t, http.StatusOK, resp.StatusCode)
}
func (s *testServer) TestCall(t *testing.T, token ...string) (string, bool) {
t.Helper()
id := strconv.Itoa(rand.Int())
request, err := http.NewRequest(http.MethodGet, s.URL+s.testEndpoint+"?id="+id, nil)
require.NoError(t, err)
if len(token) == 1 {
request.Header.Add("Authorization", token[0])
}
resp, err := s.client.Do(request)
require.NoError(t, err)
assert.True(t, resp.StatusCode == http.StatusOK ||
resp.StatusCode == http.StatusForbidden)
return id, resp.StatusCode == http.StatusOK
}
func (s *testServer) TestCallWithData(t *testing.T, data string) (string, bool) {
t.Helper()
id := strconv.Itoa(rand.Int())
request, err := http.NewRequest(http.MethodGet, s.URL+s.testEndpoint+"?id="+id, bytes.NewBufferString(data))
require.NoError(t, err)
resp, err := s.client.Do(request)
require.NoError(t, err)
assert.True(t, resp.StatusCode == http.StatusOK ||
resp.StatusCode == http.StatusForbidden)
return id, resp.StatusCode == http.StatusOK
}
func (s *testServer) Called(id string) bool {
_, ok := s.calls[id]
return ok
}
func (s *testServer) CallAndGetStatus(t *testing.T, token ...string) bool {
t.Helper()
id, ok := s.TestCall(t, token...)
assert.Equal(t, ok, s.Called(id))
return ok
}
func setTestSession(t *testing.T, c *gin.Context, privileges []string, expires int) {
t.Helper()
session := sessions.Default(c)
session.Set("uid", uint64(1))
session.Set("uhash", "testhash")
session.Set("rid", uint64(2))
session.Set("tid", models.UserTypeLocal.String())
session.Set("prm", privileges)
session.Set("gtm", time.Now().Unix())
session.Set("exp", time.Now().Add(time.Duration(expires)*time.Second).Unix())
session.Set("uuid", "uuid1")
session.Set("uname", "User 1")
session.Options(sessions.Options{
HttpOnly: true,
MaxAge: expires,
})
require.NoError(t, session.Save())
}
+890
View File
@@ -0,0 +1,890 @@
package auth_test
import (
"sync"
"testing"
"time"
"pentagi/pkg/server/auth"
"pentagi/pkg/server/models"
"github.com/gin-gonic/gin"
"github.com/golang-jwt/jwt/v5"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
// TestEndToEndAPITokenFlow tests complete flow from creation to usage
func TestEndToEndAPITokenFlow(t *testing.T) {
db := setupTestDB(t)
defer db.Close()
tokenCache := auth.NewTokenCache(db)
userCache := auth.NewUserCache(db)
authMiddleware := auth.NewAuthMiddleware("/base/url", "test_salt", tokenCache, userCache)
testCases := []struct {
name string
tokenID string
status models.TokenStatus
shouldPass bool
errorContains string
}{
{
name: "active token authenticates successfully",
tokenID: "active123",
status: models.TokenStatusActive,
shouldPass: true,
},
{
name: "revoked token is rejected",
tokenID: "revoked456",
status: models.TokenStatusRevoked,
shouldPass: false,
errorContains: "revoked",
},
}
for _, tc := range testCases {
t.Run(tc.name, func(t *testing.T) {
// Create token in database
apiToken := models.APIToken{
TokenID: tc.tokenID,
UserID: 1,
RoleID: 2,
TTL: 3600,
Status: tc.status,
}
err := db.Create(&apiToken).Error
require.NoError(t, err)
// Create JWT token
claims := models.APITokenClaims{
TokenID: tc.tokenID,
RID: 2,
UID: 1,
UHASH: "testhash",
RegisteredClaims: jwt.RegisteredClaims{
ExpiresAt: jwt.NewNumericDate(time.Now().Add(1 * time.Hour)),
IssuedAt: jwt.NewNumericDate(time.Now()),
Subject: "api_token",
},
}
token := jwt.NewWithClaims(jwt.SigningMethodHS256, claims)
tokenString, err := token.SignedString(auth.MakeJWTSigningKey("test_salt"))
require.NoError(t, err)
// Test authentication
server := newTestServer(t, "/protected", db, authMiddleware.AuthTokenRequired)
defer server.Close()
success := server.CallAndGetStatus(t, "Bearer "+tokenString)
assert.Equal(t, tc.shouldPass, success)
})
}
}
// TestAPIToken_RoleIsolation verifies that token inherits creator's role
func TestAPIToken_RoleIsolation(t *testing.T) {
testCases := []struct {
name string
creatorRole uint64
tokenRole uint64
expectMatch bool
}{
{
name: "user creates token with user role",
creatorRole: 2,
tokenRole: 2,
expectMatch: true,
},
{
name: "admin creates token with admin role",
creatorRole: 1,
tokenRole: 1,
expectMatch: true,
},
}
for _, tc := range testCases {
t.Run(tc.name, func(t *testing.T) {
tokenID, err := auth.GenerateTokenID()
require.NoError(t, err)
// Create JWT with specific role
claims := models.APITokenClaims{
TokenID: tokenID,
RID: tc.tokenRole,
UID: 1,
UHASH: "testhash",
RegisteredClaims: jwt.RegisteredClaims{
ExpiresAt: jwt.NewNumericDate(time.Now().Add(1 * time.Hour)),
IssuedAt: jwt.NewNumericDate(time.Now()),
Subject: "api_token",
},
}
token := jwt.NewWithClaims(jwt.SigningMethodHS256, claims)
tokenString, err := token.SignedString(auth.MakeJWTSigningKey("test"))
require.NoError(t, err)
// Validate and check role
validated, err := auth.ValidateAPIToken(tokenString, "test")
require.NoError(t, err)
if tc.expectMatch {
assert.Equal(t, tc.tokenRole, validated.RID)
}
})
}
}
// TestAPIToken_SignatureVerification tests various signature attacks
func TestAPIToken_SignatureVerification(t *testing.T) {
correctSalt := "correct_salt"
wrongSalt := "wrong_salt"
testCases := []struct {
name string
signSalt string
verifySalt string
expectValid bool
errorContains string
}{
{
name: "matching salt - valid",
signSalt: correctSalt,
verifySalt: correctSalt,
expectValid: true,
},
{
name: "mismatched salt - invalid",
signSalt: correctSalt,
verifySalt: wrongSalt,
expectValid: false,
errorContains: "invalid",
},
}
for _, tc := range testCases {
t.Run(tc.name, func(t *testing.T) {
tokenID, err := auth.GenerateTokenID()
require.NoError(t, err)
claims := models.APITokenClaims{
TokenID: tokenID,
RID: 2,
UID: 1,
UHASH: "testhash",
RegisteredClaims: jwt.RegisteredClaims{
ExpiresAt: jwt.NewNumericDate(time.Now().Add(1 * time.Hour)),
IssuedAt: jwt.NewNumericDate(time.Now()),
Subject: "api_token",
},
}
token := jwt.NewWithClaims(jwt.SigningMethodHS256, claims)
tokenString, err := token.SignedString(auth.MakeJWTSigningKey(tc.signSalt))
require.NoError(t, err)
validated, err := auth.ValidateAPIToken(tokenString, tc.verifySalt)
if tc.expectValid {
assert.NoError(t, err)
assert.NotNil(t, validated)
} else {
assert.Error(t, err)
if tc.errorContains != "" {
assert.Contains(t, err.Error(), tc.errorContains)
}
}
})
}
}
// TestAPIToken_CacheInvalidation verifies cache invalidation scenarios
func TestAPIToken_CacheInvalidation(t *testing.T) {
db := setupTestDB(t)
defer db.Close()
tokenCache := auth.NewTokenCache(db)
// Create token
tokenID, err := auth.GenerateTokenID()
require.NoError(t, err)
apiToken := models.APIToken{
TokenID: tokenID,
UserID: 1,
RoleID: 2,
TTL: 3600,
Status: models.TokenStatusActive,
}
err = db.Create(&apiToken).Error
require.NoError(t, err)
// Load into cache
status1, _, err := tokenCache.GetStatus(tokenID)
require.NoError(t, err)
assert.Equal(t, models.TokenStatusActive, status1)
// Update in DB
db.Model(&apiToken).Update("status", models.TokenStatusRevoked)
// Should still return active from cache
status2, _, err := tokenCache.GetStatus(tokenID)
require.NoError(t, err)
assert.Equal(t, models.TokenStatusActive, status2, "Cache should return stale value")
// Invalidate cache
tokenCache.Invalidate(tokenID)
// Should now return revoked from DB
status3, _, err := tokenCache.GetStatus(tokenID)
require.NoError(t, err)
assert.Equal(t, models.TokenStatusRevoked, status3, "Cache should be refreshed from DB")
}
// TestAPIToken_ConcurrentAccess tests thread-safety of cache
func TestAPIToken_ConcurrentAccess(t *testing.T) {
db := setupTestDB(t)
defer db.Close()
tokenCache := auth.NewTokenCache(db)
// Create multiple tokens
tokenIDs := make([]string, 10)
for i := range 10 {
tokenID, err := auth.GenerateTokenID()
require.NoError(t, err)
tokenIDs[i] = tokenID
apiToken := models.APIToken{
TokenID: tokenID,
UserID: 1,
RoleID: 2,
TTL: 3600,
Status: models.TokenStatusActive,
}
err = db.Create(&apiToken).Error
require.NoError(t, err)
}
// Verify tokens were created
var count int
db.Model(&models.APIToken{}).Where("deleted_at IS NULL").Count(&count)
require.Equal(t, 10, count)
// Warm up cache
for i := range 10 {
status, _, err := tokenCache.GetStatus(tokenIDs[i])
require.NoError(t, err)
assert.Equal(t, models.TokenStatusActive, status)
}
// Concurrent cache access using channels for error reporting
type testResult struct {
success bool
err error
}
results := make(chan testResult, 10)
var wg sync.WaitGroup
wg.Add(10)
for i := range 10 {
go func(tokenID string) {
defer wg.Done()
for range 100 {
status, _, err := tokenCache.GetStatus(tokenID)
if err != nil {
results <- testResult{success: false, err: err}
return
}
if status != models.TokenStatusActive {
results <- testResult{success: false, err: assert.AnError}
return
}
}
results <- testResult{success: true, err: nil}
}(tokenIDs[i])
}
wg.Wait()
close(results)
// Wait and check all results
for result := range results {
assert.NoError(t, result.err)
assert.True(t, result.success, "Goroutine should complete successfully")
}
}
// TestAPIToken_JSONStructure verifies JWT payload structure
func TestAPIToken_JSONStructure(t *testing.T) {
tokenID, err := auth.GenerateTokenID()
require.NoError(t, err)
claims := models.APITokenClaims{
TokenID: tokenID,
RID: 2,
UID: 1,
UHASH: "testhash",
RegisteredClaims: jwt.RegisteredClaims{
ExpiresAt: jwt.NewNumericDate(time.Now().Add(1 * time.Hour)),
IssuedAt: jwt.NewNumericDate(time.Now()),
Subject: "api_token",
},
}
token := jwt.NewWithClaims(jwt.SigningMethodHS256, claims)
tokenString, err := token.SignedString(auth.MakeJWTSigningKey("test"))
require.NoError(t, err)
// Parse and verify all fields
parsed, err := auth.ValidateAPIToken(tokenString, "test")
require.NoError(t, err)
assert.Equal(t, tokenID, parsed.TokenID, "TokenID should match")
assert.Equal(t, uint64(2), parsed.RID, "RID should match")
assert.Equal(t, uint64(1), parsed.UID, "UID should match")
assert.Equal(t, "testhash", parsed.UHASH, "UHASH should match")
assert.Equal(t, "api_token", parsed.Subject, "Subject should match")
assert.NotNil(t, parsed.ExpiresAt, "ExpiresAt should be set")
assert.NotNil(t, parsed.IssuedAt, "IssuedAt should be set")
}
// TestAPIToken_Expiration verifies TTL enforcement
func TestAPIToken_Expiration(t *testing.T) {
testCases := []struct {
name string
ttl time.Duration
expectValid bool
}{
{
name: "future expiration - valid",
ttl: 1 * time.Hour,
expectValid: true,
},
{
name: "past expiration - invalid",
ttl: -1 * time.Hour,
expectValid: false,
},
{
name: "just expired - invalid",
ttl: -1 * time.Second,
expectValid: false,
},
}
for _, tc := range testCases {
t.Run(tc.name, func(t *testing.T) {
tokenID, err := auth.GenerateTokenID()
require.NoError(t, err)
claims := models.APITokenClaims{
TokenID: tokenID,
RID: 2,
UID: 1,
UHASH: "testhash",
RegisteredClaims: jwt.RegisteredClaims{
ExpiresAt: jwt.NewNumericDate(time.Now().Add(tc.ttl)),
IssuedAt: jwt.NewNumericDate(time.Now().Add(-1 * time.Hour)),
Subject: "api_token",
},
}
token := jwt.NewWithClaims(jwt.SigningMethodHS256, claims)
tokenString, err := token.SignedString(auth.MakeJWTSigningKey("test"))
require.NoError(t, err)
validated, err := auth.ValidateAPIToken(tokenString, "test")
if tc.expectValid {
assert.NoError(t, err)
assert.NotNil(t, validated)
} else {
assert.Error(t, err)
assert.Contains(t, err.Error(), "expired")
}
})
}
}
// TestDualAuthentication verifies both cookie and token auth work together
func TestDualAuthentication(t *testing.T) {
db := setupTestDB(t)
defer db.Close()
tokenCache := auth.NewTokenCache(db)
userCache := auth.NewUserCache(db)
authMiddleware := auth.NewAuthMiddleware("/base/url", "test", tokenCache, userCache)
server := newTestServer(t, "/test", db, authMiddleware.AuthTokenRequired)
defer server.Close()
// Test 1: Cookie authentication
server.Authorize(t, []string{auth.PrivilegeAutomation})
assert.True(t, server.CallAndGetStatus(t), "Cookie auth should work")
// Test 2: Create and use API token
tokenID, err := auth.GenerateTokenID()
require.NoError(t, err)
apiToken := models.APIToken{
TokenID: tokenID,
UserID: 1,
RoleID: 2,
TTL: 3600,
Status: models.TokenStatusActive,
}
err = db.Create(&apiToken).Error
require.NoError(t, err)
claims := models.APITokenClaims{
TokenID: tokenID,
RID: 2,
UID: 1,
UHASH: "testhash",
RegisteredClaims: jwt.RegisteredClaims{
ExpiresAt: jwt.NewNumericDate(time.Now().Add(1 * time.Hour)),
IssuedAt: jwt.NewNumericDate(time.Now()),
Subject: "api_token",
},
}
token := jwt.NewWithClaims(jwt.SigningMethodHS256, claims)
tokenString, _ := token.SignedString(auth.MakeJWTSigningKey("test"))
// Unauthorize cookie
server.Unauthorize(t)
// Test 3: Token authentication should work
assert.True(t, server.CallAndGetStatus(t, "Bearer "+tokenString), "Token auth should work")
// Test 4: Both should work simultaneously
server.Authorize(t, []string{auth.PrivilegeAutomation})
assert.True(t, server.CallAndGetStatus(t, "Bearer "+tokenString), "Both auth methods should work")
}
// TestSecurityAudit_ClaimsInJWT verifies all security-critical data is in JWT
func TestSecurityAudit_ClaimsInJWT(t *testing.T) {
// Create token in DB with certain values
tokenID, err := auth.GenerateTokenID()
require.NoError(t, err)
dbToken := models.APIToken{
TokenID: tokenID,
UserID: 1,
RoleID: 2, // User role in DB
}
// Create JWT with different role (simulating compromise scenario)
jwtClaims := models.APITokenClaims{
TokenID: tokenID,
RID: 1, // Admin role in JWT (different from DB!)
UID: 1,
UHASH: "testhash",
RegisteredClaims: jwt.RegisteredClaims{
ExpiresAt: jwt.NewNumericDate(time.Now().Add(1 * time.Hour)),
IssuedAt: jwt.NewNumericDate(time.Now()),
Subject: "api_token",
},
}
token := jwt.NewWithClaims(jwt.SigningMethodHS256, jwtClaims)
tokenString, _ := token.SignedString(auth.MakeJWTSigningKey("test"))
// Validate token
validated, err := auth.ValidateAPIToken(tokenString, "test")
require.NoError(t, err)
// We trust JWT claims, not DB values
assert.Equal(t, uint64(1), validated.RID, "Should use role from JWT, not DB")
assert.NotEqual(t, dbToken.RoleID, validated.RID, "JWT role differs from DB role")
assert.Equal(t, dbToken.UserID, validated.UID)
assert.Equal(t, dbToken.TokenID, validated.TokenID)
// This is CORRECT behavior: DB only stores metadata for management
// Actual authorization data comes from signed JWT
}
// TestSecurityAudit_TokenIDUniqueness verifies token ID collision resistance
func TestSecurityAudit_TokenIDUniqueness(t *testing.T) {
iterations := 10000
tokens := make(map[string]bool, iterations)
for i := 0; i < iterations; i++ {
tokenID, err := auth.GenerateTokenID()
require.NoError(t, err)
// Check format
assert.Len(t, tokenID, 10)
// Check uniqueness
if tokens[tokenID] {
t.Fatalf("Duplicate token ID generated: %s", tokenID)
}
tokens[tokenID] = true
}
t.Logf("Generated %d unique token IDs without collision", iterations)
}
// TestSecurityAudit_SaltIsolation verifies JWT and Cookie keys are different
func TestSecurityAudit_SaltIsolation(t *testing.T) {
salts := []string{"salt1", "salt2", "production_salt"}
for _, salt := range salts {
t.Run("salt="+salt, func(t *testing.T) {
jwtKey := auth.MakeJWTSigningKey(salt)
cookieKeys := auth.MakeCookieStoreKey(salt)
// JWT key must be different from both cookie keys
assert.NotEqual(t, jwtKey, cookieKeys[0], "JWT key must differ from cookie auth key")
assert.NotEqual(t, jwtKey, cookieKeys[1], "JWT key must differ from cookie encryption key")
// Verify key lengths
assert.Len(t, jwtKey, 32, "JWT key must be 32 bytes")
assert.Len(t, cookieKeys[0], 64, "Cookie auth key must be 64 bytes")
assert.Len(t, cookieKeys[1], 32, "Cookie encryption key must be 32 bytes")
})
}
}
// TestAPIToken_ContextSetup verifies correct context values are set
func TestAPIToken_ContextSetup(t *testing.T) {
db := setupTestDB(t)
defer db.Close()
tokenCache := auth.NewTokenCache(db)
userCache := auth.NewUserCache(db)
authMiddleware := auth.NewAuthMiddleware("/base/url", "test", tokenCache, userCache)
// Create token
tokenID, err := auth.GenerateTokenID()
require.NoError(t, err)
apiToken := models.APIToken{
TokenID: tokenID,
UserID: 5,
RoleID: 3,
TTL: 3600,
Status: models.TokenStatusActive,
}
err = db.Create(&apiToken).Error
require.NoError(t, err)
user := models.User{
ID: 5,
Hash: "user5hash",
Mail: "user5@example.com",
Name: "User 5",
Status: models.UserStatusActive,
RoleID: 2,
}
err = db.Create(&user).Error
require.NoError(t, err)
claims := models.APITokenClaims{
TokenID: tokenID,
RID: 2,
UID: 5,
UHASH: "user5hash",
RegisteredClaims: jwt.RegisteredClaims{
ExpiresAt: jwt.NewNumericDate(time.Now().Add(1 * time.Hour)),
IssuedAt: jwt.NewNumericDate(time.Now()),
Subject: "api_token",
},
}
token := jwt.NewWithClaims(jwt.SigningMethodHS256, claims)
tokenString, _ := token.SignedString(auth.MakeJWTSigningKey("test"))
server := newTestServer(t, "/test", db, authMiddleware.AuthTokenRequired)
defer server.Close()
server.SetSessionCheckFunc(func(t *testing.T, c *gin.Context) {
t.Helper()
// Verify all context values are set correctly
assert.Equal(t, uint64(5), c.GetUint64("uid"), "UID from JWT")
assert.Equal(t, uint64(2), c.GetUint64("rid"), "RID from JWT")
assert.Equal(t, "user5hash", c.GetString("uhash"), "UHASH from JWT")
assert.Equal(t, "automation", c.GetString("cpt"), "CPT from JWT")
assert.Equal(t, "api", c.GetString("tid"), "TID should be 'api' for API tokens")
prms := c.GetStringSlice("prm")
assert.Contains(t, prms, auth.PrivilegeAutomation, "Should have automation privilege")
// Verify session timing fields
gtm := c.GetInt64("gtm")
assert.Greater(t, gtm, int64(0), "GTM (generation time) should be set")
exp := c.GetInt64("exp")
assert.Greater(t, exp, gtm, "EXP (expiration time) should be greater than GTM")
// UUID might be empty if hash is invalid (which is expected in tests)
uuid := c.GetString("uuid")
assert.NotNil(t, uuid, "UUID should be set (even if empty)")
})
assert.True(t, server.CallAndGetStatus(t, "Bearer "+tokenString))
}
// TestUserHashValidation_CookieAuth tests uhash validation with cookie authentication
func TestUserHashValidation_CookieAuth(t *testing.T) {
db := setupTestDB(t)
defer db.Close()
tokenCache := auth.NewTokenCache(db)
userCache := auth.NewUserCache(db)
authMiddleware := auth.NewAuthMiddleware("/base/url", "test", tokenCache, userCache)
server := newTestServer(t, "/test", db, authMiddleware.AuthUserRequired)
defer server.Close()
// Create test user with ID=1 and hash="123" to match session
var count int
db.Model(&models.User{}).Where("id = ?", 1).Count(&count)
testUser := models.User{
ID: 1,
Hash: "123",
Mail: "test_user@example.com",
Name: "Test User",
Status: models.UserStatusActive,
RoleID: 2,
}
if count == 0 {
err := db.Create(&testUser).Error
require.NoError(t, err)
} else {
db.First(&testUser, 1)
}
t.Run("correct uhash succeeds", func(t *testing.T) {
server.Authorize(t, []string{"test.permission"})
assert.True(t, server.CallAndGetStatus(t))
})
t.Run("modified uhash in database fails", func(t *testing.T) {
// Update user hash in database
db.Model(&testUser).Where("id = ?", 1).Update("hash", "modified_hash")
userCache.Invalidate(1)
// Try to authenticate with old session (has hash="123")
assert.False(t, server.CallAndGetStatus(t))
})
t.Run("blocked user fails", func(t *testing.T) {
// Restore original hash
db.Model(&testUser).Where("id = ?", 1).Update("hash", "123")
// Block user
db.Model(&testUser).Where("id = ?", 1).Update("status", models.UserStatusBlocked)
userCache.Invalidate(1)
assert.False(t, server.CallAndGetStatus(t))
})
t.Run("deleted user fails", func(t *testing.T) {
// Undelete and unblock first
db.Model(&models.User{}).Unscoped().Where("id = ?", 1).Update("deleted_at", nil)
db.Model(&testUser).Where("id = ?", 1).Update("status", models.UserStatusActive)
// Delete user
db.Delete(&testUser, 1)
userCache.Invalidate(1)
assert.False(t, server.CallAndGetStatus(t))
})
}
// TestUserHashValidation_TokenAuth tests uhash validation with token authentication
func TestUserHashValidation_TokenAuth(t *testing.T) {
db := setupTestDB(t)
defer db.Close()
tokenCache := auth.NewTokenCache(db)
userCache := auth.NewUserCache(db)
authMiddleware := auth.NewAuthMiddleware("/base/url", "test_salt", tokenCache, userCache)
server := newTestServer(t, "/protected", db, authMiddleware.AuthTokenRequired)
defer server.Close()
// Create test user
testUser := models.User{
ID: 200,
Hash: "token_test_hash",
Mail: "token_user@example.com",
Name: "Token Test User",
Status: models.UserStatusActive,
RoleID: 2,
}
err := db.Create(&testUser).Error
require.NoError(t, err)
// Create API token
tokenID, err := auth.GenerateTokenID()
require.NoError(t, err)
apiToken := models.APIToken{
TokenID: tokenID,
UserID: 200,
RoleID: 2,
TTL: 3600,
Status: models.TokenStatusActive,
}
err = db.Create(&apiToken).Error
require.NoError(t, err)
// Create JWT token with correct hash
claims := models.APITokenClaims{
TokenID: tokenID,
RID: 2,
UID: 200,
UHASH: "token_test_hash",
RegisteredClaims: jwt.RegisteredClaims{
ExpiresAt: jwt.NewNumericDate(time.Now().Add(1 * time.Hour)),
IssuedAt: jwt.NewNumericDate(time.Now()),
Subject: "api_token",
},
}
token := jwt.NewWithClaims(jwt.SigningMethodHS256, claims)
tokenString, err := token.SignedString(auth.MakeJWTSigningKey("test_salt"))
require.NoError(t, err)
t.Run("correct uhash succeeds", func(t *testing.T) {
success := server.CallAndGetStatus(t, "Bearer "+tokenString)
assert.True(t, success)
})
t.Run("modified uhash in database fails", func(t *testing.T) {
// Update user hash in database
db.Model(&testUser).Update("hash", "different_hash")
userCache.Invalidate(200)
// Try to authenticate with token (has original hash)
success := server.CallAndGetStatus(t, "Bearer "+tokenString)
assert.False(t, success)
})
t.Run("blocked user fails", func(t *testing.T) {
// Restore original hash
db.Model(&testUser).Update("hash", "token_test_hash")
// Block user
db.Model(&testUser).Update("status", models.UserStatusBlocked)
userCache.Invalidate(200)
success := server.CallAndGetStatus(t, "Bearer "+tokenString)
assert.False(t, success)
})
t.Run("deleted user fails", func(t *testing.T) {
// Unblock and restore for clean state
db.Model(&models.User{}).Unscoped().Where("id = ?", 200).Update("deleted_at", nil)
db.Model(&testUser).Update("status", models.UserStatusActive)
// Delete user
db.Delete(&testUser)
userCache.Invalidate(200)
success := server.CallAndGetStatus(t, "Bearer "+tokenString)
assert.False(t, success)
})
}
// TestUserHashValidation_CrossInstallation simulates different installations
func TestUserHashValidation_CrossInstallation(t *testing.T) {
db := setupTestDB(t)
defer db.Close()
tokenCache := auth.NewTokenCache(db)
userCache := auth.NewUserCache(db)
authMiddleware := auth.NewAuthMiddleware("/base/url", "test_salt", tokenCache, userCache)
server := newTestServer(t, "/protected", db, authMiddleware.AuthTokenRequired)
defer server.Close()
// Simulate Installation A
userInstallationA := models.User{
ID: 300,
Hash: "installation_a_hash",
Mail: "cross@example.com",
Name: "Cross Installation User",
Status: models.UserStatusActive,
RoleID: 2,
}
err := db.Create(&userInstallationA).Error
require.NoError(t, err)
// Create API token for Installation A
tokenID, err := auth.GenerateTokenID()
require.NoError(t, err)
apiToken := models.APIToken{
TokenID: tokenID,
UserID: 300,
RoleID: 2,
TTL: 3600,
Status: models.TokenStatusActive,
}
err = db.Create(&apiToken).Error
require.NoError(t, err)
// Create JWT token with Installation A hash
claimsA := models.APITokenClaims{
TokenID: tokenID,
RID: 2,
UID: 300,
UHASH: "installation_a_hash",
RegisteredClaims: jwt.RegisteredClaims{
ExpiresAt: jwt.NewNumericDate(time.Now().Add(1 * time.Hour)),
IssuedAt: jwt.NewNumericDate(time.Now()),
Subject: "api_token",
},
}
tokenA := jwt.NewWithClaims(jwt.SigningMethodHS256, claimsA)
tokenStringA, err := tokenA.SignedString(auth.MakeJWTSigningKey("test_salt"))
require.NoError(t, err)
t.Run("token works on Installation A", func(t *testing.T) {
success := server.CallAndGetStatus(t, "Bearer "+tokenStringA)
assert.True(t, success)
})
t.Run("token from Installation A fails on Installation B", func(t *testing.T) {
// Simulate Installation B - user has different hash
db.Model(&userInstallationA).Update("hash", "installation_b_hash")
userCache.Invalidate(300)
// Try to use token from Installation A (has installation_a_hash)
success := server.CallAndGetStatus(t, "Bearer "+tokenStringA)
assert.False(t, success, "Token from Installation A should not work on Installation B")
})
t.Run("new token from Installation B works", func(t *testing.T) {
// Create new token for Installation B
tokenIDB, err := auth.GenerateTokenID()
require.NoError(t, err)
apiTokenB := models.APIToken{
TokenID: tokenIDB,
UserID: 300,
RoleID: 2,
TTL: 3600,
Status: models.TokenStatusActive,
}
err = db.Create(&apiTokenB).Error
require.NoError(t, err)
// Create JWT token with Installation B hash
claimsB := models.APITokenClaims{
TokenID: tokenIDB,
RID: 2,
UID: 300,
UHASH: "installation_b_hash", // correct hash for Installation B
RegisteredClaims: jwt.RegisteredClaims{
ExpiresAt: jwt.NewNumericDate(time.Now().Add(1 * time.Hour)),
IssuedAt: jwt.NewNumericDate(time.Now()),
Subject: "api_token",
},
}
tokenB := jwt.NewWithClaims(jwt.SigningMethodHS256, claimsB)
tokenStringB, err := tokenB.SignedString(auth.MakeJWTSigningKey("test_salt"))
require.NoError(t, err)
// Token from Installation B should work
success := server.CallAndGetStatus(t, "Bearer "+tokenStringB)
assert.True(t, success, "New token from Installation B should work")
})
}
+46
View File
@@ -0,0 +1,46 @@
package auth
import (
"fmt"
"slices"
"pentagi/pkg/server/response"
"github.com/gin-gonic/gin"
)
func getPrms(c *gin.Context) ([]string, error) {
prms := c.GetStringSlice("prm")
if len(prms) == 0 {
return nil, fmt.Errorf("privileges are not set")
}
return prms, nil
}
func PrivilegesRequired(privs ...string) gin.HandlerFunc {
return func(c *gin.Context) {
if c.IsAborted() {
return
}
prms, err := getPrms(c)
if err != nil {
response.Error(c, response.ErrPrivilegesRequired, err)
c.Abort()
return
}
for _, priv := range privs {
if !LookupPerm(prms, priv) {
response.Error(c, response.ErrPrivilegesRequired, fmt.Errorf("'%s' is not set", priv))
c.Abort()
return
}
}
c.Next()
}
}
func LookupPerm(prm []string, perm string) bool {
return slices.Contains(prm, perm)
}
@@ -0,0 +1,36 @@
package auth_test
import (
"pentagi/pkg/server/auth"
"testing"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/assert"
)
func TestPrivilegesRequired(t *testing.T) {
db := setupTestDB(t)
defer db.Close()
tokenCache := auth.NewTokenCache(db)
userCache := auth.NewUserCache(db)
authMiddleware := auth.NewAuthMiddleware("/base/url", "test", tokenCache, userCache)
server := newTestServer(t, "/test", db, authMiddleware.AuthTokenRequired, auth.PrivilegesRequired("priv1", "priv2"))
defer server.Close()
server.SetSessionCheckFunc(func(t *testing.T, c *gin.Context) {
t.Helper()
assert.Equal(t, uint64(1), c.GetUint64("uid"))
})
assert.False(t, server.CallAndGetStatus(t))
server.Authorize(t, []string{"some.permission"})
assert.False(t, server.CallAndGetStatus(t))
server.Authorize(t, []string{"priv1"})
assert.False(t, server.CallAndGetStatus(t))
server.Authorize(t, []string{"priv1", "priv2"})
assert.True(t, server.CallAndGetStatus(t))
}
+71
View File
@@ -0,0 +1,71 @@
package auth
import (
"crypto/sha512"
"strings"
"sync"
"golang.org/x/crypto/pbkdf2"
)
var (
cookieStoreKeys sync.Map // cache of cookie keys per salt
jwtSigningKeys sync.Map // cache of JWT signing keys per salt
)
const (
pbkdf2Iterations = 210000 // OWASP 2023 recommendation
jwtKeyLength = 32 // 256 bits for HS256
authKeyLength = 64 // 512 bits for cookie auth key
encKeyLength = 32 // 256 bits for cookie encryption key
)
// MakeCookieStoreKey is function to generate auth and encryption keys for cookie store
func MakeCookieStoreKey(globalSalt string) [][]byte {
// Check cache for existing keys
if cached, ok := cookieStoreKeys.Load(globalSalt); ok {
return cached.([][]byte)
}
// Generate new keys for this salt using PBKDF2
password := []byte(strings.Join([]string{
"a8d0abae36f749588f4393e6fc292690",
globalSalt,
"7c9be62adec5076970fa946e78f256e2",
}, "|"))
// Auth key (64 bytes) - using salt variant 1
authSalt := []byte("pentagi.cookie.auth|" + globalSalt)
authKey := pbkdf2.Key(password, authSalt, pbkdf2Iterations, authKeyLength, sha512.New)
// Encryption key (32 bytes) - using salt variant 2
encSalt := []byte("pentagi.cookie.enc|" + globalSalt)
encKey := pbkdf2.Key(password, encSalt, pbkdf2Iterations, encKeyLength, sha512.New)
newKeys := [][]byte{authKey, encKey}
// Store in cache (LoadOrStore handles concurrent access)
actual, _ := cookieStoreKeys.LoadOrStore(globalSalt, newKeys)
return actual.([][]byte)
}
// MakeJWTSigningKey is function to generate signing key for JWT tokens
func MakeJWTSigningKey(globalSalt string) []byte {
// Check cache for existing key
if cached, ok := jwtSigningKeys.Load(globalSalt); ok {
return cached.([]byte)
}
// Generate new key for this salt using PBKDF2
password := []byte(strings.Join([]string{
"4c1e9cb77df7f9a58fcc5f52d40af685",
globalSalt,
"09784e190148d13d48885aa47cf8a297",
}, "|"))
salt := []byte("pentagi.jwt.signing|" + globalSalt)
newKey := pbkdf2.Key(password, salt, pbkdf2Iterations, jwtKeyLength, sha512.New)
// Store in cache (LoadOrStore handles concurrent access)
actual, _ := jwtSigningKeys.LoadOrStore(globalSalt, newKey)
return actual.([]byte)
}
+61
View File
@@ -0,0 +1,61 @@
package auth_test
import (
"pentagi/pkg/server/auth"
"testing"
"github.com/stretchr/testify/assert"
)
func TestMakeJWTSigningKey(t *testing.T) {
salt1 := "test_salt_1"
salt2 := "test_salt_2"
// Test that key is generated
key1 := auth.MakeJWTSigningKey(salt1)
assert.NotNil(t, key1)
assert.Len(t, key1, 32, "JWT signing key should be 32 bytes (256 bits)")
// Test that same salt produces same key (cached)
key1Again := auth.MakeJWTSigningKey(salt1)
assert.Equal(t, key1, key1Again, "Same salt should produce same key from cache")
// Test that different salts produce different keys
key2 := auth.MakeJWTSigningKey(salt2)
assert.NotEqual(t, key1, key2, "Different salts should produce different keys")
assert.Len(t, key2, 32, "JWT signing key should be 32 bytes (256 bits)")
// Verify consistency for salt2
key2Again := auth.MakeJWTSigningKey(salt2)
assert.Equal(t, key2, key2Again, "Same salt should produce same key from cache")
}
func TestMakeCookieStoreKey(t *testing.T) {
salt := "test_salt"
// Test that keys are generated
keys := auth.MakeCookieStoreKey(salt)
assert.NotNil(t, keys)
assert.Len(t, keys, 2, "Should return auth and encryption keys")
// Test that auth key is 64 bytes (SHA512)
assert.Len(t, keys[0], 64, "Auth key should be 64 bytes")
// Test that encryption key is 32 bytes (SHA256)
assert.Len(t, keys[1], 32, "Encryption key should be 32 bytes")
// Test consistency
keysAgain := auth.MakeCookieStoreKey(salt)
assert.Equal(t, keys, keysAgain, "Same salt should produce same keys")
}
func TestMakeJWTSigningKeyDifferentFromCookieKey(t *testing.T) {
salt := "test_salt"
jwtKey := auth.MakeJWTSigningKey(salt)
cookieKeys := auth.MakeCookieStoreKey(salt)
// JWT signing key should be different from both cookie keys
assert.NotEqual(t, jwtKey, cookieKeys[0], "JWT key should differ from cookie auth key")
assert.NotEqual(t, jwtKey, cookieKeys[1], "JWT key should differ from cookie encryption key")
}
+92
View File
@@ -0,0 +1,92 @@
package auth
import (
"sync"
"time"
"pentagi/pkg/server/models"
"github.com/jinzhu/gorm"
)
// userCacheEntry represents a cached user status entry
type userCacheEntry struct {
hash string
status models.UserStatus
notFound bool // negative caching
expiresAt time.Time
}
// UserCache provides caching for user hash lookups
type UserCache struct {
cache sync.Map
ttl time.Duration
db *gorm.DB
}
// NewUserCache creates a new user cache instance
func NewUserCache(db *gorm.DB) *UserCache {
return &UserCache{
ttl: 5 * time.Minute,
db: db,
}
}
// SetTTL sets the TTL for the user cache
func (uc *UserCache) SetTTL(ttl time.Duration) {
uc.ttl = ttl
}
// GetUserHash retrieves user hash and status from cache or database
func (uc *UserCache) GetUserHash(userID uint64) (string, models.UserStatus, error) {
// check cache first
if entry, ok := uc.cache.Load(userID); ok {
cached := entry.(userCacheEntry)
if time.Now().Before(cached.expiresAt) {
// return cached "not found" error
if cached.notFound {
return "", "", gorm.ErrRecordNotFound
}
return cached.hash, cached.status, nil
}
// cache entry expired, remove it
uc.cache.Delete(userID)
}
// load from database
var user models.User
if err := uc.db.Where("id = ?", userID).First(&user).Error; err != nil {
if gorm.IsRecordNotFoundError(err) {
// cache negative result (user not found)
uc.cache.Store(userID, userCacheEntry{
notFound: true,
expiresAt: time.Now().Add(uc.ttl),
})
return "", "", gorm.ErrRecordNotFound
}
return "", "", err
}
// update cache with positive result
uc.cache.Store(userID, userCacheEntry{
hash: user.Hash,
status: user.Status,
notFound: false,
expiresAt: time.Now().Add(uc.ttl),
})
return user.Hash, user.Status, nil
}
// Invalidate removes a specific user from cache
func (uc *UserCache) Invalidate(userID uint64) {
uc.cache.Delete(userID)
}
// InvalidateAll clears the entire cache
func (uc *UserCache) InvalidateAll() {
uc.cache.Range(func(key, value any) bool {
uc.cache.Delete(key)
return true
})
}
+443
View File
@@ -0,0 +1,443 @@
package auth_test
import (
"fmt"
"sync"
"testing"
"time"
"pentagi/pkg/server/auth"
"pentagi/pkg/server/models"
"github.com/jinzhu/gorm"
_ "github.com/jinzhu/gorm/dialects/sqlite"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
func setupUserTestDB(t *testing.T) *gorm.DB {
t.Helper()
db, err := gorm.Open("sqlite3", ":memory:")
require.NoError(t, err)
// Create users table
db.Exec(`
CREATE TABLE users (
id INTEGER PRIMARY KEY AUTOINCREMENT,
hash TEXT NOT NULL UNIQUE,
type TEXT NOT NULL DEFAULT 'local',
mail TEXT NOT NULL UNIQUE,
name TEXT NOT NULL DEFAULT '',
status TEXT NOT NULL DEFAULT 'active',
role_id INTEGER NOT NULL DEFAULT 2,
password TEXT,
password_change_required BOOLEAN NOT NULL DEFAULT false,
provider TEXT,
created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
deleted_at DATETIME
)
`)
// Create user_preferences table
db.Exec(`
CREATE TABLE user_preferences (
id INTEGER PRIMARY KEY AUTOINCREMENT,
user_id INTEGER NOT NULL UNIQUE,
preferences TEXT NOT NULL DEFAULT '{"favoriteFlows": []}',
created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
updated_at DATETIME DEFAULT CURRENT_TIMESTAMP,
FOREIGN KEY (user_id) REFERENCES users(id) ON DELETE CASCADE
)
`)
time.Sleep(200 * time.Millisecond) // wait for database to be ready
return db
}
func TestUserCache_GetUserHash(t *testing.T) {
db := setupUserTestDB(t)
defer db.Close()
cache := auth.NewUserCache(db)
// Insert test user
user := models.User{
ID: 1,
Hash: "test_hash_123",
Mail: "test@example.com",
Name: "Test User",
Status: models.UserStatusActive,
RoleID: 2,
}
err := db.Create(&user).Error
require.NoError(t, err)
// Test: Get user hash (should hit database)
hash, status, err := cache.GetUserHash(1)
require.NoError(t, err)
assert.Equal(t, "test_hash_123", hash)
assert.Equal(t, models.UserStatusActive, status)
// Test: Get user hash again (should hit cache)
hash, status, err = cache.GetUserHash(1)
require.NoError(t, err)
assert.Equal(t, "test_hash_123", hash)
assert.Equal(t, models.UserStatusActive, status)
// Test: Non-existent user
_, _, err = cache.GetUserHash(999)
assert.Error(t, err)
assert.Equal(t, gorm.ErrRecordNotFound, err)
}
func TestUserCache_Invalidate(t *testing.T) {
db := setupUserTestDB(t)
defer db.Close()
cache := auth.NewUserCache(db)
// Insert test user
user := models.User{
ID: 1,
Hash: "test_hash_456",
Mail: "test2@example.com",
Name: "Test User 2",
Status: models.UserStatusActive,
RoleID: 2,
}
err := db.Create(&user).Error
require.NoError(t, err)
// Get hash to populate cache
hash, status, err := cache.GetUserHash(1)
require.NoError(t, err)
assert.Equal(t, "test_hash_456", hash)
assert.Equal(t, models.UserStatusActive, status)
// Update user in database
db.Model(&user).Update("status", models.UserStatusBlocked)
// Status should still be active (from cache)
hash, status, err = cache.GetUserHash(1)
require.NoError(t, err)
assert.Equal(t, "test_hash_456", hash)
assert.Equal(t, models.UserStatusActive, status)
// Invalidate cache
cache.Invalidate(1)
// Status should now be blocked (from database)
hash, status, err = cache.GetUserHash(1)
require.NoError(t, err)
assert.Equal(t, "test_hash_456", hash)
assert.Equal(t, models.UserStatusBlocked, status)
}
func TestUserCache_Expiration(t *testing.T) {
db := setupUserTestDB(t)
defer db.Close()
// Create cache with very short TTL for testing
cache := auth.NewUserCache(db)
cache.SetTTL(300 * time.Millisecond)
// Insert test user
user := models.User{
ID: 1,
Hash: "test_hash_789",
Mail: "test3@example.com",
Name: "Test User 3",
Status: models.UserStatusActive,
RoleID: 2,
}
err := db.Create(&user).Error
require.NoError(t, err)
// Get hash to populate cache
hash, status, err := cache.GetUserHash(1)
require.NoError(t, err)
assert.Equal(t, "test_hash_789", hash)
assert.Equal(t, models.UserStatusActive, status)
// Update user in database
db.Model(&user).Update("status", models.UserStatusBlocked)
// Wait for cache to expire
time.Sleep(500 * time.Millisecond)
// Status should now be blocked (cache expired, reading from DB)
hash, status, err = cache.GetUserHash(1)
require.NoError(t, err)
assert.Equal(t, "test_hash_789", hash)
assert.Equal(t, models.UserStatusBlocked, status)
}
func TestUserCache_UserStatuses(t *testing.T) {
db := setupUserTestDB(t)
defer db.Close()
cache := auth.NewUserCache(db)
testCases := []struct {
name string
userStatus models.UserStatus
expectedStatus models.UserStatus
}{
{
name: "active user",
userStatus: models.UserStatusActive,
expectedStatus: models.UserStatusActive,
},
{
name: "blocked user",
userStatus: models.UserStatusBlocked,
expectedStatus: models.UserStatusBlocked,
},
{
name: "created user",
userStatus: models.UserStatusCreated,
expectedStatus: models.UserStatusCreated,
},
}
for i, tc := range testCases {
t.Run(tc.name, func(t *testing.T) {
user := models.User{
ID: uint64(i + 1),
Hash: "hash_" + tc.name,
Mail: tc.name + "@example.com",
Name: tc.name,
Status: tc.userStatus,
RoleID: 2,
}
err := db.Create(&user).Error
require.NoError(t, err)
hash, status, err := cache.GetUserHash(user.ID)
require.NoError(t, err)
assert.Equal(t, user.Hash, hash)
assert.Equal(t, tc.expectedStatus, status)
})
}
}
func TestUserCache_DeletedUser(t *testing.T) {
db := setupUserTestDB(t)
defer db.Close()
cache := auth.NewUserCache(db)
// Insert test user
user := models.User{
ID: 1,
Hash: "deleted_hash",
Mail: "deleted@example.com",
Name: "Deleted User",
Status: models.UserStatusActive,
RoleID: 2,
}
err := db.Create(&user).Error
require.NoError(t, err)
// Get hash to populate cache
hash, status, err := cache.GetUserHash(1)
require.NoError(t, err)
assert.Equal(t, "deleted_hash", hash)
assert.Equal(t, models.UserStatusActive, status)
// Soft delete user
db.Delete(&user)
// Invalidate cache
cache.Invalidate(1)
// Should return error for deleted user
_, _, err = cache.GetUserHash(1)
assert.Error(t, err)
assert.Equal(t, gorm.ErrRecordNotFound, err)
}
func TestUserCache_ConcurrentAccess(t *testing.T) {
db := setupUserTestDB(t)
defer db.Close()
cache := auth.NewUserCache(db)
// Insert test users
for i := 1; i <= 10; i++ {
user := models.User{
ID: uint64(i),
Hash: fmt.Sprintf("concurrent_hash_%d", i),
Mail: fmt.Sprintf("concurrent%d@example.com", i),
Name: "Concurrent User",
Status: models.UserStatusActive,
RoleID: 2,
}
err := db.Create(&user).Error
require.NoError(t, err)
}
// warm up cache
for i := range 10 {
_, _, err := cache.GetUserHash(uint64(i%10 + 1))
require.NoError(t, err)
}
var wg sync.WaitGroup
errors := make(chan error, 100)
// Concurrent reads
for i := range 10 {
wg.Add(1)
go func(userID uint64) {
defer wg.Done()
for range 10 {
_, _, err := cache.GetUserHash(userID)
if err != nil {
errors <- err
}
}
}(uint64(i%10 + 1))
}
// Concurrent invalidations
for i := range 5 {
wg.Add(1)
go func(userID uint64) {
defer wg.Done()
for range 5 {
cache.Invalidate(userID)
time.Sleep(10 * time.Millisecond)
}
}(uint64(i%10 + 1))
}
wg.Wait()
close(errors)
// Check for errors
for err := range errors {
t.Errorf("Concurrent access error: %v", err)
}
}
func TestUserCache_InvalidateAll(t *testing.T) {
db := setupUserTestDB(t)
defer db.Close()
cache := auth.NewUserCache(db)
// Insert multiple users
for i := 1; i <= 5; i++ {
user := models.User{
ID: uint64(i),
Hash: fmt.Sprintf("invalidate_all_%d", i),
Mail: fmt.Sprintf("all%d@example.com", i),
Name: fmt.Sprintf("User %d", i),
Status: models.UserStatusActive,
RoleID: 2,
}
err := db.Create(&user).Error
require.NoError(t, err)
}
// Populate cache
for i := 1; i <= 5; i++ {
_, _, err := cache.GetUserHash(uint64(i))
require.NoError(t, err)
}
// Update all users in database
db.Model(&models.User{}).Where("id > 0").Update("status", models.UserStatusBlocked)
// Invalidate all
cache.InvalidateAll()
// All users should now show blocked status
for i := 1; i <= 5; i++ {
_, status, err := cache.GetUserHash(uint64(i))
require.NoError(t, err)
assert.Equal(t, models.UserStatusBlocked, status)
}
}
func TestUserCache_NegativeCaching(t *testing.T) {
db := setupUserTestDB(t)
defer db.Close()
cache := auth.NewUserCache(db)
nonExistentUserID := uint64(9999)
// First call - should hit database and cache the "not found"
_, _, err := cache.GetUserHash(nonExistentUserID)
require.Error(t, err)
assert.Equal(t, gorm.ErrRecordNotFound, err)
// Second call - should return from cache without hitting DB
_, _, err = cache.GetUserHash(nonExistentUserID)
require.Error(t, err)
assert.Equal(t, gorm.ErrRecordNotFound, err, "Should return cached not found error")
// Now create the user in DB
user := models.User{
ID: nonExistentUserID,
Hash: "new_user_hash",
Mail: "new@example.com",
Name: "New User",
Status: models.UserStatusActive,
RoleID: 2,
}
err = db.Create(&user).Error
require.NoError(t, err)
// Should still return cached "not found" until invalidated
_, _, err = cache.GetUserHash(nonExistentUserID)
require.Error(t, err)
assert.Equal(t, gorm.ErrRecordNotFound, err, "Should still return cached not found")
// Invalidate cache
cache.Invalidate(nonExistentUserID)
// Now should find the user
hash, status, err := cache.GetUserHash(nonExistentUserID)
require.NoError(t, err)
assert.Equal(t, "new_user_hash", hash)
assert.Equal(t, models.UserStatusActive, status)
}
func TestUserCache_NegativeCachingExpiration(t *testing.T) {
db := setupUserTestDB(t)
defer db.Close()
cache := auth.NewUserCache(db)
cache.SetTTL(300 * time.Millisecond)
nonExistentUserID := uint64(8888)
// First call - cache the "not found"
_, _, err := cache.GetUserHash(nonExistentUserID)
require.Error(t, err)
assert.Equal(t, gorm.ErrRecordNotFound, err)
// Create user in DB
user := models.User{
ID: nonExistentUserID,
Hash: "temp_user_hash",
Mail: "temp@example.com",
Name: "Temp User",
Status: models.UserStatusActive,
RoleID: 2,
}
err = db.Create(&user).Error
require.NoError(t, err)
// Wait for cache to expire
time.Sleep(500 * time.Millisecond)
// Now should find the user (cache expired)
hash, status, err := cache.GetUserHash(nonExistentUserID)
require.NoError(t, err)
assert.Equal(t, "temp_user_hash", hash)
assert.Equal(t, models.UserStatusActive, status)
}
+61
View File
@@ -0,0 +1,61 @@
package context
import (
"github.com/gin-contrib/sessions"
"github.com/gin-gonic/gin"
)
// GetInt64 is function to get some int64 value from gin context
func GetInt64(c *gin.Context, key string) (int64, bool) {
if iv, ok := c.Get(key); !ok {
return 0, false
} else if v, ok := iv.(int64); !ok {
return 0, false
} else {
return v, true
}
}
// GetUint64 is function to get some uint64 value from gin context
func GetUint64(c *gin.Context, key string) (uint64, bool) {
if iv, ok := c.Get(key); !ok {
return 0, false
} else if v, ok := iv.(uint64); !ok {
return 0, false
} else {
return v, true
}
}
// GetString is function to get some string value from gin context
func GetString(c *gin.Context, key string) (string, bool) {
if iv, ok := c.Get(key); !ok {
return "", false
} else if v, ok := iv.(string); !ok {
return "", false
} else {
return v, true
}
}
// GetStringArray is function to get some string array value from gin context
func GetStringArray(c *gin.Context, key string) ([]string, bool) {
if iv, ok := c.Get(key); !ok {
return []string{}, false
} else if v, ok := iv.([]string); !ok {
return []string{}, false
} else {
return v, true
}
}
func GetStringFromSession(c *gin.Context, key string) (string, bool) {
session := sessions.Default(c)
if iv := session.Get(key); iv == nil {
return "", false
} else if v, ok := iv.(string); !ok {
return "", false
} else {
return v, true
}
}
+510
View File
@@ -0,0 +1,510 @@
package context
import (
"net/http"
"net/http/httptest"
"testing"
"github.com/gin-contrib/sessions"
"github.com/gin-contrib/sessions/cookie"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/assert"
)
func init() {
gin.SetMode(gin.TestMode)
}
func TestGetInt64(t *testing.T) {
t.Parallel()
tests := []struct {
name string
setup func(c *gin.Context)
key string
wantVal int64
wantOK bool
}{
{
name: "found",
setup: func(c *gin.Context) { c.Set("id", int64(42)) },
key: "id",
wantVal: 42,
wantOK: true,
},
{
name: "missing",
setup: func(c *gin.Context) {},
key: "id",
wantVal: 0,
wantOK: false,
},
{
name: "wrong type string",
setup: func(c *gin.Context) { c.Set("id", "not-an-int") },
key: "id",
wantVal: 0,
wantOK: false,
},
{
name: "wrong type uint64",
setup: func(c *gin.Context) { c.Set("id", uint64(99)) },
key: "id",
wantVal: 0,
wantOK: false,
},
{
name: "zero value",
setup: func(c *gin.Context) { c.Set("id", int64(0)) },
key: "id",
wantVal: 0,
wantOK: true,
},
{
name: "negative value",
setup: func(c *gin.Context) { c.Set("id", int64(-100)) },
key: "id",
wantVal: -100,
wantOK: true,
},
{
name: "max int64",
setup: func(c *gin.Context) { c.Set("id", int64(9223372036854775807)) },
key: "id",
wantVal: 9223372036854775807,
wantOK: true,
},
{
name: "different key",
setup: func(c *gin.Context) { c.Set("other", int64(123)) },
key: "id",
wantVal: 0,
wantOK: false,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
t.Parallel()
c, _ := gin.CreateTestContext(httptest.NewRecorder())
tt.setup(c)
val, ok := GetInt64(c, tt.key)
assert.Equal(t, tt.wantOK, ok)
assert.Equal(t, tt.wantVal, val)
})
}
}
func TestGetUint64(t *testing.T) {
t.Parallel()
tests := []struct {
name string
setup func(c *gin.Context)
key string
wantVal uint64
wantOK bool
}{
{
name: "found",
setup: func(c *gin.Context) { c.Set("uid", uint64(99)) },
key: "uid",
wantVal: 99,
wantOK: true,
},
{
name: "missing",
setup: func(c *gin.Context) {},
key: "uid",
wantVal: 0,
wantOK: false,
},
{
name: "wrong type int64",
setup: func(c *gin.Context) { c.Set("uid", int64(99)) },
key: "uid",
wantVal: 0,
wantOK: false,
},
{
name: "wrong type string",
setup: func(c *gin.Context) { c.Set("uid", "99") },
key: "uid",
wantVal: 0,
wantOK: false,
},
{
name: "zero value",
setup: func(c *gin.Context) { c.Set("uid", uint64(0)) },
key: "uid",
wantVal: 0,
wantOK: true,
},
{
name: "large value",
setup: func(c *gin.Context) { c.Set("uid", uint64(18446744073709551615)) },
key: "uid",
wantVal: 18446744073709551615,
wantOK: true,
},
{
name: "different key",
setup: func(c *gin.Context) { c.Set("other", uint64(456)) },
key: "uid",
wantVal: 0,
wantOK: false,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
t.Parallel()
c, _ := gin.CreateTestContext(httptest.NewRecorder())
tt.setup(c)
val, ok := GetUint64(c, tt.key)
assert.Equal(t, tt.wantOK, ok)
assert.Equal(t, tt.wantVal, val)
})
}
}
func TestGetString(t *testing.T) {
t.Parallel()
tests := []struct {
name string
setup func(c *gin.Context)
key string
wantVal string
wantOK bool
}{
{
name: "found",
setup: func(c *gin.Context) { c.Set("name", "alice") },
key: "name",
wantVal: "alice",
wantOK: true,
},
{
name: "missing",
setup: func(c *gin.Context) {},
key: "name",
wantVal: "",
wantOK: false,
},
{
name: "wrong type int",
setup: func(c *gin.Context) { c.Set("name", 123) },
key: "name",
wantVal: "",
wantOK: false,
},
{
name: "wrong type bool",
setup: func(c *gin.Context) { c.Set("name", true) },
key: "name",
wantVal: "",
wantOK: false,
},
{
name: "empty string",
setup: func(c *gin.Context) { c.Set("name", "") },
key: "name",
wantVal: "",
wantOK: true,
},
{
name: "long string",
setup: func(c *gin.Context) { c.Set("name", "very-long-string-with-special-chars-@#$%") },
key: "name",
wantVal: "very-long-string-with-special-chars-@#$%",
wantOK: true,
},
{
name: "different key",
setup: func(c *gin.Context) { c.Set("other", "value") },
key: "name",
wantVal: "",
wantOK: false,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
t.Parallel()
c, _ := gin.CreateTestContext(httptest.NewRecorder())
tt.setup(c)
val, ok := GetString(c, tt.key)
assert.Equal(t, tt.wantOK, ok)
assert.Equal(t, tt.wantVal, val)
})
}
}
func TestGetStringArray(t *testing.T) {
t.Parallel()
tests := []struct {
name string
setup func(c *gin.Context)
key string
wantVal []string
wantOK bool
}{
{
name: "found",
setup: func(c *gin.Context) { c.Set("perms", []string{"read", "write"}) },
key: "perms",
wantVal: []string{"read", "write"},
wantOK: true,
},
{
name: "missing",
setup: func(c *gin.Context) {},
key: "perms",
wantVal: []string{},
wantOK: false,
},
{
name: "wrong type string",
setup: func(c *gin.Context) { c.Set("perms", "not-a-slice") },
key: "perms",
wantVal: []string{},
wantOK: false,
},
{
name: "wrong type int slice",
setup: func(c *gin.Context) { c.Set("perms", []int{1, 2, 3}) },
key: "perms",
wantVal: []string{},
wantOK: false,
},
{
name: "empty array",
setup: func(c *gin.Context) { c.Set("perms", []string{}) },
key: "perms",
wantVal: []string{},
wantOK: true,
},
{
name: "nil array",
setup: func(c *gin.Context) { c.Set("perms", []string(nil)) },
key: "perms",
wantVal: nil,
wantOK: true,
},
{
name: "single element",
setup: func(c *gin.Context) { c.Set("perms", []string{"admin"}) },
key: "perms",
wantVal: []string{"admin"},
wantOK: true,
},
{
name: "many elements",
setup: func(c *gin.Context) { c.Set("perms", []string{"a", "b", "c", "d", "e"}) },
key: "perms",
wantVal: []string{"a", "b", "c", "d", "e"},
wantOK: true,
},
{
name: "different key",
setup: func(c *gin.Context) { c.Set("other", []string{"val"}) },
key: "perms",
wantVal: []string{},
wantOK: false,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
t.Parallel()
c, _ := gin.CreateTestContext(httptest.NewRecorder())
tt.setup(c)
val, ok := GetStringArray(c, tt.key)
assert.Equal(t, tt.wantOK, ok)
assert.Equal(t, tt.wantVal, val)
})
}
}
func TestGetStringFromSession(t *testing.T) {
t.Parallel()
tests := []struct {
name string
setup func(session sessions.Session)
key string
wantVal string
wantOK bool
}{
{
name: "found",
setup: func(s sessions.Session) {
s.Set("token", "abc123")
_ = s.Save()
},
key: "token",
wantVal: "abc123",
wantOK: true,
},
{
name: "missing",
setup: func(s sessions.Session) {},
key: "token",
wantVal: "",
wantOK: false,
},
{
name: "wrong type int",
setup: func(s sessions.Session) {
s.Set("token", 999)
_ = s.Save()
},
key: "token",
wantVal: "",
wantOK: false,
},
{
name: "wrong type bool",
setup: func(s sessions.Session) {
s.Set("token", true)
_ = s.Save()
},
key: "token",
wantVal: "",
wantOK: false,
},
{
name: "empty string",
setup: func(s sessions.Session) {
s.Set("token", "")
_ = s.Save()
},
key: "token",
wantVal: "",
wantOK: true,
},
{
name: "different key",
setup: func(s sessions.Session) {
s.Set("other", "value")
_ = s.Save()
},
key: "token",
wantVal: "",
wantOK: false,
},
{
name: "multiple values in session",
setup: func(s sessions.Session) {
s.Set("token", "abc123")
s.Set("user", "alice")
s.Set("role", "admin")
_ = s.Save()
},
key: "token",
wantVal: "abc123",
wantOK: true,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
t.Parallel()
store := cookie.NewStore([]byte("test-secret"))
router := gin.New()
router.Use(sessions.Sessions("test", store))
var val string
var ok bool
router.GET("/test", func(c *gin.Context) {
session := sessions.Default(c)
tt.setup(session)
val, ok = GetStringFromSession(c, tt.key)
})
w := httptest.NewRecorder()
req := httptest.NewRequest(http.MethodGet, "/test", nil)
router.ServeHTTP(w, req)
assert.Equal(t, tt.wantOK, ok)
assert.Equal(t, tt.wantVal, val)
})
}
}
func TestMultipleValuesInContext(t *testing.T) {
t.Parallel()
c, _ := gin.CreateTestContext(httptest.NewRecorder())
c.Set("int64_val", int64(123))
c.Set("uint64_val", uint64(456))
c.Set("string_val", "test")
c.Set("array_val", []string{"a", "b"})
// Verify all values are independently accessible
intVal, ok := GetInt64(c, "int64_val")
assert.True(t, ok)
assert.Equal(t, int64(123), intVal)
uintVal, ok := GetUint64(c, "uint64_val")
assert.True(t, ok)
assert.Equal(t, uint64(456), uintVal)
strVal, ok := GetString(c, "string_val")
assert.True(t, ok)
assert.Equal(t, "test", strVal)
arrVal, ok := GetStringArray(c, "array_val")
assert.True(t, ok)
assert.Equal(t, []string{"a", "b"}, arrVal)
}
func TestContextOverwrite(t *testing.T) {
t.Parallel()
c, _ := gin.CreateTestContext(httptest.NewRecorder())
// Set initial value
c.Set("key", "original")
val, ok := GetString(c, "key")
assert.True(t, ok)
assert.Equal(t, "original", val)
// Overwrite with new value
c.Set("key", "updated")
val, ok = GetString(c, "key")
assert.True(t, ok)
assert.Equal(t, "updated", val)
}
func TestContextTypeChange(t *testing.T) {
t.Parallel()
c, _ := gin.CreateTestContext(httptest.NewRecorder())
// Set as string
c.Set("value", "123")
strVal, ok := GetString(c, "value")
assert.True(t, ok)
assert.Equal(t, "123", strVal)
// Try to get as int64 - should fail
intVal, ok := GetInt64(c, "value")
assert.False(t, ok)
assert.Equal(t, int64(0), intVal)
// Overwrite with int64
c.Set("value", int64(123))
intVal, ok = GetInt64(c, "value")
assert.True(t, ok)
assert.Equal(t, int64(123), intVal)
}
File diff suppressed because it is too large Load Diff
File diff suppressed because it is too large Load Diff
File diff suppressed because it is too large Load Diff
+99
View File
@@ -0,0 +1,99 @@
package logger
import (
"context"
"time"
obs "pentagi/pkg/observability"
"github.com/99designs/gqlgen/graphql"
"github.com/gin-gonic/gin"
"github.com/sirupsen/logrus"
)
// FromContext is function to get logrus Entry with context
func FromContext(c *gin.Context) *logrus.Entry {
return logrus.WithContext(c.Request.Context())
}
func TraceEnabled() bool {
return logrus.IsLevelEnabled(logrus.TraceLevel)
}
func WithGinLogger(service string) gin.HandlerFunc {
return func(c *gin.Context) {
start := time.Now()
uri := c.Request.URL.Path
raw := c.Request.URL.RawQuery
if raw != "" {
uri = uri + "?" + raw
}
entry := logrus.WithFields(logrus.Fields{
"component": "api",
"net_peer_ip": c.ClientIP(),
"http_uri": uri,
"http_path": c.Request.URL.Path,
"http_host_name": c.Request.Host,
"http_method": c.Request.Method,
})
if c.FullPath() == "" {
entry = entry.WithField("request", "proxy handled")
} else {
entry = entry.WithField("request", "api handled")
}
// serve the request to the next middleware
c.Next()
if len(c.Errors) > 0 {
entry = entry.WithField("gin.errors", c.Errors.String())
}
entry = entry.WithFields(logrus.Fields{
"duration": time.Since(start).String(),
"http_status_code": c.Writer.Status(),
"http_resp_size": c.Writer.Size(),
}).WithContext(c.Request.Context())
if c.Writer.Status() >= 400 {
entry.Error("http request handled error")
} else {
entry.Debug("http request handled success")
}
}
}
func WithGqlLogger(service string) func(ctx context.Context, next graphql.ResponseHandler) *graphql.Response {
return func(ctx context.Context, next graphql.ResponseHandler) *graphql.Response {
ctx, span := obs.Observer.NewSpan(ctx, obs.SpanKindServer, "graphql.handler")
defer span.End()
start := time.Now()
entry := logrus.WithContext(ctx).WithField("component", service)
res := next(ctx)
op := graphql.GetOperationContext(ctx)
if op != nil && op.Operation != nil {
entry = entry.WithFields(logrus.Fields{
"operation_name": op.OperationName,
"operation_type": op.Operation.Operation,
})
}
entry = entry.WithField("duration", time.Since(start).String())
if res == nil {
return res
}
if len(res.Errors) > 0 {
entry = entry.WithField("gql.errors", res.Errors.Error())
entry.Error("graphql request handled with errors")
} else {
entry.Debug("graphql request handled success")
}
return res
}
}
+36
View File
@@ -0,0 +1,36 @@
package router
import (
"pentagi/pkg/server/models"
"pentagi/pkg/server/response"
"github.com/gin-contrib/sessions"
"github.com/gin-gonic/gin"
)
func localUserRequired() gin.HandlerFunc {
return func(c *gin.Context) {
if c.IsAborted() {
return
}
session := sessions.Default(c)
tid, ok := session.Get("tid").(string)
if !ok || tid != models.UserTypeLocal.String() {
response.Error(c, response.ErrLocalUserRequired, nil)
return
}
c.Next()
}
}
func noCacheMiddleware() gin.HandlerFunc {
return func(c *gin.Context) {
c.Header("Cache-Control", "no-cache, no-store, must-revalidate") // HTTP 1.1
c.Header("Pragma", "no-cache") // HTTP 1.0
c.Header("Expires", "0") // prevents caching at the proxy server
c.Next()
}
}
+38
View File
@@ -0,0 +1,38 @@
package models
import (
"time"
"github.com/jinzhu/gorm"
)
// Agentlog is model to contain agent task and result information
// nolint:lll
type Agentlog struct {
ID uint64 `form:"id" json:"id" validate:"min=0,numeric" gorm:"type:BIGINT;NOT NULL;PRIMARY_KEY;AUTO_INCREMENT"`
Initiator MsgchainType `json:"initiator" validate:"valid,required" gorm:"type:MSGCHAIN_TYPE;NOT NULL"`
Executor MsgchainType `json:"executor" validate:"valid,required" gorm:"type:MSGCHAIN_TYPE;NOT NULL"`
Task string `json:"task" validate:"required" gorm:"type:TEXT;NOT NULL"`
Result string `json:"result" validate:"omitempty" gorm:"type:TEXT;NOT NULL;default:''"`
FlowID uint64 `form:"flow_id" json:"flow_id" validate:"min=0,numeric,required" gorm:"type:BIGINT;NOT NULL"`
TaskID *uint64 `form:"task_id,omitempty" json:"task_id,omitempty" validate:"omitnil,min=0" gorm:"type:BIGINT;NOT NULL"`
SubtaskID *uint64 `form:"subtask_id,omitempty" json:"subtask_id,omitempty" validate:"omitnil,min=0" gorm:"type:BIGINT;NOT NULL"`
CreatedAt time.Time `form:"created_at,omitempty" json:"created_at,omitempty" validate:"omitempty" gorm:"type:TIMESTAMPTZ;default:CURRENT_TIMESTAMP"`
}
// TableName returns the table name string to guaranty use correct table
func (ml *Agentlog) TableName() string {
return "agentlogs"
}
// Valid is function to control input/output data
func (ml Agentlog) Valid() error {
return validate.Struct(ml)
}
// Validate is function to use callback to control input/output data
func (ml Agentlog) Validate(db *gorm.DB) {
if err := ml.Valid(); err != nil {
db.AddError(err)
}
}
+546
View File
@@ -0,0 +1,546 @@
package models
import (
"fmt"
"time"
"github.com/jinzhu/gorm"
)
// UsageStatsPeriod represents time period enum for analytics
type UsageStatsPeriod string
const (
UsageStatsPeriodWeek UsageStatsPeriod = "week"
UsageStatsPeriodMonth UsageStatsPeriod = "month"
UsageStatsPeriodQuarter UsageStatsPeriod = "quarter"
)
func (p UsageStatsPeriod) String() string {
return string(p)
}
// Valid is function to control input/output data
func (p UsageStatsPeriod) Valid() error {
switch p {
case UsageStatsPeriodWeek,
UsageStatsPeriodMonth,
UsageStatsPeriodQuarter:
return nil
default:
return fmt.Errorf("invalid UsageStatsPeriod: %s", p)
}
}
// Validate is function to use callback to control input/output data
func (p UsageStatsPeriod) Validate(db *gorm.DB) {
if err := p.Valid(); err != nil {
db.AddError(err)
}
}
// ==================== Basic Statistics Structures ====================
// UsageStats represents token usage statistics
// nolint:lll
type UsageStats struct {
TotalUsageIn int `json:"total_usage_in" validate:"min=0"`
TotalUsageOut int `json:"total_usage_out" validate:"min=0"`
TotalUsageCacheIn int `json:"total_usage_cache_in" validate:"min=0"`
TotalUsageCacheOut int `json:"total_usage_cache_out" validate:"min=0"`
TotalUsageCostIn float64 `json:"total_usage_cost_in" validate:"min=0"`
TotalUsageCostOut float64 `json:"total_usage_cost_out" validate:"min=0"`
}
// Valid is function to control input/output data
func (u UsageStats) Valid() error {
return validate.Struct(u)
}
// Validate is function to use callback to control input/output data
func (u UsageStats) Validate(db *gorm.DB) {
if err := u.Valid(); err != nil {
db.AddError(err)
}
}
// ToolcallsStats represents toolcalls statistics
// nolint:lll
type ToolcallsStats struct {
TotalCount int `json:"total_count" validate:"min=0"`
TotalDurationSeconds float64 `json:"total_duration_seconds" validate:"min=0"`
}
// Valid is function to control input/output data
func (t ToolcallsStats) Valid() error {
return validate.Struct(t)
}
// Validate is function to use callback to control input/output data
func (t ToolcallsStats) Validate(db *gorm.DB) {
if err := t.Valid(); err != nil {
db.AddError(err)
}
}
// FlowsStats represents flows/tasks/subtasks counts
// nolint:lll
type FlowsStats struct {
TotalFlowsCount int `json:"total_flows_count" validate:"min=0"`
TotalTasksCount int `json:"total_tasks_count" validate:"min=0"`
TotalSubtasksCount int `json:"total_subtasks_count" validate:"min=0"`
TotalAssistantsCount int `json:"total_assistants_count" validate:"min=0"`
}
// Valid is function to control input/output data
func (f FlowsStats) Valid() error {
return validate.Struct(f)
}
// Validate is function to use callback to control input/output data
func (f FlowsStats) Validate(db *gorm.DB) {
if err := f.Valid(); err != nil {
db.AddError(err)
}
}
// FlowStats represents single flow statistics
// nolint:lll
type FlowStats struct {
TotalTasksCount int `json:"total_tasks_count" validate:"min=0"`
TotalSubtasksCount int `json:"total_subtasks_count" validate:"min=0"`
TotalAssistantsCount int `json:"total_assistants_count" validate:"min=0"`
}
// Valid is function to control input/output data
func (f FlowStats) Valid() error {
return validate.Struct(f)
}
// Validate is function to use callback to control input/output data
func (f FlowStats) Validate(db *gorm.DB) {
if err := f.Valid(); err != nil {
db.AddError(err)
}
}
// ==================== Time-series Statistics ====================
// DailyUsageStats for time-series usage data
// nolint:lll
type DailyUsageStats struct {
Date time.Time `json:"date" validate:"required"`
Stats *UsageStats `json:"stats" validate:"required"`
}
// Valid is function to control input/output data
func (d DailyUsageStats) Valid() error {
if err := validate.Struct(d); err != nil {
return err
}
if d.Stats != nil {
return d.Stats.Valid()
}
return nil
}
// Validate is function to use callback to control input/output data
func (d DailyUsageStats) Validate(db *gorm.DB) {
if err := d.Valid(); err != nil {
db.AddError(err)
}
}
// DailyToolcallsStats for time-series toolcalls data
// nolint:lll
type DailyToolcallsStats struct {
Date time.Time `json:"date" validate:"required"`
Stats *ToolcallsStats `json:"stats" validate:"required"`
}
// Valid is function to control input/output data
func (d DailyToolcallsStats) Valid() error {
if err := validate.Struct(d); err != nil {
return err
}
if d.Stats != nil {
return d.Stats.Valid()
}
return nil
}
// Validate is function to use callback to control input/output data
func (d DailyToolcallsStats) Validate(db *gorm.DB) {
if err := d.Valid(); err != nil {
db.AddError(err)
}
}
// DailyFlowsStats for time-series flows data
// nolint:lll
type DailyFlowsStats struct {
Date time.Time `json:"date" validate:"required"`
Stats *FlowsStats `json:"stats" validate:"required"`
}
// Valid is function to control input/output data
func (d DailyFlowsStats) Valid() error {
if err := validate.Struct(d); err != nil {
return err
}
if d.Stats != nil {
return d.Stats.Valid()
}
return nil
}
// Validate is function to use callback to control input/output data
func (d DailyFlowsStats) Validate(db *gorm.DB) {
if err := d.Valid(); err != nil {
db.AddError(err)
}
}
// ==================== Grouped Statistics ====================
// ProviderUsageStats for provider-specific usage statistics
// nolint:lll
type ProviderUsageStats struct {
Provider string `json:"provider" validate:"required"`
Stats *UsageStats `json:"stats" validate:"required"`
}
// Valid is function to control input/output data
func (p ProviderUsageStats) Valid() error {
if err := validate.Struct(p); err != nil {
return err
}
if p.Stats != nil {
return p.Stats.Valid()
}
return nil
}
// Validate is function to use callback to control input/output data
func (p ProviderUsageStats) Validate(db *gorm.DB) {
if err := p.Valid(); err != nil {
db.AddError(err)
}
}
// ModelUsageStats for model-specific usage statistics
// nolint:lll
type ModelUsageStats struct {
Model string `json:"model" validate:"required"`
Provider string `json:"provider" validate:"required"`
Stats *UsageStats `json:"stats" validate:"required"`
}
// Valid is function to control input/output data
func (m ModelUsageStats) Valid() error {
if err := validate.Struct(m); err != nil {
return err
}
if m.Stats != nil {
return m.Stats.Valid()
}
return nil
}
// Validate is function to use callback to control input/output data
func (m ModelUsageStats) Validate(db *gorm.DB) {
if err := m.Valid(); err != nil {
db.AddError(err)
}
}
// AgentTypeUsageStats for agent type usage statistics
// nolint:lll
type AgentTypeUsageStats struct {
AgentType MsgchainType `json:"agent_type" validate:"valid,required"`
Stats *UsageStats `json:"stats" validate:"required"`
}
// Valid is function to control input/output data
func (a AgentTypeUsageStats) Valid() error {
if err := validate.Struct(a); err != nil {
return err
}
if a.Stats != nil {
return a.Stats.Valid()
}
return nil
}
// Validate is function to use callback to control input/output data
func (a AgentTypeUsageStats) Validate(db *gorm.DB) {
if err := a.Valid(); err != nil {
db.AddError(err)
}
}
// FunctionToolcallsStats for function-specific toolcalls statistics
// nolint:lll
type FunctionToolcallsStats struct {
FunctionName string `json:"function_name" validate:"required"`
IsAgent bool `json:"is_agent"`
TotalCount int `json:"total_count" validate:"min=0"`
TotalDurationSeconds float64 `json:"total_duration_seconds" validate:"min=0"`
AvgDurationSeconds float64 `json:"avg_duration_seconds" validate:"min=0"`
}
// Valid is function to control input/output data
func (f FunctionToolcallsStats) Valid() error {
return validate.Struct(f)
}
// Validate is function to use callback to control input/output data
func (f FunctionToolcallsStats) Validate(db *gorm.DB) {
if err := f.Valid(); err != nil {
db.AddError(err)
}
}
// ==================== Execution Statistics ====================
// SubtaskExecutionStats represents execution statistics for a subtask
// nolint:lll
type SubtaskExecutionStats struct {
SubtaskID int64 `json:"subtask_id" validate:"min=0"`
SubtaskTitle string `json:"subtask_title" validate:"required"`
TotalDurationSeconds float64 `json:"total_duration_seconds" validate:"min=0"`
TotalToolcallsCount int `json:"total_toolcalls_count" validate:"min=0"`
}
// Valid is function to control input/output data
func (s SubtaskExecutionStats) Valid() error {
return validate.Struct(s)
}
// Validate is function to use callback to control input/output data
func (s SubtaskExecutionStats) Validate(db *gorm.DB) {
if err := s.Valid(); err != nil {
db.AddError(err)
}
}
// TaskExecutionStats represents execution statistics for a task
// nolint:lll
type TaskExecutionStats struct {
TaskID int64 `json:"task_id" validate:"min=0"`
TaskTitle string `json:"task_title" validate:"required"`
TotalDurationSeconds float64 `json:"total_duration_seconds" validate:"min=0"`
TotalToolcallsCount int `json:"total_toolcalls_count" validate:"min=0"`
Subtasks []SubtaskExecutionStats `json:"subtasks" validate:"omitempty"`
}
// Valid is function to control input/output data
func (t TaskExecutionStats) Valid() error {
if err := validate.Struct(t); err != nil {
return err
}
for i := range t.Subtasks {
if err := t.Subtasks[i].Valid(); err != nil {
return err
}
}
return nil
}
// Validate is function to use callback to control input/output data
func (t TaskExecutionStats) Validate(db *gorm.DB) {
if err := t.Valid(); err != nil {
db.AddError(err)
}
}
// FlowExecutionStats represents execution statistics for a flow
// nolint:lll
type FlowExecutionStats struct {
FlowID int64 `json:"flow_id" validate:"min=0"`
FlowTitle string `json:"flow_title" validate:"required"`
TotalDurationSeconds float64 `json:"total_duration_seconds" validate:"min=0"`
TotalToolcallsCount int `json:"total_toolcalls_count" validate:"min=0"`
TotalAssistantsCount int `json:"total_assistants_count" validate:"min=0"`
Tasks []TaskExecutionStats `json:"tasks" validate:"omitempty"`
}
// Valid is function to control input/output data
func (f FlowExecutionStats) Valid() error {
if err := validate.Struct(f); err != nil {
return err
}
for i := range f.Tasks {
if err := f.Tasks[i].Valid(); err != nil {
return err
}
}
return nil
}
// Validate is function to use callback to control input/output data
func (f FlowExecutionStats) Validate(db *gorm.DB) {
if err := f.Valid(); err != nil {
db.AddError(err)
}
}
// ==================== Aggregated Response Models ====================
// SystemUsageResponse represents system-wide analytics response
// nolint:lll
type SystemUsageResponse struct {
UsageStatsTotal *UsageStats `json:"usage_stats_total" validate:"required"`
ToolcallsStatsTotal *ToolcallsStats `json:"toolcalls_stats_total" validate:"required"`
FlowsStatsTotal *FlowsStats `json:"flows_stats_total" validate:"required"`
UsageStatsByProvider []ProviderUsageStats `json:"usage_stats_by_provider" validate:"omitempty"`
UsageStatsByModel []ModelUsageStats `json:"usage_stats_by_model" validate:"omitempty"`
UsageStatsByAgentType []AgentTypeUsageStats `json:"usage_stats_by_agent_type" validate:"omitempty"`
ToolcallsStatsByFunction []FunctionToolcallsStats `json:"toolcalls_stats_by_function" validate:"omitempty"`
}
// Valid is function to control input/output data
func (s SystemUsageResponse) Valid() error {
if err := validate.Struct(s); err != nil {
return err
}
if s.UsageStatsTotal != nil {
if err := s.UsageStatsTotal.Valid(); err != nil {
return err
}
}
if s.ToolcallsStatsTotal != nil {
if err := s.ToolcallsStatsTotal.Valid(); err != nil {
return err
}
}
if s.FlowsStatsTotal != nil {
if err := s.FlowsStatsTotal.Valid(); err != nil {
return err
}
}
for i := range s.UsageStatsByProvider {
if err := s.UsageStatsByProvider[i].Valid(); err != nil {
return err
}
}
for i := range s.UsageStatsByModel {
if err := s.UsageStatsByModel[i].Valid(); err != nil {
return err
}
}
for i := range s.UsageStatsByAgentType {
if err := s.UsageStatsByAgentType[i].Valid(); err != nil {
return err
}
}
for i := range s.ToolcallsStatsByFunction {
if err := s.ToolcallsStatsByFunction[i].Valid(); err != nil {
return err
}
}
return nil
}
// Validate is function to use callback to control input/output data
func (s SystemUsageResponse) Validate(db *gorm.DB) {
if err := s.Valid(); err != nil {
db.AddError(err)
}
}
// PeriodUsageResponse represents period-based analytics response
// nolint:lll
type PeriodUsageResponse struct {
Period string `json:"period" validate:"required"`
UsageStatsByPeriod []DailyUsageStats `json:"usage_stats_by_period" validate:"omitempty"`
ToolcallsStatsByPeriod []DailyToolcallsStats `json:"toolcalls_stats_by_period" validate:"omitempty"`
FlowsStatsByPeriod []DailyFlowsStats `json:"flows_stats_by_period" validate:"omitempty"`
FlowsExecutionStatsByPeriod []FlowExecutionStats `json:"flows_execution_stats_by_period" validate:"omitempty"`
}
// Valid is function to control input/output data
func (p PeriodUsageResponse) Valid() error {
if err := validate.Struct(p); err != nil {
return err
}
for i := range p.UsageStatsByPeriod {
if err := p.UsageStatsByPeriod[i].Valid(); err != nil {
return err
}
}
for i := range p.ToolcallsStatsByPeriod {
if err := p.ToolcallsStatsByPeriod[i].Valid(); err != nil {
return err
}
}
for i := range p.FlowsStatsByPeriod {
if err := p.FlowsStatsByPeriod[i].Valid(); err != nil {
return err
}
}
for i := range p.FlowsExecutionStatsByPeriod {
if err := p.FlowsExecutionStatsByPeriod[i].Valid(); err != nil {
return err
}
}
return nil
}
// Validate is function to use callback to control input/output data
func (p PeriodUsageResponse) Validate(db *gorm.DB) {
if err := p.Valid(); err != nil {
db.AddError(err)
}
}
// FlowUsageResponse represents flow-specific analytics response
// nolint:lll
type FlowUsageResponse struct {
FlowID int64 `json:"flow_id" validate:"min=0"`
UsageStatsByFlow *UsageStats `json:"usage_stats_by_flow" validate:"required"`
UsageStatsByAgentTypeForFlow []AgentTypeUsageStats `json:"usage_stats_by_agent_type_for_flow" validate:"omitempty"`
ToolcallsStatsByFlow *ToolcallsStats `json:"toolcalls_stats_by_flow" validate:"required"`
ToolcallsStatsByFunctionForFlow []FunctionToolcallsStats `json:"toolcalls_stats_by_function_for_flow" validate:"omitempty"`
FlowStatsByFlow *FlowStats `json:"flow_stats_by_flow" validate:"required"`
}
// Valid is function to control input/output data
func (f FlowUsageResponse) Valid() error {
if err := validate.Struct(f); err != nil {
return err
}
if f.UsageStatsByFlow != nil {
if err := f.UsageStatsByFlow.Valid(); err != nil {
return err
}
}
for i := range f.UsageStatsByAgentTypeForFlow {
if err := f.UsageStatsByAgentTypeForFlow[i].Valid(); err != nil {
return err
}
}
if f.ToolcallsStatsByFlow != nil {
if err := f.ToolcallsStatsByFlow.Valid(); err != nil {
return err
}
}
for i := range f.ToolcallsStatsByFunctionForFlow {
if err := f.ToolcallsStatsByFunctionForFlow[i].Valid(); err != nil {
return err
}
}
if f.FlowStatsByFlow != nil {
if err := f.FlowStatsByFlow.Valid(); err != nil {
return err
}
}
return nil
}
// Validate is function to use callback to control input/output data
func (f FlowUsageResponse) Validate(db *gorm.DB) {
if err := f.Valid(); err != nil {
db.AddError(err)
}
}
+133
View File
@@ -0,0 +1,133 @@
package models
import (
"fmt"
"time"
"github.com/golang-jwt/jwt/v5"
"github.com/jinzhu/gorm"
)
// TokenStatus represents the status of an API token
type TokenStatus string
const (
TokenStatusActive TokenStatus = "active"
TokenStatusRevoked TokenStatus = "revoked"
TokenStatusExpired TokenStatus = "expired"
)
func (s TokenStatus) String() string {
return string(s)
}
// Valid is function to control input/output data
func (s TokenStatus) Valid() error {
switch s {
case TokenStatusActive, TokenStatusRevoked, TokenStatusExpired:
return nil
default:
return fmt.Errorf("invalid TokenStatus: %s", s)
}
}
// Validate is function to use callback to control input/output data
func (s TokenStatus) Validate(db *gorm.DB) {
if err := s.Valid(); err != nil {
db.AddError(err)
}
}
// APIToken is model to contain API token metadata
// nolint:lll
type APIToken struct {
ID uint64 `form:"id" json:"id" validate:"min=0,numeric" gorm:"type:BIGINT;NOT NULL;PRIMARY_KEY;AUTO_INCREMENT"`
TokenID string `form:"token_id" json:"token_id" validate:"required,len=10" gorm:"type:TEXT;NOT NULL;UNIQUE_INDEX"`
UserID uint64 `form:"user_id" json:"user_id" validate:"min=0,numeric" gorm:"type:BIGINT;NOT NULL"`
RoleID uint64 `form:"role_id" json:"role_id" validate:"min=0,numeric" gorm:"type:BIGINT;NOT NULL"`
Name *string `form:"name,omitempty" json:"name,omitempty" validate:"omitempty,max=100" gorm:"type:TEXT"`
TTL uint64 `form:"ttl" json:"ttl" validate:"required,min=60,max=94608000" gorm:"type:BIGINT;NOT NULL"`
Status TokenStatus `form:"status" json:"status" validate:"valid,required" gorm:"type:TOKEN_STATUS;NOT NULL;default:'active'"`
CreatedAt time.Time `form:"created_at" json:"created_at" validate:"required" gorm:"type:TIMESTAMPTZ;NOT NULL;default:CURRENT_TIMESTAMP"`
UpdatedAt time.Time `form:"updated_at" json:"updated_at" validate:"required" gorm:"type:TIMESTAMPTZ;NOT NULL;default:CURRENT_TIMESTAMP"`
DeletedAt *time.Time `form:"deleted_at,omitempty" json:"deleted_at,omitempty" validate:"omitempty" sql:"index" gorm:"type:TIMESTAMPTZ"`
}
// TableName returns the table name string to guaranty use correct table
func (at *APIToken) TableName() string {
return "api_tokens"
}
// Valid is function to control input/output data
func (at APIToken) Valid() error {
if err := at.Status.Valid(); err != nil {
return err
}
return validate.Struct(at)
}
// Validate is function to use callback to control input/output data
func (at APIToken) Validate(db *gorm.DB) {
if err := at.Valid(); err != nil {
db.AddError(err)
}
}
// APITokenWithSecret is model to contain API token with the JWT token string (returned only on creation)
// nolint:lll
type APITokenWithSecret struct {
APIToken `form:"" json:""`
Token string `form:"token" json:"token" validate:"required,jwt" gorm:"-"`
}
// Valid is function to control input/output data
func (ats APITokenWithSecret) Valid() error {
if err := ats.APIToken.Valid(); err != nil {
return err
}
return validate.Struct(ats)
}
// CreateAPITokenRequest is model to contain request data for creating an API token
// nolint:lll
type CreateAPITokenRequest struct {
Name *string `form:"name,omitempty" json:"name,omitempty" validate:"omitempty,max=100"`
TTL uint64 `form:"ttl" json:"ttl" validate:"required,min=60,max=94608000"` // from 1 minute to 3 years
}
// Valid is function to control input/output data
func (catr CreateAPITokenRequest) Valid() error {
return validate.Struct(catr)
}
// UpdateAPITokenRequest is model to contain request data for updating an API token
// nolint:lll
type UpdateAPITokenRequest struct {
Name *string `form:"name,omitempty" json:"name,omitempty" validate:"omitempty,max=100"`
Status TokenStatus `form:"status,omitempty" json:"status,omitempty" validate:"omitempty,valid"`
}
// Valid is function to control input/output data
func (uatr UpdateAPITokenRequest) Valid() error {
if uatr.Status != "" {
if err := uatr.Status.Valid(); err != nil {
return err
}
}
return validate.Struct(uatr)
}
// APITokenClaims is model to contain JWT claims for API tokens
// nolint:lll
type APITokenClaims struct {
TokenID string `json:"tid" validate:"required,len=10"`
RID uint64 `json:"rid" validate:"min=0,max=10000"`
UID uint64 `json:"uid" validate:"min=0,max=10000"`
UHASH string `json:"uhash" validate:"required"`
jwt.RegisteredClaims
}
// Valid is function to control input/output data
func (atc APITokenClaims) Valid() error {
return validate.Struct(atc)
}
@@ -0,0 +1,255 @@
package models
import (
"testing"
"time"
"github.com/stretchr/testify/assert"
)
func TestTokenStatusValid(t *testing.T) {
t.Parallel()
tests := []struct {
name string
status TokenStatus
wantErr bool
}{
{"valid active", TokenStatusActive, false},
{"valid revoked", TokenStatusRevoked, false},
{"valid expired", TokenStatusExpired, false},
{"invalid empty", TokenStatus(""), true},
{"invalid unknown", TokenStatus("unknown"), true},
{"invalid deleted", TokenStatus("deleted"), true},
}
for _, tt := range tests {
tt := tt
t.Run(tt.name, func(t *testing.T) {
t.Parallel()
err := tt.status.Valid()
if tt.wantErr {
assert.Error(t, err)
assert.Contains(t, err.Error(), "invalid TokenStatus")
} else {
assert.NoError(t, err)
}
})
}
}
func TestTokenStatusString(t *testing.T) {
t.Parallel()
assert.Equal(t, "active", TokenStatusActive.String())
assert.Equal(t, "revoked", TokenStatusRevoked.String())
assert.Equal(t, "expired", TokenStatusExpired.String())
}
func TestAPITokenValid(t *testing.T) {
t.Parallel()
now := time.Now()
validToken := APIToken{
TokenID: "abcdefghij",
UserID: 1,
RoleID: 1,
TTL: 3600,
Status: TokenStatusActive,
CreatedAt: now,
UpdatedAt: now,
}
t.Run("valid token", func(t *testing.T) {
t.Parallel()
assert.NoError(t, validToken.Valid())
})
t.Run("invalid status", func(t *testing.T) {
t.Parallel()
token := validToken
token.Status = TokenStatus("invalid")
assert.Error(t, token.Valid())
})
t.Run("token id wrong length", func(t *testing.T) {
t.Parallel()
token := validToken
token.TokenID = "short"
assert.Error(t, token.Valid())
})
t.Run("ttl too small", func(t *testing.T) {
t.Parallel()
token := validToken
token.TTL = 10
assert.Error(t, token.Valid())
})
t.Run("ttl too large", func(t *testing.T) {
t.Parallel()
token := validToken
token.TTL = 94608001
assert.Error(t, token.Valid())
})
t.Run("ttl minimum boundary", func(t *testing.T) {
t.Parallel()
token := validToken
token.TTL = 60
assert.NoError(t, token.Valid())
})
t.Run("ttl maximum boundary", func(t *testing.T) {
t.Parallel()
token := validToken
token.TTL = 94608000
assert.NoError(t, token.Valid())
})
}
func TestAPITokenTableName(t *testing.T) {
t.Parallel()
at := &APIToken{}
assert.Equal(t, "api_tokens", at.TableName())
}
func TestCreateAPITokenRequestValid(t *testing.T) {
t.Parallel()
t.Run("valid request", func(t *testing.T) {
t.Parallel()
req := CreateAPITokenRequest{TTL: 3600}
assert.NoError(t, req.Valid())
})
t.Run("valid request with name", func(t *testing.T) {
t.Parallel()
name := "my-token"
req := CreateAPITokenRequest{Name: &name, TTL: 3600}
assert.NoError(t, req.Valid())
})
t.Run("ttl too small", func(t *testing.T) {
t.Parallel()
req := CreateAPITokenRequest{TTL: 30}
assert.Error(t, req.Valid())
})
t.Run("zero ttl", func(t *testing.T) {
t.Parallel()
req := CreateAPITokenRequest{TTL: 0}
assert.Error(t, req.Valid())
})
}
func TestUpdateAPITokenRequestValid(t *testing.T) {
t.Parallel()
t.Run("valid update with status", func(t *testing.T) {
t.Parallel()
req := UpdateAPITokenRequest{Status: TokenStatusRevoked}
assert.NoError(t, req.Valid())
})
t.Run("valid empty update", func(t *testing.T) {
t.Parallel()
req := UpdateAPITokenRequest{}
assert.NoError(t, req.Valid())
})
t.Run("invalid status", func(t *testing.T) {
t.Parallel()
req := UpdateAPITokenRequest{Status: TokenStatus("invalid")}
assert.Error(t, req.Valid())
})
}
func TestAPITokenWithSecretValid(t *testing.T) {
t.Parallel()
now := time.Now()
validJWT := "eyJhbGciOiJSUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiIxMjM0NTY3ODkwIn0.POstGetfAytaZS82wHcjoTyoqhMyxXiWdR7Nn7A29DNSl0EiXLdwJ6xC6AfgZWF1bOsS_TuYI3OG85AmiExREkrS6tDfTQ2B3WXlrr-wp5AokiRbz3_oB4OxG-W9KcEEbDRcZc0nH3L7LzYptiy1PtAylQGxHTWZXtGz4ht0bAecBgmpdgXMguEIcoqPJ1n3pIWk_dUZegpqx0Lka21H6XxUTxiy8OcaarA8zdnPUnV6AmNP3ecFawIFYdvJB_cm-GvpCSbr8G8y_Mllj8f4x9nBH8pQux89_6gUY618iYv7tuPWBFfEbLxtF2pZS6YC1aSfLQxaOoaBSTNRg"
baseToken := APIToken{
TokenID: "abcdefghij",
UserID: 1,
RoleID: 1,
TTL: 3600,
Status: TokenStatusActive,
CreatedAt: now,
UpdatedAt: now,
}
t.Run("valid token with secret", func(t *testing.T) {
t.Parallel()
ats := APITokenWithSecret{
APIToken: baseToken,
Token: validJWT,
}
assert.NoError(t, ats.Valid())
})
t.Run("invalid embedded api token", func(t *testing.T) {
t.Parallel()
badToken := baseToken
badToken.Status = TokenStatus("invalid")
ats := APITokenWithSecret{
APIToken: badToken,
Token: validJWT,
}
assert.Error(t, ats.Valid())
})
t.Run("invalid jwt token string", func(t *testing.T) {
t.Parallel()
ats := APITokenWithSecret{
APIToken: baseToken,
Token: "not-a-jwt",
}
assert.Error(t, ats.Valid())
})
}
func TestAPITokenClaimsValid(t *testing.T) {
t.Parallel()
t.Run("valid claims", func(t *testing.T) {
t.Parallel()
claims := APITokenClaims{
TokenID: "abcdefghij",
RID: 1,
UID: 1,
UHASH: "somehash",
}
assert.NoError(t, claims.Valid())
})
t.Run("missing token id", func(t *testing.T) {
t.Parallel()
claims := APITokenClaims{
TokenID: "",
UHASH: "somehash",
}
assert.Error(t, claims.Valid())
})
t.Run("missing uhash", func(t *testing.T) {
t.Parallel()
claims := APITokenClaims{
TokenID: "abcdefghij",
UHASH: "",
}
assert.Error(t, claims.Valid())
})
t.Run("uid too large", func(t *testing.T) {
t.Parallel()
claims := APITokenClaims{
TokenID: "abcdefghij",
UID: 10001,
UHASH: "somehash",
}
assert.Error(t, claims.Valid())
})
}
@@ -0,0 +1,38 @@
package models
import (
"time"
"github.com/jinzhu/gorm"
)
// Assistantlog is model to contain log record information from agents about their actions
// nolint:lll
type Assistantlog struct {
ID uint64 `form:"id" json:"id" validate:"min=0,numeric" gorm:"type:BIGINT;NOT NULL;PRIMARY_KEY;AUTO_INCREMENT"`
Type MsglogType `form:"type" json:"type" validate:"valid,required" gorm:"type:MSGLOG_TYPE;NOT NULL"`
Message string `form:"message" json:"message" validate:"omitempty" gorm:"type:TEXT;NOT NULL"`
Thinking string `form:"thinking" json:"thinking" validate:"omitempty" gorm:"type:TEXT;NULL"`
Result string `form:"result" json:"result" validate:"omitempty" gorm:"type:TEXT;NOT NULL;default:''"`
ResultFormat MsglogResultFormat `form:"result_format" json:"result_format" validate:"valid,required" gorm:"type:MSGLOG_RESULT_FORMAT;NOT NULL;default:plain"`
FlowID uint64 `form:"flow_id" json:"flow_id" validate:"min=0,numeric,required" gorm:"type:BIGINT;NOT NULL"`
AssistantID uint64 `form:"assistant_id" json:"assistant_id" validate:"min=0,numeric,required" gorm:"type:BIGINT;NOT NULL"`
CreatedAt time.Time `form:"created_at,omitempty" json:"created_at,omitempty" validate:"omitempty" gorm:"type:TIMESTAMPTZ;default:CURRENT_TIMESTAMP"`
}
// TableName returns the table name string to guaranty use correct table
func (al *Assistantlog) TableName() string {
return "assistantlogs"
}
// Valid is function to control input/output data
func (al Assistantlog) Valid() error {
return validate.Struct(al)
}
// Validate is function to use callback to control input/output data
func (al Assistantlog) Validate(db *gorm.DB) {
if err := al.Valid(); err != nil {
db.AddError(err)
}
}
+133
View File
@@ -0,0 +1,133 @@
package models
import (
"fmt"
"pentagi/pkg/tools"
"time"
"github.com/jinzhu/gorm"
)
type AssistantStatus string
const (
AssistantStatusCreated AssistantStatus = "created"
AssistantStatusRunning AssistantStatus = "running"
AssistantStatusWaiting AssistantStatus = "waiting"
AssistantStatusFinished AssistantStatus = "finished"
AssistantStatusFailed AssistantStatus = "failed"
)
func (s AssistantStatus) String() string {
return string(s)
}
// Valid is function to control input/output data
func (s AssistantStatus) Valid() error {
switch s {
case AssistantStatusCreated,
AssistantStatusRunning,
AssistantStatusWaiting,
AssistantStatusFinished,
AssistantStatusFailed:
return nil
default:
return fmt.Errorf("invalid AssistantStatus: %s", s)
}
}
// Validate is function to use callback to control input/output data
func (s AssistantStatus) Validate(db *gorm.DB) {
if err := s.Valid(); err != nil {
db.AddError(err)
}
}
// Assistant is model to contain assistant information
// nolint:lll
type Assistant struct {
ID uint64 `form:"id" json:"id" validate:"min=0,numeric" gorm:"type:BIGINT;NOT NULL;PRIMARY_KEY;AUTO_INCREMENT"`
Status AssistantStatus `form:"status" json:"status" validate:"valid,required" gorm:"type:ASSISTANT_STATUS;NOT NULL;default:'created'"`
Title string `form:"title" json:"title" validate:"required" gorm:"type:TEXT;NOT NULL;default:'untitled'"`
Model string `form:"model" json:"model" validate:"max=70,required" gorm:"type:TEXT;NOT NULL"`
ModelProviderName string `form:"model_provider_name" json:"model_provider_name" validate:"max=70,required" gorm:"type:TEXT;NOT NULL"`
ModelProviderType ProviderType `form:"model_provider_type" json:"model_provider_type" validate:"valid,required" gorm:"type:PROVIDER_TYPE;NOT NULL"`
Language string `form:"language" json:"language" validate:"max=70,required" gorm:"type:TEXT;NOT NULL"`
Functions *tools.Functions `form:"functions,omitempty" json:"functions,omitempty" validate:"omitempty,valid" gorm:"type:JSON;NOT NULL;default:'{}'"`
FlowID uint64 `form:"flow_id" json:"flow_id" validate:"min=0,numeric,required" gorm:"type:BIGINT;NOT NULL"`
MsgchainID *uint64 `form:"msgchain_id" json:"msgchain_id" validate:"min=0,numeric" gorm:"type:BIGINT;NOT NULL"`
TraceID *string `form:"trace_id" json:"trace_id" validate:"max=70,required" gorm:"type:TEXT;NOT NULL"`
ToolCallIDTemplate string `form:"tool_call_id_template" json:"tool_call_id_template" validate:"max=70,required" gorm:"type:TEXT;NOT NULL"`
UseAgents bool `form:"use_agents" json:"use_agents" validate:"omitempty" gorm:"type:BOOLEAN;NOT NULL;default:false"`
CreatedAt time.Time `form:"created_at,omitempty" json:"created_at,omitempty" validate:"omitempty" gorm:"type:TIMESTAMPTZ;default:CURRENT_TIMESTAMP"`
UpdatedAt time.Time `form:"updated_at,omitempty" json:"updated_at,omitempty" validate:"omitempty" gorm:"type:TIMESTAMPTZ;default:CURRENT_TIMESTAMP"`
DeletedAt *time.Time `form:"deleted_at,omitempty" json:"deleted_at,omitempty" validate:"omitempty" sql:"index" gorm:"type:TIMESTAMPTZ"`
}
// TableName returns the table name string to guaranty use correct table
func (a *Assistant) TableName() string {
return "assistants"
}
// Valid is function to control input/output data
func (a Assistant) Valid() error {
return validate.Struct(a)
}
// Validate is function to use callback to control input/output data
func (a Assistant) Validate(db *gorm.DB) {
if err := a.Valid(); err != nil {
db.AddError(err)
}
}
// CreateAssistant is model to contain assistant creation paylaod
// nolint:lll
type CreateAssistant struct {
Input string `form:"input" json:"input" validate:"required" example:"user input for running assistant"`
Provider string `form:"provider" json:"provider" validate:"required" example:"openai"`
UseAgents bool `form:"use_agents" json:"use_agents" validate:"omitempty" example:"true"`
Functions *tools.Functions `form:"functions,omitempty" json:"functions,omitempty" validate:"omitempty,valid"`
ResourceIDs []uint64 `form:"resource_ids,omitempty" json:"resource_ids,omitempty" validate:"omitempty" swaggertype:"array,integer"`
}
// Valid is function to control input/output data
func (ca CreateAssistant) Valid() error {
return validate.Struct(ca)
}
// PatchAssistant is model to contain assistant patching paylaod
// nolint:lll
type PatchAssistant struct {
Action string `form:"action" json:"action" validate:"required,oneof=stop input" enums:"stop,input" default:"stop"`
Input *string `form:"input,omitempty" json:"input,omitempty" validate:"required_if=Action input" example:"user input for waiting assistant"`
UseAgents bool `form:"use_agents" json:"use_agents" validate:"omitempty" example:"true"`
ResourceIDs []uint64 `form:"resource_ids,omitempty" json:"resource_ids,omitempty" validate:"omitempty" swaggertype:"array,integer"`
}
// Valid is function to control input/output data
func (pa PatchAssistant) Valid() error {
return validate.Struct(pa)
}
// AssistantFlow is model to contain assistant information linked with flow
// nolint:lll
type AssistantFlow struct {
Flow Flow `form:"flow,omitempty" json:"flow,omitempty" gorm:"association_autoupdate:false;association_autocreate:false"`
Assistant `form:"" json:""`
}
// Valid is function to control input/output data
func (af AssistantFlow) Valid() error {
if err := af.Flow.Valid(); err != nil {
return err
}
return af.Assistant.Valid()
}
// Validate is function to use callback to control input/output data
func (af AssistantFlow) Validate(db *gorm.DB) {
if err := af.Valid(); err != nil {
db.AddError(err)
}
}
@@ -0,0 +1,211 @@
package models
import (
"testing"
"github.com/stretchr/testify/assert"
)
func TestAssistantStatusValid(t *testing.T) {
t.Parallel()
tests := []struct {
name string
status AssistantStatus
wantErr bool
}{
{"valid created", AssistantStatusCreated, false},
{"valid running", AssistantStatusRunning, false},
{"valid waiting", AssistantStatusWaiting, false},
{"valid finished", AssistantStatusFinished, false},
{"valid failed", AssistantStatusFailed, false},
{"invalid empty", AssistantStatus(""), true},
{"invalid unknown", AssistantStatus("unknown"), true},
{"invalid paused", AssistantStatus("paused"), true},
}
for _, tt := range tests {
tt := tt
t.Run(tt.name, func(t *testing.T) {
t.Parallel()
err := tt.status.Valid()
if tt.wantErr {
assert.Error(t, err)
assert.Contains(t, err.Error(), "invalid AssistantStatus")
} else {
assert.NoError(t, err)
}
})
}
}
func TestAssistantStatusString(t *testing.T) {
t.Parallel()
assert.Equal(t, "created", AssistantStatusCreated.String())
assert.Equal(t, "running", AssistantStatusRunning.String())
assert.Equal(t, "finished", AssistantStatusFinished.String())
}
func TestAssistantValid(t *testing.T) {
t.Parallel()
traceID := "trace-123"
msgchainID := uint64(1)
validAssistant := Assistant{
Status: AssistantStatusCreated,
Title: "test assistant",
Model: "gpt-4",
ModelProviderName: "openai",
ModelProviderType: ProviderType("openai"),
Language: "en",
ToolCallIDTemplate: "call_{id}",
TraceID: &traceID,
FlowID: 1,
MsgchainID: &msgchainID,
}
t.Run("valid assistant", func(t *testing.T) {
t.Parallel()
assert.NoError(t, validAssistant.Valid())
})
t.Run("invalid status", func(t *testing.T) {
t.Parallel()
a := validAssistant
a.Status = AssistantStatus("invalid")
assert.Error(t, a.Valid())
})
t.Run("missing title", func(t *testing.T) {
t.Parallel()
a := validAssistant
a.Title = ""
assert.Error(t, a.Valid())
})
t.Run("missing model", func(t *testing.T) {
t.Parallel()
a := validAssistant
a.Model = ""
assert.Error(t, a.Valid())
})
}
func TestAssistantTableName(t *testing.T) {
t.Parallel()
a := &Assistant{}
assert.Equal(t, "assistants", a.TableName())
}
func TestCreateAssistantValid(t *testing.T) {
t.Parallel()
t.Run("valid create assistant", func(t *testing.T) {
t.Parallel()
ca := CreateAssistant{Input: "hello", Provider: "openai"}
assert.NoError(t, ca.Valid())
})
t.Run("missing input", func(t *testing.T) {
t.Parallel()
ca := CreateAssistant{Input: "", Provider: "openai"}
assert.Error(t, ca.Valid())
})
t.Run("missing provider", func(t *testing.T) {
t.Parallel()
ca := CreateAssistant{Input: "hello", Provider: ""}
assert.Error(t, ca.Valid())
})
}
func TestPatchAssistantValid(t *testing.T) {
t.Parallel()
t.Run("valid stop action", func(t *testing.T) {
t.Parallel()
pa := PatchAssistant{Action: "stop"}
assert.NoError(t, pa.Valid())
})
t.Run("valid input action with input", func(t *testing.T) {
t.Parallel()
input := "user response"
pa := PatchAssistant{Action: "input", Input: &input}
assert.NoError(t, pa.Valid())
})
t.Run("invalid action", func(t *testing.T) {
t.Parallel()
pa := PatchAssistant{Action: "restart"}
assert.Error(t, pa.Valid())
})
t.Run("input action without input", func(t *testing.T) {
t.Parallel()
pa := PatchAssistant{Action: "input"}
assert.Error(t, pa.Valid())
})
}
func TestAssistantFlowValid(t *testing.T) {
t.Parallel()
traceID := "trace-123"
mcID := uint64(1)
validAssistant := Assistant{
Status: AssistantStatusCreated,
Title: "test",
Model: "gpt-4",
ModelProviderName: "openai",
ModelProviderType: ProviderType("openai"),
Language: "en",
ToolCallIDTemplate: "call_{id}",
TraceID: &traceID,
FlowID: 1,
MsgchainID: &mcID,
}
validFlow := Flow{
Status: FlowStatusCreated,
Title: "flow",
Model: "gpt-4",
ModelProviderName: "openai",
ModelProviderType: ProviderType("openai"),
Language: "en",
ToolCallIDTemplate: "call_{id}",
TraceID: &traceID,
UserID: 1,
}
t.Run("valid assistant flow", func(t *testing.T) {
t.Parallel()
af := AssistantFlow{
Flow: validFlow,
Assistant: validAssistant,
}
assert.NoError(t, af.Valid())
})
t.Run("invalid flow", func(t *testing.T) {
t.Parallel()
badFlow := validFlow
badFlow.Title = ""
af := AssistantFlow{
Flow: badFlow,
Assistant: validAssistant,
}
assert.Error(t, af.Valid())
})
t.Run("invalid assistant", func(t *testing.T) {
t.Parallel()
badAssistant := validAssistant
badAssistant.Title = ""
af := AssistantFlow{
Flow: validFlow,
Assistant: badAssistant,
}
assert.Error(t, af.Valid())
})
}
+103
View File
@@ -0,0 +1,103 @@
package models
import (
"fmt"
"time"
"github.com/jinzhu/gorm"
)
type ContainerStatus string
const (
ContainerStatusStarting ContainerStatus = "starting"
ContainerStatusRunning ContainerStatus = "running"
ContainerStatusStopped ContainerStatus = "stopped"
ContainerStatusDeleted ContainerStatus = "deleted"
ContainerStatusFailed ContainerStatus = "failed"
)
func (s ContainerStatus) String() string {
return string(s)
}
// Valid is function to control input/output data
func (s ContainerStatus) Valid() error {
switch s {
case ContainerStatusStarting,
ContainerStatusRunning,
ContainerStatusStopped,
ContainerStatusDeleted,
ContainerStatusFailed:
return nil
default:
return fmt.Errorf("invalid ContainerStatus: %s", s)
}
}
// Validate is function to use callback to control input/output data
func (s ContainerStatus) Validate(db *gorm.DB) {
if err := s.Valid(); err != nil {
db.AddError(err)
}
}
type ContainerType string
const (
ContainerTypePrimary ContainerType = "primary"
ContainerTypeSecondary ContainerType = "secondary"
)
func (t ContainerType) String() string {
return string(t)
}
// Valid is function to control input/output data
func (t ContainerType) Valid() error {
switch t {
case ContainerTypePrimary, ContainerTypeSecondary:
return nil
default:
return fmt.Errorf("invalid ContainerType: %s", t)
}
}
// Validate is function to use callback to control input/output data
func (t ContainerType) Validate(db *gorm.DB) {
if err := t.Valid(); err != nil {
db.AddError(err)
}
}
// Container is model to contain container information
// nolint:lll
type Container struct {
ID uint64 `form:"id" json:"id" validate:"min=0,numeric" gorm:"type:BIGINT;NOT NULL;PRIMARY_KEY;AUTO_INCREMENT"`
Type ContainerType `form:"type" json:"type" validate:"valid,required" gorm:"type:CONTAINER_TYPE;NOT NULL;default:'primary'"`
Name string `form:"name" json:"name" validate:"required" gorm:"type:TEXT;NOT NULL;default:MD5(RANDOM()::text)"`
Image string `form:"image" json:"image" validate:"required" gorm:"type:TEXT;NOT NULL"`
Status ContainerStatus `form:"status" json:"status" validate:"valid,required" gorm:"type:CONTAINER_STATUS;NOT NULL;default:'starting'"`
LocalID string `form:"local_id" json:"local_id" validate:"required" gorm:"type:TEXT;NOT NULL"`
LocalDir string `form:"local_dir" json:"local_dir" validate:"required" gorm:"type:TEXT;NOT NULL"`
FlowID uint64 `form:"flow_id" json:"flow_id" validate:"min=0,numeric,required" gorm:"type:BIGINT;NOT NULL"`
CreatedAt time.Time `form:"created_at,omitempty" json:"created_at,omitempty" validate:"omitempty" gorm:"type:TIMESTAMPTZ;default:CURRENT_TIMESTAMP"`
UpdatedAt time.Time `form:"updated_at,omitempty" json:"updated_at,omitempty" validate:"omitempty" gorm:"type:TIMESTAMPTZ;default:CURRENT_TIMESTAMP"`
}
// TableName returns the table name string to guaranty use correct table
func (c *Container) TableName() string {
return "containers"
}
// Valid is function to control input/output data
func (c Container) Valid() error {
return validate.Struct(c)
}
// Validate is function to use callback to control input/output data
func (c Container) Validate(db *gorm.DB) {
if err := c.Valid(); err != nil {
db.AddError(err)
}
}
+79
View File
@@ -0,0 +1,79 @@
package models
import "time"
// FlowFile represents a single entry in the flow's local file cache.
type FlowFile struct {
ID string `json:"id"`
Name string `json:"name"`
Path string `json:"path"` // relative: "uploads/<x>" or resources/<x> or container/<x>"
Size int64 `json:"size"`
IsDir bool `json:"is_dir"`
ModifiedAt time.Time `json:"modified_at"`
}
// FlowFiles is the list response for flow file operations.
type FlowFiles struct {
Files []FlowFile `json:"files"`
Total uint64 `json:"total"`
}
// ContainerFile represents a single entry in the running container's filesystem.
type ContainerFile struct {
ID string `json:"id"`
Name string `json:"name"`
Path string `json:"path"`
Size int64 `json:"size"`
IsDir bool `json:"is_dir"`
ModifiedAt time.Time `json:"modified_at"`
}
// ContainerFiles is the list response for container directory listing.
type ContainerFiles struct {
Path string `json:"path"`
Files []ContainerFile `json:"files"`
Total uint64 `json:"total"`
}
// PullFlowFilesRequest is the request body for pulling files from a container.
type PullFlowFilesRequest struct {
// Path is a single container path (maintained for backward compatibility).
// Combined with Paths if both are provided.
Path string `json:"path"`
// Paths is a list of container paths to pull. Combined with Path if provided.
Paths []string `json:"paths"`
// Force overwrites local cache entries that already exist.
Force bool `json:"force"`
}
// AddResourcesRequest is the request body for copying user resources into a flow.
type AddResourcesRequest struct {
// IDs is the list of user resource IDs to copy into the flow resources directory.
IDs []uint64 `json:"ids" binding:"required,min=1"`
// Force overwrites files that already exist in the flow resources directory.
Force bool `json:"force"`
}
// AddResourceFromFlowRequest is the request body for promoting one or more flow
// files / directories into the user's global resource store.
//
// At least one of Source or Sources must be non-empty after deduplication.
//
// - Single source (Source XOR one entry in Sources):
// Destination is the exact target path (file) or root directory (dir tree).
// - Multiple sources:
// Destination is treated as a base directory; each source's base name is
// appended automatically, e.g. sources ["uploads/a.txt", "container/b.txt"]
// with destination "results" → "results/a.txt" and "results/b.txt".
type AddResourceFromFlowRequest struct {
// Source is a single relative path within the flow cache (kept for backward
// compatibility). Combined with Sources when both are provided.
Source string `json:"source"`
// Sources is a list of relative paths within the flow cache. Combined with
// Source when both are provided; duplicate entries are silently removed.
Sources []string `json:"sources"`
// Destination is the virtual path prefix in the user's resource tree.
Destination string `json:"destination" binding:"required"`
// Force overwrites existing resources at the target paths when true.
Force bool `json:"force"`
}
+168
View File
@@ -0,0 +1,168 @@
package models
import (
"fmt"
"time"
"pentagi/pkg/tools"
"github.com/jinzhu/gorm"
)
type FlowStatus string
const (
FlowStatusCreated FlowStatus = "created"
FlowStatusRunning FlowStatus = "running"
FlowStatusWaiting FlowStatus = "waiting"
FlowStatusFinished FlowStatus = "finished"
FlowStatusFailed FlowStatus = "failed"
)
func (s FlowStatus) String() string {
return string(s)
}
// Valid is function to control input/output data
func (s FlowStatus) Valid() error {
switch s {
case FlowStatusCreated,
FlowStatusRunning,
FlowStatusWaiting,
FlowStatusFinished,
FlowStatusFailed:
return nil
default:
return fmt.Errorf("invalid FlowStatus: %s", s)
}
}
// Validate is function to use callback to control input/output data
func (s FlowStatus) Validate(db *gorm.DB) {
if err := s.Valid(); err != nil {
db.AddError(err)
}
}
// Flow is model to contain flow information
// nolint:lll
type Flow struct {
ID uint64 `form:"id" json:"id" validate:"min=0,numeric" gorm:"type:BIGINT;NOT NULL;PRIMARY_KEY;AUTO_INCREMENT"`
Status FlowStatus `form:"status" json:"status" validate:"valid,required" gorm:"type:FLOW_STATUS;NOT NULL;default:'created'"`
Title string `form:"title" json:"title" validate:"required" gorm:"type:TEXT;NOT NULL;default:'untitled'"`
Model string `form:"model" json:"model" validate:"max=70,required" gorm:"type:TEXT;NOT NULL"`
ModelProviderName string `form:"model_provider_name" json:"model_provider_name" validate:"max=70,required" gorm:"type:TEXT;NOT NULL"`
ModelProviderType ProviderType `form:"model_provider_type" json:"model_provider_type" validate:"valid,required" gorm:"type:PROVIDER_TYPE;NOT NULL"`
Language string `form:"language" json:"language" validate:"max=70,required" gorm:"type:TEXT;NOT NULL"`
Functions *tools.Functions `form:"functions,omitempty" json:"functions,omitempty" validate:"omitempty,valid" gorm:"type:JSON;NOT NULL;default:'{}'"`
ToolCallIDTemplate string `form:"tool_call_id_template" json:"tool_call_id_template" validate:"max=70,required" gorm:"type:TEXT;NOT NULL"`
TraceID *string `form:"trace_id" json:"trace_id" validate:"max=70,required" gorm:"type:TEXT;NOT NULL"`
UserID uint64 `form:"user_id" json:"user_id" validate:"min=0,numeric,required" gorm:"type:BIGINT;NOT NULL"`
CreatedAt time.Time `form:"created_at,omitempty" json:"created_at,omitempty" validate:"omitempty" gorm:"type:TIMESTAMPTZ;default:CURRENT_TIMESTAMP"`
UpdatedAt time.Time `form:"updated_at,omitempty" json:"updated_at,omitempty" validate:"omitempty" gorm:"type:TIMESTAMPTZ;default:CURRENT_TIMESTAMP"`
DeletedAt *time.Time `form:"deleted_at,omitempty" json:"deleted_at,omitempty" validate:"omitempty" sql:"index" gorm:"type:TIMESTAMPTZ"`
}
// TableName returns the table name string to guaranty use correct table
func (f *Flow) TableName() string {
return "flows"
}
// Valid is function to control input/output data
func (f Flow) Valid() error {
return validate.Struct(f)
}
// Validate is function to use callback to control input/output data
func (f Flow) Validate(db *gorm.DB) {
if err := f.Valid(); err != nil {
db.AddError(err)
}
}
// CreateFlow is model to contain flow creation paylaod
// nolint:lll
type CreateFlow struct {
Input string `form:"input" json:"input" validate:"required" example:"user input for first task in the flow"`
Provider string `form:"provider" json:"provider" validate:"required" example:"openai"`
Functions *tools.Functions `form:"functions,omitempty" json:"functions,omitempty" validate:"omitempty,valid"`
ResourceIDs []uint64 `form:"resource_ids,omitempty" json:"resource_ids,omitempty" validate:"omitempty" swaggertype:"array,integer"`
}
// Valid is function to control input/output data
func (cf CreateFlow) Valid() error {
return validate.Struct(cf)
}
// PatchFlow is model to contain flow patching paylaod
// nolint:lll
type PatchFlow struct {
Action string `form:"action" json:"action" validate:"required,oneof=stop finish input rename" enums:"stop,finish,input,rename" default:"stop"`
Input *string `form:"input,omitempty" json:"input,omitempty" validate:"required_if=Action input" example:"user input for waiting flow"`
Provider *string `form:"provider,omitempty" json:"provider,omitempty" validate:"omitempty"`
Name *string `form:"name,omitempty" json:"name,omitempty" validate:"required_if=Action rename" example:"new flow name"`
ResourceIDs []uint64 `form:"resource_ids,omitempty" json:"resource_ids,omitempty" validate:"omitempty" swaggertype:"array,integer"`
}
// Valid is function to control input/output data
func (pf PatchFlow) Valid() error {
return validate.Struct(pf)
}
// FlowTasksSubtasks is model to contain flow, linded tasks and linked subtasks information
// nolint:lll
type FlowTasksSubtasks struct {
Tasks []TaskSubtasks `form:"tasks" json:"tasks" validate:"required" gorm:"foreignkey:FlowID;association_autoupdate:false;association_autocreate:false"`
Flow `form:"" json:""`
}
// TableName returns the table name string to guaranty use correct table
func (fts *FlowTasksSubtasks) TableName() string {
return "flows"
}
// Valid is function to control input/output data
func (fts FlowTasksSubtasks) Valid() error {
for i := range fts.Tasks {
if err := fts.Tasks[i].Valid(); err != nil {
return err
}
}
return fts.Flow.Valid()
}
// Validate is function to use callback to control input/output data
func (fts FlowTasksSubtasks) Validate(db *gorm.DB) {
if err := fts.Valid(); err != nil {
db.AddError(err)
}
}
// FlowContainers is model to contain flow and linked containers information
// nolint:lll
type FlowContainers struct {
Containers []Container `form:"containers" json:"containers" validate:"required" gorm:"foreignkey:FlowID;association_autoupdate:false;association_autocreate:false"`
Flow `form:"" json:""`
}
// TableName returns the table name string to guaranty use correct table
func (fc *FlowContainers) TableName() string {
return "flows"
}
// Valid is function to control input/output data
func (fc FlowContainers) Valid() error {
for i := range fc.Containers {
if err := fc.Containers[i].Valid(); err != nil {
return err
}
}
return fc.Flow.Valid()
}
// Validate is function to use callback to control input/output data
func (fc FlowContainers) Validate(db *gorm.DB) {
if err := fc.Valid(); err != nil {
db.AddError(err)
}
}
+711
View File
@@ -0,0 +1,711 @@
package models
import (
"testing"
"github.com/stretchr/testify/assert"
)
func TestFlowStatusValid(t *testing.T) {
t.Parallel()
tests := []struct {
name string
status FlowStatus
wantErr bool
}{
{"valid created", FlowStatusCreated, false},
{"valid running", FlowStatusRunning, false},
{"valid waiting", FlowStatusWaiting, false},
{"valid finished", FlowStatusFinished, false},
{"valid failed", FlowStatusFailed, false},
{"invalid empty", FlowStatus(""), true},
{"invalid unknown", FlowStatus("unknown"), true},
{"invalid paused", FlowStatus("paused"), true},
}
for _, tt := range tests {
tt := tt
t.Run(tt.name, func(t *testing.T) {
t.Parallel()
err := tt.status.Valid()
if tt.wantErr {
assert.Error(t, err)
assert.Contains(t, err.Error(), "invalid FlowStatus")
} else {
assert.NoError(t, err)
}
})
}
}
func TestFlowStatusString(t *testing.T) {
t.Parallel()
assert.Equal(t, "created", FlowStatusCreated.String())
assert.Equal(t, "running", FlowStatusRunning.String())
assert.Equal(t, "finished", FlowStatusFinished.String())
}
func TestFlowTableName(t *testing.T) {
t.Parallel()
f := &Flow{}
assert.Equal(t, "flows", f.TableName())
}
func TestCreateFlowValid(t *testing.T) {
t.Parallel()
t.Run("valid create flow", func(t *testing.T) {
t.Parallel()
cf := CreateFlow{Input: "scan target", Provider: "openai"}
assert.NoError(t, cf.Valid())
})
t.Run("missing input", func(t *testing.T) {
t.Parallel()
cf := CreateFlow{Input: "", Provider: "openai"}
assert.Error(t, cf.Valid())
})
t.Run("missing provider", func(t *testing.T) {
t.Parallel()
cf := CreateFlow{Input: "scan target", Provider: ""}
assert.Error(t, cf.Valid())
})
}
func TestPatchFlowValid(t *testing.T) {
t.Parallel()
t.Run("valid stop action", func(t *testing.T) {
t.Parallel()
pf := PatchFlow{Action: "stop"}
assert.NoError(t, pf.Valid())
})
t.Run("valid finish action", func(t *testing.T) {
t.Parallel()
pf := PatchFlow{Action: "finish"}
assert.NoError(t, pf.Valid())
})
t.Run("valid input action with input", func(t *testing.T) {
t.Parallel()
input := "new input"
pf := PatchFlow{Action: "input", Input: &input}
assert.NoError(t, pf.Valid())
})
t.Run("valid rename action with name", func(t *testing.T) {
t.Parallel()
name := "new name"
pf := PatchFlow{Action: "rename", Name: &name}
assert.NoError(t, pf.Valid())
})
t.Run("invalid action", func(t *testing.T) {
t.Parallel()
pf := PatchFlow{Action: "invalid"}
assert.Error(t, pf.Valid())
})
t.Run("empty action", func(t *testing.T) {
t.Parallel()
pf := PatchFlow{Action: ""}
assert.Error(t, pf.Valid())
})
t.Run("input action without input", func(t *testing.T) {
t.Parallel()
pf := PatchFlow{Action: "input"}
assert.Error(t, pf.Valid())
})
t.Run("rename action without name", func(t *testing.T) {
t.Parallel()
pf := PatchFlow{Action: "rename"}
assert.Error(t, pf.Valid())
})
}
func TestTaskStatusValid(t *testing.T) {
t.Parallel()
tests := []struct {
name string
status TaskStatus
wantErr bool
}{
{"valid created", TaskStatusCreated, false},
{"valid running", TaskStatusRunning, false},
{"valid waiting", TaskStatusWaiting, false},
{"valid finished", TaskStatusFinished, false},
{"valid failed", TaskStatusFailed, false},
{"invalid empty", TaskStatus(""), true},
{"invalid unknown", TaskStatus("cancelled"), true},
}
for _, tt := range tests {
tt := tt
t.Run(tt.name, func(t *testing.T) {
t.Parallel()
err := tt.status.Valid()
if tt.wantErr {
assert.Error(t, err)
assert.Contains(t, err.Error(), "invalid TaskStatus")
} else {
assert.NoError(t, err)
}
})
}
}
func TestTaskTableName(t *testing.T) {
t.Parallel()
tk := &Task{}
assert.Equal(t, "tasks", tk.TableName())
}
func TestSubtaskStatusValid(t *testing.T) {
t.Parallel()
tests := []struct {
name string
status SubtaskStatus
wantErr bool
}{
{"valid created", SubtaskStatusCreated, false},
{"valid running", SubtaskStatusRunning, false},
{"valid waiting", SubtaskStatusWaiting, false},
{"valid finished", SubtaskStatusFinished, false},
{"valid failed", SubtaskStatusFailed, false},
{"invalid empty", SubtaskStatus(""), true},
{"invalid unknown", SubtaskStatus("pending"), true},
}
for _, tt := range tests {
tt := tt
t.Run(tt.name, func(t *testing.T) {
t.Parallel()
err := tt.status.Valid()
if tt.wantErr {
assert.Error(t, err)
assert.Contains(t, err.Error(), "invalid SubtaskStatus")
} else {
assert.NoError(t, err)
}
})
}
}
func TestSubtaskTableName(t *testing.T) {
t.Parallel()
s := &Subtask{}
assert.Equal(t, "subtasks", s.TableName())
}
func TestContainerStatusValid(t *testing.T) {
t.Parallel()
tests := []struct {
name string
status ContainerStatus
wantErr bool
}{
{"valid starting", ContainerStatusStarting, false},
{"valid running", ContainerStatusRunning, false},
{"valid stopped", ContainerStatusStopped, false},
{"valid deleted", ContainerStatusDeleted, false},
{"valid failed", ContainerStatusFailed, false},
{"invalid empty", ContainerStatus(""), true},
{"invalid unknown", ContainerStatus("restarting"), true},
}
for _, tt := range tests {
tt := tt
t.Run(tt.name, func(t *testing.T) {
t.Parallel()
err := tt.status.Valid()
if tt.wantErr {
assert.Error(t, err)
assert.Contains(t, err.Error(), "invalid ContainerStatus")
} else {
assert.NoError(t, err)
}
})
}
}
func TestContainerTypeValid(t *testing.T) {
t.Parallel()
tests := []struct {
name string
ct ContainerType
wantErr bool
}{
{"valid primary", ContainerTypePrimary, false},
{"valid secondary", ContainerTypeSecondary, false},
{"invalid empty", ContainerType(""), true},
{"invalid tertiary", ContainerType("tertiary"), true},
}
for _, tt := range tests {
tt := tt
t.Run(tt.name, func(t *testing.T) {
t.Parallel()
err := tt.ct.Valid()
if tt.wantErr {
assert.Error(t, err)
assert.Contains(t, err.Error(), "invalid ContainerType")
} else {
assert.NoError(t, err)
}
})
}
}
func TestContainerTableName(t *testing.T) {
t.Parallel()
c := &Container{}
assert.Equal(t, "containers", c.TableName())
}
func TestRoleValid(t *testing.T) {
t.Parallel()
t.Run("valid role", func(t *testing.T) {
t.Parallel()
r := Role{Name: "admin"}
assert.NoError(t, r.Valid())
})
t.Run("empty name", func(t *testing.T) {
t.Parallel()
r := Role{Name: ""}
assert.Error(t, r.Valid())
})
}
func TestRoleTableName(t *testing.T) {
t.Parallel()
r := &Role{}
assert.Equal(t, "roles", r.TableName())
}
func TestPrivilegeValid(t *testing.T) {
t.Parallel()
t.Run("valid privilege", func(t *testing.T) {
t.Parallel()
p := Privilege{Name: "read"}
assert.NoError(t, p.Valid())
})
t.Run("empty name", func(t *testing.T) {
t.Parallel()
p := Privilege{Name: ""}
assert.Error(t, p.Valid())
})
}
func TestPrivilegeTableName(t *testing.T) {
t.Parallel()
p := &Privilege{}
assert.Equal(t, "privileges", p.TableName())
}
func TestRolePrivilegesValid(t *testing.T) {
t.Parallel()
t.Run("valid role privileges", func(t *testing.T) {
t.Parallel()
rp := RolePrivileges{
Privileges: []Privilege{{Name: "read"}, {Name: "write"}},
Role: Role{Name: "admin"},
}
assert.NoError(t, rp.Valid())
})
t.Run("invalid privilege", func(t *testing.T) {
t.Parallel()
rp := RolePrivileges{
Privileges: []Privilege{{Name: ""}, {Name: "write"}},
Role: Role{Name: "admin"},
}
assert.Error(t, rp.Valid())
})
t.Run("invalid role", func(t *testing.T) {
t.Parallel()
rp := RolePrivileges{
Privileges: []Privilege{{Name: "read"}},
Role: Role{Name: ""},
}
assert.Error(t, rp.Valid())
})
}
func TestRolePrivilegesTableName(t *testing.T) {
t.Parallel()
rp := &RolePrivileges{}
assert.Equal(t, "roles", rp.TableName())
}
func TestFlowValid(t *testing.T) {
t.Parallel()
traceID := "trace-123"
validFlow := Flow{
Status: FlowStatusCreated,
Title: "test flow",
Model: "gpt-4",
ModelProviderName: "openai",
ModelProviderType: ProviderType("openai"),
Language: "en",
ToolCallIDTemplate: "call_{id}",
TraceID: &traceID,
UserID: 1,
}
t.Run("valid flow", func(t *testing.T) {
t.Parallel()
assert.NoError(t, validFlow.Valid())
})
t.Run("missing title", func(t *testing.T) {
t.Parallel()
f := validFlow
f.Title = ""
assert.Error(t, f.Valid())
})
t.Run("invalid status", func(t *testing.T) {
t.Parallel()
f := validFlow
f.Status = FlowStatus("invalid")
assert.Error(t, f.Valid())
})
t.Run("invalid provider type", func(t *testing.T) {
t.Parallel()
f := validFlow
f.ModelProviderType = ProviderType("invalid")
assert.Error(t, f.Valid())
})
}
func TestFlowTasksSubtasksValid(t *testing.T) {
t.Parallel()
traceID := "trace-123"
validFlow := Flow{
Status: FlowStatusCreated,
Title: "test flow",
Model: "gpt-4",
ModelProviderName: "openai",
ModelProviderType: ProviderType("openai"),
Language: "en",
ToolCallIDTemplate: "call_{id}",
TraceID: &traceID,
UserID: 1,
}
validTask := Task{
Status: TaskStatusCreated,
Title: "task1",
Input: "input1",
FlowID: 1,
}
validSubtask := Subtask{
Status: SubtaskStatusCreated,
Title: "subtask1",
Description: "desc1",
TaskID: 1,
}
t.Run("valid flow tasks subtasks", func(t *testing.T) {
t.Parallel()
fts := FlowTasksSubtasks{
Tasks: []TaskSubtasks{{
Subtasks: []Subtask{validSubtask},
Task: validTask,
}},
Flow: validFlow,
}
assert.NoError(t, fts.Valid())
})
t.Run("invalid flow", func(t *testing.T) {
t.Parallel()
badFlow := validFlow
badFlow.Title = ""
fts := FlowTasksSubtasks{
Tasks: []TaskSubtasks{{
Subtasks: []Subtask{validSubtask},
Task: validTask,
}},
Flow: badFlow,
}
assert.Error(t, fts.Valid())
})
t.Run("invalid nested subtask", func(t *testing.T) {
t.Parallel()
badSubtask := validSubtask
badSubtask.Title = ""
fts := FlowTasksSubtasks{
Tasks: []TaskSubtasks{{
Subtasks: []Subtask{badSubtask},
Task: validTask,
}},
Flow: validFlow,
}
assert.Error(t, fts.Valid())
})
}
func TestFlowTasksSubtasksTableName(t *testing.T) {
t.Parallel()
fts := &FlowTasksSubtasks{}
assert.Equal(t, "flows", fts.TableName())
}
func TestFlowContainersValid(t *testing.T) {
t.Parallel()
traceID := "trace-123"
validFlow := Flow{
Status: FlowStatusCreated,
Title: "test flow",
Model: "gpt-4",
ModelProviderName: "openai",
ModelProviderType: ProviderType("openai"),
Language: "en",
ToolCallIDTemplate: "call_{id}",
TraceID: &traceID,
UserID: 1,
}
validContainer := Container{
Type: ContainerTypePrimary,
Name: "test-container",
Image: "alpine:latest",
Status: ContainerStatusRunning,
LocalID: "abc123",
LocalDir: "/tmp/test",
FlowID: 1,
}
t.Run("valid flow containers", func(t *testing.T) {
t.Parallel()
fc := FlowContainers{
Containers: []Container{validContainer},
Flow: validFlow,
}
assert.NoError(t, fc.Valid())
})
t.Run("invalid flow", func(t *testing.T) {
t.Parallel()
badFlow := validFlow
badFlow.Title = ""
fc := FlowContainers{
Containers: []Container{validContainer},
Flow: badFlow,
}
assert.Error(t, fc.Valid())
})
t.Run("invalid container", func(t *testing.T) {
t.Parallel()
badContainer := validContainer
badContainer.Name = ""
fc := FlowContainers{
Containers: []Container{badContainer},
Flow: validFlow,
}
assert.Error(t, fc.Valid())
})
}
func TestFlowContainersTableName(t *testing.T) {
t.Parallel()
fc := &FlowContainers{}
assert.Equal(t, "flows", fc.TableName())
}
func TestTaskValid(t *testing.T) {
t.Parallel()
validTask := Task{
Status: TaskStatusCreated,
Title: "test task",
Input: "test input",
FlowID: 1,
}
t.Run("valid task", func(t *testing.T) {
t.Parallel()
assert.NoError(t, validTask.Valid())
})
t.Run("missing title", func(t *testing.T) {
t.Parallel()
tk := validTask
tk.Title = ""
assert.Error(t, tk.Valid())
})
t.Run("missing input", func(t *testing.T) {
t.Parallel()
tk := validTask
tk.Input = ""
assert.Error(t, tk.Valid())
})
t.Run("invalid status", func(t *testing.T) {
t.Parallel()
tk := validTask
tk.Status = TaskStatus("invalid")
assert.Error(t, tk.Valid())
})
}
func TestTaskSubtasksValid(t *testing.T) {
t.Parallel()
validTask := Task{
Status: TaskStatusCreated,
Title: "task1",
Input: "input1",
FlowID: 1,
}
validSubtask := Subtask{
Status: SubtaskStatusCreated,
Title: "subtask1",
Description: "desc1",
TaskID: 1,
}
t.Run("valid task subtasks", func(t *testing.T) {
t.Parallel()
ts := TaskSubtasks{
Subtasks: []Subtask{validSubtask},
Task: validTask,
}
assert.NoError(t, ts.Valid())
})
t.Run("invalid task", func(t *testing.T) {
t.Parallel()
badTask := validTask
badTask.Title = ""
ts := TaskSubtasks{
Subtasks: []Subtask{validSubtask},
Task: badTask,
}
assert.Error(t, ts.Valid())
})
t.Run("invalid subtask", func(t *testing.T) {
t.Parallel()
badSubtask := validSubtask
badSubtask.Title = ""
ts := TaskSubtasks{
Subtasks: []Subtask{badSubtask},
Task: validTask,
}
assert.Error(t, ts.Valid())
})
}
func TestTaskSubtasksTableName(t *testing.T) {
t.Parallel()
ts := &TaskSubtasks{}
assert.Equal(t, "tasks", ts.TableName())
}
func TestSubtaskValid(t *testing.T) {
t.Parallel()
validSubtask := Subtask{
Status: SubtaskStatusCreated,
Title: "subtask1",
Description: "description1",
TaskID: 1,
}
t.Run("valid subtask", func(t *testing.T) {
t.Parallel()
assert.NoError(t, validSubtask.Valid())
})
t.Run("missing title", func(t *testing.T) {
t.Parallel()
s := validSubtask
s.Title = ""
assert.Error(t, s.Valid())
})
t.Run("missing description", func(t *testing.T) {
t.Parallel()
s := validSubtask
s.Description = ""
assert.Error(t, s.Valid())
})
t.Run("invalid status", func(t *testing.T) {
t.Parallel()
s := validSubtask
s.Status = SubtaskStatus("invalid")
assert.Error(t, s.Valid())
})
}
func TestContainerValid(t *testing.T) {
t.Parallel()
validContainer := Container{
Type: ContainerTypePrimary,
Name: "test-container",
Image: "alpine:latest",
Status: ContainerStatusRunning,
LocalID: "abc123",
LocalDir: "/tmp/test",
FlowID: 1,
}
t.Run("valid container", func(t *testing.T) {
t.Parallel()
assert.NoError(t, validContainer.Valid())
})
t.Run("missing name", func(t *testing.T) {
t.Parallel()
c := validContainer
c.Name = ""
assert.Error(t, c.Valid())
})
t.Run("missing image", func(t *testing.T) {
t.Parallel()
c := validContainer
c.Image = ""
assert.Error(t, c.Valid())
})
t.Run("invalid status", func(t *testing.T) {
t.Parallel()
c := validContainer
c.Status = ContainerStatus("invalid")
assert.Error(t, c.Valid())
})
t.Run("invalid type", func(t *testing.T) {
t.Parallel()
c := validContainer
c.Type = ContainerType("invalid")
assert.Error(t, c.Valid())
})
}
+231
View File
@@ -0,0 +1,231 @@
package models
import (
"encoding/json"
"fmt"
"reflect"
"regexp"
"strings"
"github.com/go-playground/validator/v10"
"github.com/xeipuuv/gojsonschema"
)
const (
solidRegexString = "^[a-z0-9_\\-]+$"
clDateRegexString = "^[0-9]{2}[.-][0-9]{2}[.-][0-9]{4}$"
semverRegexString = "^[0-9]+\\.[0-9]+(\\.[0-9]+)?$"
semverexRegexString = "^(v)?[0-9]+\\.[0-9]+(\\.[0-9]+)?(\\.[0-9]+)?(-[a-zA-Z0-9]+)?$"
)
var (
validate *validator.Validate
)
func GetValidator() *validator.Validate {
return validate
}
// IValid is interface to control all models from user code
type IValid interface {
Valid() error
}
func templateValidatorString(regexpString string) validator.Func {
regexpValue := regexp.MustCompile(regexpString)
return func(fl validator.FieldLevel) bool {
field := fl.Field()
matchString := func(str string) bool {
if str == "" && fl.Param() == "omitempty" {
return true
}
return regexpValue.MatchString(str)
}
switch field.Kind() {
case reflect.String:
return matchString(fl.Field().String())
case reflect.Slice, reflect.Array:
for i := 0; i < field.Len(); i++ {
if !matchString(field.Index(i).String()) {
return false
}
}
return true
case reflect.Map:
for _, k := range field.MapKeys() {
if !matchString(field.MapIndex(k).String()) {
return false
}
}
return true
default:
return false
}
}
}
func strongPasswordValidatorString() validator.Func {
numberRegex := regexp.MustCompile("[0-9]")
alphaLRegex := regexp.MustCompile("[a-z]")
alphaURegex := regexp.MustCompile("[A-Z]")
specRegex := regexp.MustCompile("[!@#$&*]")
return func(fl validator.FieldLevel) bool {
field := fl.Field()
switch field.Kind() {
case reflect.String:
password := fl.Field().String()
return len(password) > 15 || (len(password) >= 8 &&
numberRegex.MatchString(password) &&
alphaLRegex.MatchString(password) &&
alphaURegex.MatchString(password) &&
specRegex.MatchString(password))
default:
return false
}
}
}
func emailValidatorString() validator.Func {
emailRegex := regexp.MustCompile(`^[a-z0-9._%+\-]+@[a-z0-9.\-]+\.[a-z]{2,4}$`)
return func(fl validator.FieldLevel) bool {
field := fl.Field()
switch field.Kind() {
case reflect.String:
email := fl.Field().String()
if email == "admin" {
return true
}
if err := validate.Var(email, "required,uuid"); err == nil {
return true
}
return len(email) > 4 && emailRegex.MatchString(email)
default:
return false
}
}
}
func oauthMinScope() validator.Func {
scopeParts := []string{
"openid",
"email",
}
return func(fl validator.FieldLevel) bool {
field := fl.Field()
switch field.Kind() {
case reflect.String:
scope := strings.ToLower(fl.Field().String())
for _, part := range scopeParts {
if !strings.Contains(scope, part) {
return false
}
}
return true
default:
return false
}
}
}
func deepValidator() validator.Func {
return func(fl validator.FieldLevel) bool {
if iv, ok := fl.Field().Interface().(IValid); ok {
if err := iv.Valid(); err != nil {
return false
}
}
return true
}
}
func getMapKeys(kvmap interface{}) string {
kl := []interface{}{}
val := reflect.ValueOf(kvmap)
if val.Kind() == reflect.Map {
for _, e := range val.MapKeys() {
v := val.MapIndex(e)
kl = append(kl, v.Interface())
}
}
kld, _ := json.Marshal(kl)
return string(kld)
}
func mismatchLenError(tag string, wants, current int) string {
return fmt.Sprintf("%s wants len %d but current is %d", tag, wants, current)
}
func keyIsNotExtistInMap(tag, key string, kvmap interface{}) string {
return fmt.Sprintf("%s must present key %s in keys list %s", tag, key, getMapKeys(kvmap))
}
func keyIsNotExtistInSlice(tag, key string, klist interface{}) string {
kld, _ := json.Marshal(klist)
return fmt.Sprintf("%s must present key %s in keys list %s", tag, key, string(kld))
}
func keysAreNotExtistInSlice(tag, lkeys, rkeys interface{}) string {
lkeysd, _ := json.Marshal(lkeys)
rkeysd, _ := json.Marshal(rkeys)
return fmt.Sprintf("%s must all keys present %s in keys list %s", tag, string(lkeysd), string(rkeysd))
}
func contextError(tag string, id string, ctx interface{}) string {
ctxd, _ := json.Marshal(ctx)
return fmt.Sprintf("%s with %s ctx %s", tag, id, string(ctxd))
}
func caughtValidationError(tag string, err error) string {
return fmt.Sprintf("%s caught error %s", tag, err.Error())
}
func caughtSchemaValidationError(tag string, errs []gojsonschema.ResultError) string {
var arr []string
for _, err := range errs {
arr = append(arr, err.String())
}
errd, _ := json.Marshal(arr)
return fmt.Sprintf("%s caught errors %s", tag, string(errd))
}
func scanFromJSON(input interface{}, output interface{}) error {
if v, ok := input.(string); ok {
return json.Unmarshal([]byte(v), output)
} else if v, ok := input.([]byte); ok {
if err := json.Unmarshal(v, output); err != nil {
return err
}
return nil
}
return fmt.Errorf("unsupported type of input value to scan")
}
func init() {
validate = validator.New()
_ = validate.RegisterValidation("solid", templateValidatorString(solidRegexString))
_ = validate.RegisterValidation("cldate", templateValidatorString(clDateRegexString))
_ = validate.RegisterValidation("semver", templateValidatorString(semverRegexString))
_ = validate.RegisterValidation("semverex", templateValidatorString(semverexRegexString))
_ = validate.RegisterValidation("stpass", strongPasswordValidatorString())
_ = validate.RegisterValidation("vmail", emailValidatorString())
_ = validate.RegisterValidation("oauth_min_scope", oauthMinScope())
_ = validate.RegisterValidation("valid", deepValidator())
// Check validation interface for all models
_, _ = reflect.ValueOf(Login{}).Interface().(IValid)
_, _ = reflect.ValueOf(AuthCallback{}).Interface().(IValid)
_, _ = reflect.ValueOf(User{}).Interface().(IValid)
_, _ = reflect.ValueOf(Password{}).Interface().(IValid)
_, _ = reflect.ValueOf(Role{}).Interface().(IValid)
_, _ = reflect.ValueOf(Prompt{}).Interface().(IValid)
_, _ = reflect.ValueOf(Assistant{}).Interface().(IValid)
_, _ = reflect.ValueOf(Flow{}).Interface().(IValid)
_, _ = reflect.ValueOf(Provider{}).Interface().(IValid)
}
+250
View File
@@ -0,0 +1,250 @@
package models
import (
"testing"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
func TestGetValidator(t *testing.T) {
t.Parallel()
v := GetValidator()
require.NotNil(t, v)
}
func TestStrongPasswordValidator(t *testing.T) {
t.Parallel()
tests := []struct {
name string
pw string
wantErr bool
}{
{"long password over 15 chars", "abcdefghijklmnop", false},
{"exactly 16 chars", "abcdefghijklmnop", false},
{"8 chars with all requirements", "Pass1!ab", false},
{"short no special", "pass1", true},
{"7 chars with requirements", "Pa1!abc", true},
{"8 chars no number", "Pass!abc", true},
{"8 chars no uppercase", "pass1!ab", true},
{"8 chars no lowercase", "PASS1!AB", true},
{"8 chars no special", "Pass1abc", true},
{"empty password", "", true},
}
for _, tt := range tests {
tt := tt
t.Run(tt.name, func(t *testing.T) {
t.Parallel()
err := GetValidator().Var(tt.pw, "stpass")
if tt.wantErr {
assert.Error(t, err)
} else {
assert.NoError(t, err)
}
})
}
}
func TestEmailValidator(t *testing.T) {
t.Parallel()
tests := []struct {
name string
email string
wantErr bool
}{
{"valid email", "test@example.com", false},
{"admin special case", "admin", false},
{"email with subdomain", "user@mail.example.com", false},
{"email too short", "a@b", true},
{"empty email", "", true},
{"no at sign", "testexample.com", true},
{"no domain", "test@", true},
}
for _, tt := range tests {
tt := tt
t.Run(tt.name, func(t *testing.T) {
t.Parallel()
err := GetValidator().Var(tt.email, "vmail")
if tt.wantErr {
assert.Error(t, err)
} else {
assert.NoError(t, err)
}
})
}
}
func TestOAuthMinScopeValidator(t *testing.T) {
t.Parallel()
tests := []struct {
name string
scope string
wantErr bool
}{
{"valid openid email", "openid email", false},
{"valid with extra scopes", "openid email profile", false},
{"valid case insensitive", "OpenID Email", false},
{"missing openid", "email profile", true},
{"missing email", "openid profile", true},
{"empty scope", "", true},
}
for _, tt := range tests {
tt := tt
t.Run(tt.name, func(t *testing.T) {
t.Parallel()
err := GetValidator().Var(tt.scope, "oauth_min_scope")
if tt.wantErr {
assert.Error(t, err)
} else {
assert.NoError(t, err)
}
})
}
}
func TestSolidValidator(t *testing.T) {
t.Parallel()
tests := []struct {
name string
input string
wantErr bool
}{
{"valid lowercase", "hello", false},
{"valid with numbers", "abc123", false},
{"valid with underscore", "my_name", false},
{"valid with dash", "my-name", false},
{"invalid uppercase", "Hello", true},
{"invalid spaces", "hello world", true},
{"invalid special chars", "hello@world", true},
{"empty string", "", true},
}
for _, tt := range tests {
tt := tt
t.Run(tt.name, func(t *testing.T) {
t.Parallel()
err := GetValidator().Var(tt.input, "solid")
if tt.wantErr {
assert.Error(t, err)
} else {
assert.NoError(t, err)
}
})
}
}
func TestSemverValidator(t *testing.T) {
t.Parallel()
tests := []struct {
name string
input string
wantErr bool
}{
{"valid two part", "1.0", false},
{"valid three part", "1.2.3", false},
{"valid zero", "0.0.0", false},
{"invalid prefix v", "v1.0.0", true},
{"invalid single", "1", true},
{"invalid text", "abc", true},
{"empty", "", true},
}
for _, tt := range tests {
tt := tt
t.Run(tt.name, func(t *testing.T) {
t.Parallel()
err := GetValidator().Var(tt.input, "semver")
if tt.wantErr {
assert.Error(t, err)
} else {
assert.NoError(t, err)
}
})
}
}
func TestSemverExValidator(t *testing.T) {
t.Parallel()
tests := []struct {
name string
input string
wantErr bool
}{
{"valid two part", "1.0", false},
{"valid three part", "1.2.3", false},
{"valid with v prefix", "v1.0.0", false},
{"valid with prerelease", "1.2.3-beta", false},
{"valid four part", "1.2.3.4", false},
{"invalid text only", "abc", true},
{"empty", "", true},
}
for _, tt := range tests {
tt := tt
t.Run(tt.name, func(t *testing.T) {
t.Parallel()
err := GetValidator().Var(tt.input, "semverex")
if tt.wantErr {
assert.Error(t, err)
} else {
assert.NoError(t, err)
}
})
}
}
func TestScanFromJSON(t *testing.T) {
t.Parallel()
type testStruct struct {
Name string `json:"name"`
}
t.Run("scan from string", func(t *testing.T) {
t.Parallel()
var out testStruct
err := scanFromJSON(`{"name":"test"}`, &out)
require.NoError(t, err)
assert.Equal(t, "test", out.Name)
})
t.Run("scan from bytes", func(t *testing.T) {
t.Parallel()
var out testStruct
err := scanFromJSON([]byte(`{"name":"hello"}`), &out)
require.NoError(t, err)
assert.Equal(t, "hello", out.Name)
})
t.Run("scan unsupported type", func(t *testing.T) {
t.Parallel()
var out testStruct
err := scanFromJSON(12345, &out)
assert.Error(t, err)
assert.Contains(t, err.Error(), "unsupported type")
})
t.Run("scan invalid json string", func(t *testing.T) {
t.Parallel()
var out testStruct
err := scanFromJSON("not-json", &out)
assert.Error(t, err)
})
t.Run("scan invalid json bytes", func(t *testing.T) {
t.Parallel()
var out testStruct
err := scanFromJSON([]byte("not-json"), &out)
assert.Error(t, err)
})
}
+476
View File
@@ -0,0 +1,476 @@
package models
import (
"encoding/json"
"fmt"
gqlmodel "pentagi/pkg/graph/model"
)
// knowledgeRawMeta mirrors the flat JSON stored in langchain_pg_embedding.cmetadata.
// Used only within the models package for REST → gorm conversion.
type knowledgeRawMeta struct {
DocType string `json:"doc_type"`
UserID int64 `json:"user_id"`
FlowID *int64 `json:"flow_id"`
TaskID *int64 `json:"task_id"`
SubtaskID *int64 `json:"subtask_id"`
Question string `json:"question"`
Description string `json:"description"`
GuideType string `json:"guide_type"`
AnswerType string `json:"answer_type"`
CodeLang string `json:"code_lang"`
PartSize int `json:"part_size"`
TotalSize int `json:"total_size"`
Manual bool `json:"manual"`
}
func knowledgeMetaFromJSON(raw string) knowledgeRawMeta {
var m knowledgeRawMeta
if raw != "" && raw != "{}" {
_ = json.Unmarshal([]byte(raw), &m)
}
return m
}
func knowledgeMetaToEntry(id, content string, m knowledgeRawMeta) KnowledgeDocEntry {
entry := KnowledgeDocEntry{
ID: id,
DocType: KnowledgeDocType(m.DocType),
Content: content,
Question: m.Question,
UserID: m.UserID,
FlowID: m.FlowID,
TaskID: m.TaskID,
SubtaskID: m.SubtaskID,
PartSize: m.PartSize,
TotalSize: m.TotalSize,
Manual: m.Manual,
}
if m.Description != "" {
d := m.Description
entry.Description = &d
}
if m.GuideType != "" {
gt := KnowledgeGuideType(m.GuideType)
entry.GuideType = &gt
}
if m.AnswerType != "" {
at := KnowledgeAnswerType(m.AnswerType)
entry.AnswerType = &at
}
if m.CodeLang != "" {
cl := m.CodeLang
entry.CodeLang = &cl
}
return entry
}
// ==================== Enum types =============================================
// KnowledgeDocType mirrors model.KnowledgeDocType for the REST layer.
type KnowledgeDocType string
const (
KnowledgeDocTypeAnswer KnowledgeDocType = "answer"
KnowledgeDocTypeGuide KnowledgeDocType = "guide"
KnowledgeDocTypeCode KnowledgeDocType = "code"
)
// Valid implements IValid.
func (t KnowledgeDocType) Valid() error {
switch t {
case KnowledgeDocTypeAnswer, KnowledgeDocTypeGuide, KnowledgeDocTypeCode:
return nil
default:
return fmt.Errorf("invalid KnowledgeDocType: %q (expected answer|guide|code)", t)
}
}
// KnowledgeGuideType mirrors model.KnowledgeGuideType for the REST layer.
type KnowledgeGuideType string
const (
KnowledgeGuideTypeInstall KnowledgeGuideType = "install"
KnowledgeGuideTypeConfigure KnowledgeGuideType = "configure"
KnowledgeGuideTypeUse KnowledgeGuideType = "use"
KnowledgeGuideTypePentest KnowledgeGuideType = "pentest"
KnowledgeGuideTypeDevelopment KnowledgeGuideType = "development"
KnowledgeGuideTypeOther KnowledgeGuideType = "other"
)
// Valid implements IValid.
func (t KnowledgeGuideType) Valid() error {
switch t {
case KnowledgeGuideTypeInstall,
KnowledgeGuideTypeConfigure,
KnowledgeGuideTypeUse,
KnowledgeGuideTypePentest,
KnowledgeGuideTypeDevelopment,
KnowledgeGuideTypeOther:
return nil
default:
return fmt.Errorf("invalid KnowledgeGuideType: %q", t)
}
}
// KnowledgeAnswerType mirrors model.KnowledgeAnswerType for the REST layer.
type KnowledgeAnswerType string
const (
KnowledgeAnswerTypeGuide KnowledgeAnswerType = "guide"
KnowledgeAnswerTypeVulnerability KnowledgeAnswerType = "vulnerability"
KnowledgeAnswerTypeCode KnowledgeAnswerType = "code"
KnowledgeAnswerTypeTool KnowledgeAnswerType = "tool"
KnowledgeAnswerTypeOther KnowledgeAnswerType = "other"
)
// Valid implements IValid.
func (t KnowledgeAnswerType) Valid() error {
switch t {
case KnowledgeAnswerTypeGuide,
KnowledgeAnswerTypeVulnerability,
KnowledgeAnswerTypeCode,
KnowledgeAnswerTypeTool,
KnowledgeAnswerTypeOther:
return nil
default:
return fmt.Errorf("invalid KnowledgeAnswerType: %q", t)
}
}
// ==================== Response models =========================================
// KnowledgeDocEntry is the REST response representation of a knowledge document.
//
//nolint:lll
type KnowledgeDocEntry struct {
ID string `json:"id"`
DocType KnowledgeDocType `json:"doc_type"`
Content string `json:"content"`
Question string `json:"question"`
Description *string `json:"description,omitempty"`
UserID int64 `json:"user_id"`
FlowID *int64 `json:"flow_id,omitempty"`
TaskID *int64 `json:"task_id,omitempty"`
SubtaskID *int64 `json:"subtask_id,omitempty"`
GuideType *KnowledgeGuideType `json:"guide_type,omitempty"`
AnswerType *KnowledgeAnswerType `json:"answer_type,omitempty"`
CodeLang *string `json:"code_lang,omitempty"`
PartSize int `json:"part_size"`
TotalSize int `json:"total_size"`
Manual bool `json:"manual"`
}
// KnowledgeDocList is the REST list response (Total matches rdb.TableQuery return type).
type KnowledgeDocList struct {
Items []KnowledgeDocEntry `json:"items"`
Total uint64 `json:"total"`
}
// KnowledgeEmbeddingRow is a gorm scan target for langchain_pg_embedding rows.
// Only the columns that the knowledge API exposes are selected (no embedding vector).
type KnowledgeEmbeddingRow struct {
ID string `gorm:"column:id"`
Document string `gorm:"column:document"`
Cmetadata string `gorm:"column:cmetadata"`
}
// KnowledgeDocWithScore wraps a document with its semantic similarity score.
type KnowledgeDocWithScore struct {
Score float64 `json:"score"`
Document KnowledgeDocEntry `json:"document"`
}
// KnowledgeSearchResult is the REST search response.
type KnowledgeSearchResult struct {
Items []KnowledgeDocWithScore `json:"items"`
Total int `json:"total"`
}
// ==================== Query / request models =================================
// KnowledgeListQuery holds query-string parameters for the list endpoint.
//
//nolint:lll
type KnowledgeListQuery struct {
WithContent bool `form:"with_content" json:"with_content"`
DocTypes []KnowledgeDocType `form:"doc_types[]" json:"doc_types,omitempty"`
GuideTypes []KnowledgeGuideType `form:"guide_types[]" json:"guide_types,omitempty"`
AnswerTypes []KnowledgeAnswerType `form:"answer_types[]" json:"answer_types,omitempty"`
CodeLangs []string `form:"code_langs[]" json:"code_langs,omitempty"`
FlowID *int64 `form:"flow_id" json:"flow_id,omitempty"`
Manual *bool `form:"manual" json:"manual,omitempty"`
}
// Valid validates all enum values supplied in the list query.
func (q KnowledgeListQuery) Valid() error {
for _, dt := range q.DocTypes {
if err := dt.Valid(); err != nil {
return err
}
}
for _, gt := range q.GuideTypes {
if err := gt.Valid(); err != nil {
return err
}
}
for _, at := range q.AnswerTypes {
if err := at.Valid(); err != nil {
return err
}
}
return nil
}
// CreateKnowledgeDocRequest is the POST body for creating a knowledge document.
//
//nolint:lll
type CreateKnowledgeDocRequest struct {
DocType KnowledgeDocType `json:"doc_type" validate:"required,valid"`
Content string `json:"content" validate:"required,min=1,max=65536"`
Question string `json:"question" validate:"required,min=1,max=2048"`
Description *string `json:"description,omitempty" validate:"omitempty,max=1000"`
GuideType *KnowledgeGuideType `json:"guide_type,omitempty" validate:"omitempty,valid"`
AnswerType *KnowledgeAnswerType `json:"answer_type,omitempty" validate:"omitempty,valid"`
CodeLang *string `json:"code_lang,omitempty" validate:"omitempty,max=100"`
}
// Valid implements IValid.
func (r CreateKnowledgeDocRequest) Valid() error {
if err := r.DocType.Valid(); err != nil {
return err
}
if r.GuideType != nil {
if err := r.GuideType.Valid(); err != nil {
return err
}
}
if r.AnswerType != nil {
if err := r.AnswerType.Valid(); err != nil {
return err
}
}
return validate.Struct(r)
}
// UpdateKnowledgeDocRequest is the PUT body for updating a knowledge document.
// content is mandatory (triggers re-embedding); all other fields are optional overrides.
//
//nolint:lll
type UpdateKnowledgeDocRequest struct {
Content string `json:"content" validate:"required,min=1,max=65536"`
Question *string `json:"question,omitempty" validate:"omitempty,min=1,max=2048"`
Description *string `json:"description,omitempty" validate:"omitempty,max=1000"`
GuideType *KnowledgeGuideType `json:"guide_type,omitempty" validate:"omitempty,valid"`
AnswerType *KnowledgeAnswerType `json:"answer_type,omitempty" validate:"omitempty,valid"`
CodeLang *string `json:"code_lang,omitempty" validate:"omitempty,max=100"`
}
// Valid implements IValid.
func (r UpdateKnowledgeDocRequest) Valid() error {
if r.GuideType != nil {
if err := r.GuideType.Valid(); err != nil {
return err
}
}
if r.AnswerType != nil {
if err := r.AnswerType.Valid(); err != nil {
return err
}
}
return validate.Struct(r)
}
// KnowledgeSearchRequest is the POST body for semantic search.
//
//nolint:lll
type KnowledgeSearchRequest struct {
Query string `json:"query" validate:"required,min=1,max=2048"`
Limit int `json:"limit,omitempty" validate:"omitempty,min=1,max=100"`
DocTypes []KnowledgeDocType `json:"doc_types,omitempty" validate:"omitempty,dive,valid"`
GuideTypes []KnowledgeGuideType `json:"guide_types,omitempty" validate:"omitempty,dive,valid"`
AnswerTypes []KnowledgeAnswerType `json:"answer_types,omitempty" validate:"omitempty,dive,valid"`
CodeLangs []string `json:"code_langs,omitempty" validate:"omitempty,dive,max=100"`
FlowID *int64 `json:"flow_id,omitempty"`
Manual *bool `json:"manual,omitempty"`
}
// Valid implements IValid.
func (r KnowledgeSearchRequest) Valid() error {
for _, dt := range r.DocTypes {
if err := dt.Valid(); err != nil {
return err
}
}
for _, gt := range r.GuideTypes {
if err := gt.Valid(); err != nil {
return err
}
}
for _, at := range r.AnswerTypes {
if err := at.Valid(); err != nil {
return err
}
}
return validate.Struct(r)
}
// ==================== Converters (REST ↔ GraphQL model) ======================
// KnowledgeDocFromGQL converts a GraphQL KnowledgeDocument to a REST KnowledgeDocEntry.
func KnowledgeDocFromGQL(doc *gqlmodel.KnowledgeDocument) KnowledgeDocEntry {
entry := KnowledgeDocEntry{
ID: doc.ID,
DocType: KnowledgeDocType(doc.DocType),
Content: doc.Content,
Question: doc.Question,
UserID: doc.UserID,
FlowID: doc.FlowID,
TaskID: doc.TaskID,
SubtaskID: doc.SubtaskID,
PartSize: doc.PartSize,
TotalSize: doc.TotalSize,
Manual: doc.Manual,
}
if doc.Description != nil {
entry.Description = doc.Description
}
if doc.GuideType != nil {
gt := KnowledgeGuideType(*doc.GuideType)
entry.GuideType = &gt
}
if doc.AnswerType != nil {
at := KnowledgeAnswerType(*doc.AnswerType)
entry.AnswerType = &at
}
if doc.CodeLang != nil {
entry.CodeLang = doc.CodeLang
}
return entry
}
// KnowledgeDocListFromGQL converts a slice of GraphQL documents to a REST list response.
func KnowledgeDocListFromGQL(docs []*gqlmodel.KnowledgeDocument) KnowledgeDocList {
items := make([]KnowledgeDocEntry, 0, len(docs))
for _, d := range docs {
items = append(items, KnowledgeDocFromGQL(d))
}
return KnowledgeDocList{Items: items, Total: uint64(len(items))}
}
// KnowledgeDocListFromRows converts raw gorm scan rows to a REST list response.
func KnowledgeDocListFromRows(rows []KnowledgeEmbeddingRow, total uint64, withContent bool) KnowledgeDocList {
items := make([]KnowledgeDocEntry, 0, len(rows))
for _, r := range rows {
items = append(items, KnowledgeDocEntryFromRow(r, withContent))
}
return KnowledgeDocList{Items: items, Total: total}
}
// KnowledgeDocEntryFromRow converts a single gorm scan row to a REST KnowledgeDocEntry.
func KnowledgeDocEntryFromRow(r KnowledgeEmbeddingRow, withContent bool) KnowledgeDocEntry {
meta := knowledgeMetaFromJSON(r.Cmetadata)
content := r.Document
if !withContent {
content = ""
}
return knowledgeMetaToEntry(r.ID, content, meta)
}
// KnowledgeSearchResultFromGQL converts GraphQL search results to the REST response.
func KnowledgeSearchResultFromGQL(results []*gqlmodel.KnowledgeDocumentWithScore) KnowledgeSearchResult {
items := make([]KnowledgeDocWithScore, 0, len(results))
for _, r := range results {
items = append(items, KnowledgeDocWithScore{
Score: r.Score,
Document: KnowledgeDocFromGQL(r.Document),
})
}
return KnowledgeSearchResult{Items: items, Total: len(items)}
}
// CreateRequestToGQL converts a REST create request to the GraphQL input type.
func (r CreateKnowledgeDocRequest) ToGQL() gqlmodel.CreateKnowledgeDocumentInput {
input := gqlmodel.CreateKnowledgeDocumentInput{
DocType: gqlmodel.KnowledgeDocType(r.DocType),
Content: r.Content,
Question: r.Question,
Description: r.Description,
CodeLang: r.CodeLang,
}
if r.GuideType != nil {
gt := gqlmodel.KnowledgeGuideType(*r.GuideType)
input.GuideType = &gt
}
if r.AnswerType != nil {
at := gqlmodel.KnowledgeAnswerType(*r.AnswerType)
input.AnswerType = &at
}
return input
}
// UpdateRequestToGQL converts a REST update request to the GraphQL input type.
func (r UpdateKnowledgeDocRequest) ToGQL() gqlmodel.UpdateKnowledgeDocumentInput {
input := gqlmodel.UpdateKnowledgeDocumentInput{
Content: r.Content,
Question: r.Question,
Description: r.Description,
CodeLang: r.CodeLang,
}
if r.GuideType != nil {
gt := gqlmodel.KnowledgeGuideType(*r.GuideType)
input.GuideType = &gt
}
if r.AnswerType != nil {
at := gqlmodel.KnowledgeAnswerType(*r.AnswerType)
input.AnswerType = &at
}
return input
}
// ListQueryToGQLFilter converts a REST list query to a GraphQL KnowledgeFilter.
func (q KnowledgeListQuery) ToGQLFilter() *gqlmodel.KnowledgeFilter {
f := &gqlmodel.KnowledgeFilter{
FlowID: q.FlowID,
Manual: q.Manual,
}
for _, dt := range q.DocTypes {
f.DocTypes = append(f.DocTypes, gqlmodel.KnowledgeDocType(dt))
}
for _, gt := range q.GuideTypes {
f.GuideTypes = append(f.GuideTypes, gqlmodel.KnowledgeGuideType(gt))
}
for _, at := range q.AnswerTypes {
f.AnswerTypes = append(f.AnswerTypes, gqlmodel.KnowledgeAnswerType(at))
}
f.CodeLangs = append(f.CodeLangs, q.CodeLangs...)
if f.FlowID == nil && f.Manual == nil && len(f.DocTypes) == 0 &&
len(f.GuideTypes) == 0 && len(f.AnswerTypes) == 0 && len(f.CodeLangs) == 0 {
return nil
}
return f
}
// SearchRequestToGQLFilter converts a REST search request to a GraphQL KnowledgeFilter.
func (r KnowledgeSearchRequest) ToGQLFilter() *gqlmodel.KnowledgeFilter {
f := &gqlmodel.KnowledgeFilter{
FlowID: r.FlowID,
Manual: r.Manual,
}
for _, dt := range r.DocTypes {
f.DocTypes = append(f.DocTypes, gqlmodel.KnowledgeDocType(dt))
}
for _, gt := range r.GuideTypes {
f.GuideTypes = append(f.GuideTypes, gqlmodel.KnowledgeGuideType(gt))
}
for _, at := range r.AnswerTypes {
f.AnswerTypes = append(f.AnswerTypes, gqlmodel.KnowledgeAnswerType(at))
}
f.CodeLangs = append(f.CodeLangs, r.CodeLangs...)
if f.FlowID == nil && f.Manual == nil && len(f.DocTypes) == 0 &&
len(f.GuideTypes) == 0 && len(f.AnswerTypes) == 0 && len(f.CodeLangs) == 0 {
return nil
}
return f
}
+67
View File
@@ -0,0 +1,67 @@
package models
import (
"fmt"
"github.com/jinzhu/gorm"
)
type MsgchainType string
const (
MsgchainTypePrimaryAgent MsgchainType = "primary_agent"
MsgchainTypeReporter MsgchainType = "reporter"
MsgchainTypeGenerator MsgchainType = "generator"
MsgchainTypeRefiner MsgchainType = "refiner"
MsgchainTypeReflector MsgchainType = "reflector"
MsgchainTypeEnricher MsgchainType = "enricher"
MsgchainTypeAdviser MsgchainType = "adviser"
MsgchainTypeCoder MsgchainType = "coder"
MsgchainTypeMemorist MsgchainType = "memorist"
MsgchainTypeSearcher MsgchainType = "searcher"
MsgchainTypeInstaller MsgchainType = "installer"
MsgchainTypePentester MsgchainType = "pentester"
MsgchainTypeSummarizer MsgchainType = "summarizer"
MsgchainTypeToolCallFixer MsgchainType = "tool_call_fixer"
MsgchainTypeAssistant MsgchainType = "assistant"
)
func (e *MsgchainType) Scan(src interface{}) error {
switch s := src.(type) {
case []byte:
*e = MsgchainType(s)
case string:
*e = MsgchainType(s)
default:
return fmt.Errorf("unsupported scan type for MsgchainType: %T", src)
}
return nil
}
func (s MsgchainType) String() string {
return string(s)
}
// Valid is function to control input/output data
func (ml MsgchainType) Valid() error {
switch ml {
case MsgchainTypePrimaryAgent, MsgchainTypeReporter,
MsgchainTypeGenerator, MsgchainTypeRefiner,
MsgchainTypeReflector, MsgchainTypeEnricher,
MsgchainTypeAdviser, MsgchainTypeCoder,
MsgchainTypeMemorist, MsgchainTypeSearcher,
MsgchainTypeInstaller, MsgchainTypePentester,
MsgchainTypeSummarizer, MsgchainTypeToolCallFixer,
MsgchainTypeAssistant:
return nil
default:
return fmt.Errorf("invalid MsgchainType: %s", ml)
}
}
// Validate is function to use callback to control input/output data
func (ml MsgchainType) Validate(db *gorm.DB) {
if err := ml.Valid(); err != nil {
db.AddError(err)
}
}
+111
View File
@@ -0,0 +1,111 @@
package models
import (
"fmt"
"time"
"github.com/jinzhu/gorm"
)
type MsglogType string
const (
MsglogTypeAnswer MsglogType = "answer"
MsglogTypeReport MsglogType = "report"
MsglogTypeThoughts MsglogType = "thoughts"
MsglogTypeBrowser MsglogType = "browser"
MsglogTypeTerminal MsglogType = "terminal"
MsglogTypeFile MsglogType = "file"
MsglogTypeSearch MsglogType = "search"
MsglogTypeAdvice MsglogType = "advice"
MsglogTypeAsk MsglogType = "ask"
MsglogTypeInput MsglogType = "input"
MsglogTypeDone MsglogType = "done"
)
func (s MsglogType) String() string {
return string(s)
}
// Valid is function to control input/output data
func (s MsglogType) Valid() error {
switch s {
case MsglogTypeAnswer, MsglogTypeReport, MsglogTypeThoughts,
MsglogTypeBrowser, MsglogTypeTerminal, MsglogTypeFile,
MsglogTypeSearch, MsglogTypeAdvice, MsglogTypeAsk,
MsglogTypeInput, MsglogTypeDone:
return nil
default:
return fmt.Errorf("invalid MsglogType: %s", s)
}
}
// Validate is function to use callback to control input/output data
func (s MsglogType) Validate(db *gorm.DB) {
if err := s.Valid(); err != nil {
db.AddError(err)
}
}
type MsglogResultFormat string
const (
MsglogResultFormatPlain MsglogResultFormat = "plain"
MsglogResultFormatMarkdown MsglogResultFormat = "markdown"
MsglogResultFormatTerminal MsglogResultFormat = "terminal"
)
func (s MsglogResultFormat) String() string {
return string(s)
}
// Valid is function to control input/output data
func (s MsglogResultFormat) Valid() error {
switch s {
case MsglogResultFormatPlain,
MsglogResultFormatMarkdown,
MsglogResultFormatTerminal:
return nil
default:
return fmt.Errorf("invalid MsglogResultFormat: %s", s)
}
}
// Validate is function to use callback to control input/output data
func (s MsglogResultFormat) Validate(db *gorm.DB) {
if err := s.Valid(); err != nil {
db.AddError(err)
}
}
// Msglog is model to contain log record information from agents about their actions
// nolint:lll
type Msglog struct {
ID uint64 `form:"id" json:"id" validate:"min=0,numeric" gorm:"type:BIGINT;NOT NULL;PRIMARY_KEY;AUTO_INCREMENT"`
Type MsglogType `form:"type" json:"type" validate:"valid,required" gorm:"type:MSGLOG_TYPE;NOT NULL"`
Message string `form:"message" json:"message" validate:"required" gorm:"type:TEXT;NOT NULL"`
Thinking string `form:"thinking" json:"thinking" validate:"omitempty" gorm:"type:TEXT;NULL"`
Result string `form:"result" json:"result" validate:"omitempty" gorm:"type:TEXT;NOT NULL;default:''"`
ResultFormat MsglogResultFormat `form:"result_format" json:"result_format" validate:"valid,required" gorm:"type:MSGLOG_RESULT_FORMAT;NOT NULL;default:plain"`
FlowID uint64 `form:"flow_id" json:"flow_id" validate:"min=0,numeric,required" gorm:"type:BIGINT;NOT NULL"`
TaskID *uint64 `form:"task_id,omitempty" json:"task_id,omitempty" validate:"omitnil,min=0" gorm:"type:BIGINT;NOT NULL"`
SubtaskID *uint64 `form:"subtask_id,omitempty" json:"subtask_id,omitempty" validate:"omitnil,min=0" gorm:"type:BIGINT;NOT NULL"`
CreatedAt time.Time `form:"created_at,omitempty" json:"created_at,omitempty" validate:"omitempty" gorm:"type:TIMESTAMPTZ;default:CURRENT_TIMESTAMP"`
}
// TableName returns the table name string to guaranty use correct table
func (ml *Msglog) TableName() string {
return "msglogs"
}
// Valid is function to control input/output data
func (ml Msglog) Valid() error {
return validate.Struct(ml)
}
// Validate is function to use callback to control input/output data
func (ml Msglog) Validate(db *gorm.DB) {
if err := ml.Valid(); err != nil {
db.AddError(err)
}
}
+92
View File
@@ -0,0 +1,92 @@
package models
import (
"fmt"
"time"
"pentagi/pkg/templates"
"github.com/jinzhu/gorm"
)
// PromptType is an alias for templates.PromptType with validation methods for GORM
type PromptType templates.PromptType
// String returns the string representation of PromptType
func (s PromptType) String() string {
return string(s)
}
// Valid is function to control input/output data
func (s PromptType) Valid() error {
// Convert to templates.PromptType and validate against known constants
templateType := templates.PromptType(s)
switch templateType {
case templates.PromptTypePrimaryAgent, templates.PromptTypeAssistant,
templates.PromptTypePentester, templates.PromptTypeQuestionPentester,
templates.PromptTypeCoder, templates.PromptTypeQuestionCoder,
templates.PromptTypeInstaller, templates.PromptTypeQuestionInstaller,
templates.PromptTypeSearcher, templates.PromptTypeQuestionSearcher,
templates.PromptTypeMemorist, templates.PromptTypeQuestionMemorist,
templates.PromptTypeAdviser, templates.PromptTypeQuestionAdviser,
templates.PromptTypeGenerator, templates.PromptTypeSubtasksGenerator,
templates.PromptTypeRefiner, templates.PromptTypeSubtasksRefiner,
templates.PromptTypeReporter, templates.PromptTypeTaskReporter,
templates.PromptTypeReflector, templates.PromptTypeQuestionReflector,
templates.PromptTypeEnricher, templates.PromptTypeQuestionEnricher,
templates.PromptTypeToolCallFixer, templates.PromptTypeInputToolCallFixer,
templates.PromptTypeSummarizer, templates.PromptTypeImageChooser,
templates.PromptTypeLanguageChooser, templates.PromptTypeFlowDescriptor,
templates.PromptTypeTaskDescriptor, templates.PromptTypeExecutionLogs,
templates.PromptTypeFullExecutionContext, templates.PromptTypeShortExecutionContext,
templates.PromptTypeToolCallIDCollector, templates.PromptTypeToolCallIDDetector:
return nil
default:
return fmt.Errorf("invalid PromptType: %s", s)
}
}
// Validate is function to use callback to control input/output data
func (s PromptType) Validate(db *gorm.DB) {
if err := s.Valid(); err != nil {
db.AddError(err)
}
}
// Prompt is model to contain prompt information
// nolint:lll
type Prompt struct {
ID uint64 `form:"id" json:"id" validate:"min=0,numeric" gorm:"type:BIGINT;NOT NULL;PRIMARY_KEY;AUTO_INCREMENT"`
Type PromptType `form:"type" json:"type" validate:"valid,required" gorm:"type:PROMPT_TYPE;NOT NULL"`
UserID uint64 `form:"user_id" json:"user_id" validate:"min=0,numeric" gorm:"type:BIGINT;NOT NULL"`
Prompt string `form:"prompt" json:"prompt" validate:"required" gorm:"type:TEXT;NOT NULL"`
CreatedAt time.Time `form:"created_at,omitempty" json:"created_at,omitempty" validate:"omitempty" gorm:"type:TIMESTAMPTZ;default:CURRENT_TIMESTAMP"`
UpdatedAt time.Time `form:"updated_at,omitempty" json:"updated_at,omitempty" validate:"omitempty" gorm:"type:TIMESTAMPTZ;default:CURRENT_TIMESTAMP"`
}
// TableName returns the table name string to guaranty use correct table
func (p *Prompt) TableName() string {
return "prompts"
}
// Valid is function to control input/output data
func (p Prompt) Valid() error {
return validate.Struct(p)
}
// Validate is function to use callback to control input/output data
func (p Prompt) Validate(db *gorm.DB) {
if err := p.Valid(); err != nil {
db.AddError(err)
}
}
// PatchPrompt is model to contain prompt patching paylaod
type PatchPrompt struct {
Prompt string `form:"prompt" json:"prompt" validate:"required"`
}
// Valid is function to control input/output data
func (pp PatchPrompt) Valid() error {
return validate.Struct(pp)
}
+111
View File
@@ -0,0 +1,111 @@
package models
import (
"testing"
"pentagi/pkg/templates"
"github.com/stretchr/testify/assert"
)
func TestPromptTypeValid(t *testing.T) {
t.Parallel()
validTypes := []struct {
name string
pt PromptType
}{
{"primary_agent", PromptType(templates.PromptTypePrimaryAgent)},
{"assistant", PromptType(templates.PromptTypeAssistant)},
{"pentester", PromptType(templates.PromptTypePentester)},
{"coder", PromptType(templates.PromptTypeCoder)},
{"searcher", PromptType(templates.PromptTypeSearcher)},
{"summarizer", PromptType(templates.PromptTypeSummarizer)},
{"reporter", PromptType(templates.PromptTypeReporter)},
{"tool_call_id_detector", PromptType(templates.PromptTypeToolCallIDDetector)},
}
for _, tt := range validTypes {
tt := tt
t.Run("valid_"+tt.name, func(t *testing.T) {
t.Parallel()
assert.NoError(t, tt.pt.Valid())
})
}
invalidTypes := []struct {
name string
pt PromptType
}{
{"empty", PromptType("")},
{"unknown", PromptType("unknown")},
{"typo", PromptType("primarry_agent")},
}
for _, tt := range invalidTypes {
tt := tt
t.Run("invalid_"+tt.name, func(t *testing.T) {
t.Parallel()
err := tt.pt.Valid()
assert.Error(t, err)
assert.Contains(t, err.Error(), "invalid PromptType")
})
}
}
func TestPromptTypeString(t *testing.T) {
t.Parallel()
assert.Equal(t, "primary_agent", PromptType(templates.PromptTypePrimaryAgent).String())
assert.Equal(t, "assistant", PromptType(templates.PromptTypeAssistant).String())
}
func TestPromptValid(t *testing.T) {
t.Parallel()
validPrompt := Prompt{
Type: PromptType(templates.PromptTypePrimaryAgent),
Prompt: "You are a security assistant.",
}
t.Run("valid prompt", func(t *testing.T) {
t.Parallel()
assert.NoError(t, validPrompt.Valid())
})
t.Run("invalid type", func(t *testing.T) {
t.Parallel()
p := validPrompt
p.Type = PromptType("invalid")
assert.Error(t, p.Valid())
})
t.Run("missing prompt text", func(t *testing.T) {
t.Parallel()
p := validPrompt
p.Prompt = ""
assert.Error(t, p.Valid())
})
}
func TestPromptTableName(t *testing.T) {
t.Parallel()
p := &Prompt{}
assert.Equal(t, "prompts", p.TableName())
}
func TestPatchPromptValid(t *testing.T) {
t.Parallel()
t.Run("valid patch", func(t *testing.T) {
t.Parallel()
pp := PatchPrompt{Prompt: "updated prompt text"}
assert.NoError(t, pp.Valid())
})
t.Run("missing prompt", func(t *testing.T) {
t.Parallel()
pp := PatchPrompt{Prompt: ""}
assert.Error(t, pp.Valid())
})
}
+133
View File
@@ -0,0 +1,133 @@
package models
import (
"encoding/json"
"fmt"
"time"
"pentagi/pkg/providers/provider"
"github.com/jinzhu/gorm"
)
type ProviderType provider.ProviderType
func (s ProviderType) String() string {
return string(s)
}
// Valid is function to control input/output data
func (s ProviderType) Valid() error {
providerType := provider.ProviderType(s)
switch providerType {
case provider.ProviderOpenAI,
provider.ProviderAnthropic,
provider.ProviderGemini,
provider.ProviderBedrock,
provider.ProviderOllama,
provider.ProviderCustom,
provider.ProviderDeepSeek,
provider.ProviderGLM,
provider.ProviderKimi,
provider.ProviderQwen:
return nil
default:
return fmt.Errorf("invalid ProviderType: %s", s)
}
}
// Validate is function to use callback to control input/output data
func (s ProviderType) Validate(db *gorm.DB) {
if err := s.Valid(); err != nil {
db.AddError(err)
}
}
// Provider is model to contain provider configuration information
// nolint:lll
type Provider struct {
ID uint64 `form:"id" json:"id" validate:"min=0,numeric" gorm:"type:BIGINT;NOT NULL;PRIMARY_KEY;AUTO_INCREMENT"`
UserID uint64 `form:"user_id" json:"user_id" validate:"min=0,numeric" gorm:"type:BIGINT;NOT NULL"`
Type ProviderType `form:"type" json:"type" validate:"valid,required" gorm:"type:PROVIDER_TYPE;NOT NULL"`
Name string `form:"name" json:"name" validate:"required" gorm:"type:TEXT;NOT NULL"`
Config json.RawMessage `form:"config" json:"config" validate:"required" gorm:"type:JSON;NOT NULL"`
CreatedAt time.Time `form:"created_at,omitempty" json:"created_at,omitempty" validate:"omitempty" gorm:"type:TIMESTAMPTZ;default:CURRENT_TIMESTAMP"`
UpdatedAt time.Time `form:"updated_at,omitempty" json:"updated_at,omitempty" validate:"omitempty" gorm:"type:TIMESTAMPTZ;default:CURRENT_TIMESTAMP"`
DeletedAt *time.Time `form:"deleted_at,omitempty" json:"deleted_at,omitempty" validate:"omitempty" sql:"index" gorm:"type:TIMESTAMPTZ"`
}
// TableName returns the table name string to guaranty use correct table
func (p *Provider) TableName() string {
return "providers"
}
// Valid is function to control input/output data
func (p Provider) Valid() error {
return validate.Struct(p)
}
// Validate is function to use callback to control input/output data
func (p Provider) Validate(db *gorm.DB) {
if err := p.Valid(); err != nil {
db.AddError(err)
}
}
// CreateProvider is model to contain provider creation payload
// nolint:lll
type CreateProvider struct {
Config json.RawMessage `form:"config" json:"config" validate:"required" example:"{}"`
}
// Valid is function to control input/output data
func (cp CreateProvider) Valid() error {
return validate.Struct(cp)
}
// PatchProvider is model to contain provider patching payload
// nolint:lll
type PatchProvider struct {
Name *string `form:"name,omitempty" json:"name,omitempty" validate:"omitempty" example:"updated provider name"`
Config *json.RawMessage `form:"config,omitempty" json:"config,omitempty" validate:"omitempty" example:"{}"`
}
// Valid is function to control input/output data
func (pp PatchProvider) Valid() error {
return validate.Struct(pp)
}
// ModelPriceInfo is model to contain price information for a model
// nolint:lll
type ModelPriceInfo struct {
Input float64 `form:"input" json:"input" validate:"omitempty,numeric,min=0" example:"1.1"`
Output float64 `form:"output" json:"output" validate:"omitempty,numeric,min=0" example:"3.0"`
CacheRead float64 `form:"cache_read" json:"cache_read" validate:"omitempty,numeric,min=0" example:"0.1"`
CacheWrite float64 `form:"cache_write" json:"cache_write" validate:"omitempty,numeric,min=0" example:"0.3"`
}
// Valid is function to control input/output data
func (mpi ModelPriceInfo) Valid() error {
return validate.Struct(mpi)
}
// ModelInfo is model to contain model short information for display
// nolint:lll
type ModelInfo struct {
Name string `form:"name" json:"name" validate:"required" example:"gpt-4o"`
AgentTypes []string `form:"agent_types" json:"agent_types" validate:"required"`
PriceInfo *ModelPriceInfo `form:"price_info" json:"price_info" validate:"omitempty,valid"`
}
// ProviderInfo is model to contain provider short information for display
// nolint:lll
type ProviderInfo struct {
Name string `form:"name" json:"name" validate:"required" example:"my openai provider"`
Type ProviderType `form:"type" json:"type" validate:"valid,required" example:"openai"`
DefaultModel string `form:"default_model" json:"default_model" validate:"required" example:"gpt-4o"`
Models []ModelInfo `form:"models" json:"models" validate:"required"`
}
// Valid is function to control input/output data
func (p ProviderInfo) Valid() error {
return validate.Struct(p)
}
+183
View File
@@ -0,0 +1,183 @@
package models
import (
"encoding/json"
"testing"
"github.com/stretchr/testify/assert"
)
func TestProviderTypeValid(t *testing.T) {
t.Parallel()
validTypes := []struct {
name string
pt ProviderType
}{
{"openai", ProviderType("openai")},
{"anthropic", ProviderType("anthropic")},
{"gemini", ProviderType("gemini")},
{"bedrock", ProviderType("bedrock")},
{"ollama", ProviderType("ollama")},
{"custom", ProviderType("custom")},
{"deepseek", ProviderType("deepseek")},
{"glm", ProviderType("glm")},
{"kimi", ProviderType("kimi")},
{"qwen", ProviderType("qwen")},
}
for _, tt := range validTypes {
tt := tt
t.Run("valid_"+tt.name, func(t *testing.T) {
t.Parallel()
assert.NoError(t, tt.pt.Valid())
})
}
invalidTypes := []struct {
name string
pt ProviderType
}{
{"empty", ProviderType("")},
{"unknown", ProviderType("unknown")},
{"gpt4", ProviderType("gpt4")},
{"azure", ProviderType("azure")},
}
for _, tt := range invalidTypes {
tt := tt
t.Run("invalid_"+tt.name, func(t *testing.T) {
t.Parallel()
err := tt.pt.Valid()
assert.Error(t, err)
assert.Contains(t, err.Error(), "invalid ProviderType")
})
}
}
func TestProviderTypeString(t *testing.T) {
t.Parallel()
assert.Equal(t, "openai", ProviderType("openai").String())
assert.Equal(t, "anthropic", ProviderType("anthropic").String())
}
func TestProviderValid(t *testing.T) {
t.Parallel()
validProvider := Provider{
UserID: 1,
Type: ProviderType("openai"),
Name: "my-openai",
Config: json.RawMessage(`{"api_key":"test"}`),
}
t.Run("valid provider", func(t *testing.T) {
t.Parallel()
assert.NoError(t, validProvider.Valid())
})
t.Run("missing name", func(t *testing.T) {
t.Parallel()
p := validProvider
p.Name = ""
assert.Error(t, p.Valid())
})
t.Run("missing config", func(t *testing.T) {
t.Parallel()
p := validProvider
p.Config = nil
assert.Error(t, p.Valid())
})
t.Run("invalid type", func(t *testing.T) {
t.Parallel()
p := validProvider
p.Type = ProviderType("invalid")
assert.Error(t, p.Valid())
})
}
func TestProviderTableName(t *testing.T) {
t.Parallel()
p := &Provider{}
assert.Equal(t, "providers", p.TableName())
}
func TestCreateProviderValid(t *testing.T) {
t.Parallel()
t.Run("valid create provider", func(t *testing.T) {
t.Parallel()
cp := CreateProvider{Config: json.RawMessage(`{"key":"val"}`)}
assert.NoError(t, cp.Valid())
})
t.Run("missing config", func(t *testing.T) {
t.Parallel()
cp := CreateProvider{Config: nil}
assert.Error(t, cp.Valid())
})
}
func TestPatchProviderValid(t *testing.T) {
t.Parallel()
t.Run("valid patch with name", func(t *testing.T) {
t.Parallel()
name := "updated"
pp := PatchProvider{Name: &name}
assert.NoError(t, pp.Valid())
})
t.Run("valid empty patch", func(t *testing.T) {
t.Parallel()
pp := PatchProvider{}
assert.NoError(t, pp.Valid())
})
}
func TestProviderInfoValid(t *testing.T) {
t.Parallel()
validInfo := ProviderInfo{
Name: "my-provider",
Type: ProviderType("openai"),
DefaultModel: "gpt-4o",
Models: []ModelInfo{{Name: "gpt-4o", AgentTypes: []string{"primary_agent"}}},
}
t.Run("valid provider info", func(t *testing.T) {
t.Parallel()
assert.NoError(t, validInfo.Valid())
})
t.Run("missing name", func(t *testing.T) {
t.Parallel()
pi := validInfo
pi.Name = ""
assert.Error(t, pi.Valid())
})
t.Run("invalid type", func(t *testing.T) {
t.Parallel()
pi := validInfo
pi.Type = ProviderType("invalid")
assert.Error(t, pi.Valid())
})
t.Run("missing default model", func(t *testing.T) {
t.Parallel()
pi := validInfo
pi.DefaultModel = ""
assert.Error(t, pi.Valid())
})
t.Run("missing models", func(t *testing.T) {
t.Parallel()
pi := validInfo
pi.Models = nil
assert.Error(t, pi.Valid())
})
}
+91
View File
@@ -0,0 +1,91 @@
package models
import "time"
// UserResource is the GORM model for a user-owned resource (file or directory)
// whose metadata lives in PostgreSQL and whose content (for files) is stored as
// a .blob file on disk keyed by MD5 hash.
//
// nolint:lll
type UserResource struct {
ID uint64 `form:"id" json:"id" gorm:"type:BIGINT;NOT NULL;PRIMARY_KEY;AUTO_INCREMENT"`
UserID uint64 `form:"user_id" json:"user_id" gorm:"type:BIGINT;NOT NULL;index"`
Hash string `form:"hash" json:"hash" gorm:"type:TEXT;NOT NULL;default:''"`
Name string `form:"name" json:"name" gorm:"type:TEXT;NOT NULL"`
Path string `form:"path" json:"path" gorm:"type:TEXT;NOT NULL"`
Size int64 `form:"size" json:"size" gorm:"type:BIGINT;NOT NULL;default:0"`
IsDir bool `form:"is_dir" json:"is_dir" gorm:"type:BOOLEAN;NOT NULL;default:false"`
CreatedAt time.Time `form:"created_at" json:"created_at"`
UpdatedAt time.Time `form:"updated_at" json:"updated_at"`
}
// ResourceEntry is the REST response representation of a single resource.
type ResourceEntry struct {
ID uint64 `json:"id"`
UserID uint64 `json:"user_id"`
Name string `json:"name"`
Path string `json:"path"`
Size int64 `json:"size"`
IsDir bool `json:"is_dir"`
CreatedAt time.Time `json:"created_at"`
UpdatedAt time.Time `json:"updated_at"`
}
// ResourceList is the REST list response for resource operations.
type ResourceList struct {
Items []ResourceEntry `json:"items"`
Total uint64 `json:"total"`
}
// MkdirResourceRequest is the request body for creating a virtual directory.
type MkdirResourceRequest struct {
Path string `json:"path" binding:"required"`
}
// MoveResourceRequest is the request body for moving or renaming a resource.
//
// At least one of Source or Sources must be non-empty after deduplication.
//
// - Single source: Destination is the exact target path (file) or root
// directory (dir tree), identical to the original single-source behaviour.
// - Multiple sources: Destination is treated as a base directory; each
// source's base name is appended automatically, e.g. moving "a.txt" and
// "b.txt" to "archive" produces "archive/a.txt" and "archive/b.txt".
// All moves happen atomically inside a single DB transaction.
type MoveResourceRequest struct {
// Source is kept for backward compatibility (single source).
// Combined with Sources when both are provided; duplicates are removed.
Source string `json:"source"`
// Sources is a list of virtual resource paths to move.
// Combined with Source when both are provided; duplicates are removed.
Sources []string `json:"sources"`
// Destination is the exact target path (single source) or base directory
// (multiple sources). Empty string means the root directory.
Destination string `json:"destination"`
// Force overwrites existing resources at the target paths when true.
Force bool `json:"force"`
}
// CopyResourceRequest is the request body for copying a resource.
//
// At least one of Source or Sources must be non-empty after deduplication.
//
// - Single source: Destination is the exact target path (file) or root
// directory (dir tree), identical to the original single-source behaviour.
// - Multiple sources: Destination is treated as a base directory; each
// source's base name is appended automatically, e.g. copying "a.txt" and
// "b.txt" to "backup" produces "backup/a.txt" and "backup/b.txt".
// All copies happen atomically inside a single DB transaction.
type CopyResourceRequest struct {
// Source is kept for backward compatibility (single source).
// Combined with Sources when both are provided; duplicates are removed.
Source string `json:"source"`
// Sources is a list of virtual resource paths to copy.
// Combined with Source when both are provided; duplicates are removed.
Sources []string `json:"sources"`
// Destination is the exact target path (single source) or base directory
// (multiple sources). Required.
Destination string `json:"destination" binding:"required"`
// Force overwrites existing resources at the target paths when true.
Force bool `json:"force"`
}
+74
View File
@@ -0,0 +1,74 @@
package models
import "github.com/jinzhu/gorm"
// Role is model to contain user role information
// nolint:lll
type Role struct {
ID uint64 `form:"id" json:"id" validate:"min=0,numeric" gorm:"type:BIGINT;NOT NULL;PRIMARY_KEY;AUTO_INCREMENT"`
Name string `form:"name" json:"name" validate:"max=50,required" gorm:"type:TEXT;NOT NULL;UNIQUE_INDEX"`
}
// TableName returns the table name string to guaranty use correct table
func (r *Role) TableName() string {
return "roles"
}
// Valid is function to control input/output data
func (r Role) Valid() error {
return validate.Struct(r)
}
// Validate is function to use callback to control input/output data
func (r Role) Validate(db *gorm.DB) {
if err := r.Valid(); err != nil {
db.AddError(err)
}
}
// Privilege is model to contain user privileges
// nolint:lll
type Privilege struct {
ID uint64 `form:"id" json:"id" validate:"min=0,numeric" gorm:"type:BIGINT;NOT NULL;PRIMARY_KEY;AUTO_INCREMENT"`
RoleID uint64 `form:"role_id" json:"role_id" validate:"min=0,numeric" gorm:"type:BIGINT;NOT NULL"`
Name string `form:"name" json:"name" validate:"max=70,required" gorm:"type:TEXT;NOT NULL"`
}
// TableName returns the table name string to guaranty use correct table
func (p *Privilege) TableName() string {
return "privileges"
}
// Valid is function to control input/output data
func (p Privilege) Valid() error {
return validate.Struct(p)
}
// RolePrivileges is model to contain user role privileges
// nolint:lll
type RolePrivileges struct {
Privileges []Privilege `form:"privileges" json:"privileges" validate:"required" gorm:"foreignkey:RoleID;association_autoupdate:false;association_autocreate:false"`
Role `form:"" json:""`
}
// TableName returns the table name string to guaranty use correct table
func (rp *RolePrivileges) TableName() string {
return "roles"
}
// Valid is function to control input/output data
func (rp RolePrivileges) Valid() error {
for i := range rp.Privileges {
if err := rp.Privileges[i].Valid(); err != nil {
return err
}
}
return rp.Role.Valid()
}
// Validate is function to use callback to control input/output data
func (rp RolePrivileges) Validate(db *gorm.DB) {
if err := rp.Valid(); err != nil {
db.AddError(err)
}
}
+36
View File
@@ -0,0 +1,36 @@
package models
import (
"time"
"github.com/jinzhu/gorm"
)
// Screenshot is model to contain screenshot information
// nolint:lll
type Screenshot struct {
ID uint64 `form:"id" json:"id" validate:"min=0,numeric" gorm:"type:BIGINT;NOT NULL;PRIMARY_KEY;AUTO_INCREMENT"`
Name string `form:"name" json:"name" validate:"required" gorm:"type:TEXT;NOT NULL"`
URL string `form:"url" json:"url" validate:"required" gorm:"type:TEXT;NOT NULL"`
FlowID uint64 `form:"flow_id" json:"flow_id" validate:"min=0,numeric,required" gorm:"type:BIGINT;NOT NULL"`
TaskID *uint64 `form:"task_id,omitempty" json:"task_id,omitempty" validate:"omitnil,min=0" gorm:"type:BIGINT"`
SubtaskID *uint64 `form:"subtask_id,omitempty" json:"subtask_id,omitempty" validate:"omitnil,min=0" gorm:"type:BIGINT"`
CreatedAt time.Time `form:"created_at,omitempty" json:"created_at,omitempty" validate:"omitempty" gorm:"type:TIMESTAMPTZ;default:CURRENT_TIMESTAMP"`
}
// TableName returns the table name string to guaranty use correct table
func (s *Screenshot) TableName() string {
return "screenshots"
}
// Valid is function to control input/output data
func (s Screenshot) Valid() error {
return validate.Struct(s)
}
// Validate is function to use callback to control input/output data
func (s Screenshot) Validate(db *gorm.DB) {
if err := s.Valid(); err != nil {
db.AddError(err)
}
}
+79
View File
@@ -0,0 +1,79 @@
package models
import (
"fmt"
"time"
"github.com/jinzhu/gorm"
)
type SearchEngineType string
const (
SearchEngineTypeGoogle SearchEngineType = "google"
SearchEngineTypeDuckduckgo SearchEngineType = "duckduckgo"
SearchEngineTypeTavily SearchEngineType = "tavily"
SearchEngineTypeTraversaal SearchEngineType = "traversaal"
SearchEngineTypePerplexity SearchEngineType = "perplexity"
SearchEngineTypeBrowser SearchEngineType = "browser"
SearchEngineTypeSploitus SearchEngineType = "sploitus"
)
func (s SearchEngineType) String() string {
return string(s)
}
// Valid is function to control input/output data
func (s SearchEngineType) Valid() error {
switch s {
case SearchEngineTypeGoogle,
SearchEngineTypeDuckduckgo,
SearchEngineTypeTavily,
SearchEngineTypeTraversaal,
SearchEngineTypePerplexity,
SearchEngineTypeBrowser,
SearchEngineTypeSploitus:
return nil
default:
return fmt.Errorf("invalid SearchEngineType: %s", s)
}
}
// Validate is function to use callback to control input/output data
func (s SearchEngineType) Validate(db *gorm.DB) {
if err := s.Valid(); err != nil {
db.AddError(err)
}
}
// Searchlog is model to contain search action information in the internet or local network
// nolint:lll
type Searchlog struct {
ID uint64 `form:"id" json:"id" validate:"min=0,numeric" gorm:"type:BIGINT;NOT NULL;PRIMARY_KEY;AUTO_INCREMENT"`
Initiator MsgchainType `json:"initiator" validate:"valid,required" gorm:"type:MSGCHAIN_TYPE;NOT NULL"`
Executor MsgchainType `json:"executor" validate:"valid,required" gorm:"type:MSGCHAIN_TYPE;NOT NULL"`
Engine SearchEngineType `json:"engine" validate:"valid,required" gorm:"type:SEARCHENGINE_TYPE;NOT NULL"`
Query string `json:"query" validate:"required" gorm:"type:TEXT;NOT NULL"`
Result string `json:"result" validate:"omitempty" gorm:"type:TEXT;NOT NULL;default:''"`
FlowID uint64 `form:"flow_id" json:"flow_id" validate:"min=0,numeric,required" gorm:"type:BIGINT;NOT NULL"`
TaskID *uint64 `form:"task_id,omitempty" json:"task_id,omitempty" validate:"omitnil,min=0" gorm:"type:BIGINT;NOT NULL"`
SubtaskID *uint64 `form:"subtask_id,omitempty" json:"subtask_id,omitempty" validate:"omitnil,min=0" gorm:"type:BIGINT;NOT NULL"`
CreatedAt time.Time `form:"created_at,omitempty" json:"created_at,omitempty" validate:"omitempty" gorm:"type:TIMESTAMPTZ;default:CURRENT_TIMESTAMP"`
}
// TableName returns the table name string to guaranty use correct table
func (ml *Searchlog) TableName() string {
return "searchlogs"
}
// Valid is function to control input/output data
func (ml Searchlog) Valid() error {
return validate.Struct(ml)
}
// Validate is function to use callback to control input/output data
func (ml Searchlog) Validate(db *gorm.DB) {
if err := ml.Valid(); err != nil {
db.AddError(err)
}
}
+17
View File
@@ -0,0 +1,17 @@
package models
// Settings is model to contain application settings information
// nolint:lll
type Settings struct {
Debug bool `json:"debug" example:"false"`
AskUser bool `json:"ask_user" example:"false"`
Version string `json:"version" example:"v1.0.0"`
DockerInside bool `json:"docker_inside" example:"false"`
IsDevelopMode bool `json:"is_develop_mode" example:"false"`
AssistantUseAgents bool `json:"assistant_use_agents" example:"false"`
}
// Valid is function to control input/output data
func (s Settings) Valid() error {
return validate.Struct(s)
}
+74
View File
@@ -0,0 +1,74 @@
package models
import (
"fmt"
"time"
"github.com/jinzhu/gorm"
)
type SubtaskStatus string
const (
SubtaskStatusCreated SubtaskStatus = "created"
SubtaskStatusRunning SubtaskStatus = "running"
SubtaskStatusWaiting SubtaskStatus = "waiting"
SubtaskStatusFinished SubtaskStatus = "finished"
SubtaskStatusFailed SubtaskStatus = "failed"
)
func (s SubtaskStatus) String() string {
return string(s)
}
// Valid is function to control input/output data
func (s SubtaskStatus) Valid() error {
switch s {
case SubtaskStatusCreated,
SubtaskStatusRunning,
SubtaskStatusWaiting,
SubtaskStatusFinished,
SubtaskStatusFailed:
return nil
default:
return fmt.Errorf("invalid SubtaskStatus: %s", s)
}
}
// Validate is function to use callback to control input/output data
func (s SubtaskStatus) Validate(db *gorm.DB) {
if err := s.Valid(); err != nil {
db.AddError(err)
}
}
// Subtask is model to contain subtask information
// nolint:lll
type Subtask struct {
ID uint64 `form:"id" json:"id" validate:"min=0,numeric" gorm:"type:BIGINT;NOT NULL;PRIMARY_KEY;AUTO_INCREMENT"`
Status SubtaskStatus `form:"status" json:"status" validate:"valid,required" gorm:"type:SUBTASK_STATUS;NOT NULL;default:'created'"`
Title string `form:"title" json:"title" validate:"required" gorm:"type:TEXT;NOT NULL"`
Description string `form:"description" json:"description" validate:"required" gorm:"type:TEXT;NOT NULL"`
Context string `form:"context" json:"context" validate:"omitempty" gorm:"type:TEXT;NOT NULL;default:''"`
Result string `form:"result" json:"result" validate:"omitempty" gorm:"type:TEXT;NOT NULL;default:''"`
TaskID uint64 `form:"task_id" json:"task_id" validate:"min=0,numeric,required" gorm:"type:BIGINT;NOT NULL"`
CreatedAt time.Time `form:"created_at,omitempty" json:"created_at,omitempty" validate:"omitempty" gorm:"type:TIMESTAMPTZ;default:CURRENT_TIMESTAMP"`
UpdatedAt time.Time `form:"updated_at,omitempty" json:"updated_at,omitempty" validate:"omitempty" gorm:"type:TIMESTAMPTZ;default:CURRENT_TIMESTAMP"`
}
// TableName returns the table name string to guaranty use correct table
func (s *Subtask) TableName() string {
return "subtasks"
}
// Valid is function to control input/output data
func (s Subtask) Valid() error {
return validate.Struct(s)
}
// Validate is function to use callback to control input/output data
func (s Subtask) Validate(db *gorm.DB) {
if err := s.Valid(); err != nil {
db.AddError(err)
}
}
+102
View File
@@ -0,0 +1,102 @@
package models
import (
"fmt"
"time"
"github.com/jinzhu/gorm"
)
type TaskStatus string
const (
TaskStatusCreated TaskStatus = "created"
TaskStatusRunning TaskStatus = "running"
TaskStatusWaiting TaskStatus = "waiting"
TaskStatusFinished TaskStatus = "finished"
TaskStatusFailed TaskStatus = "failed"
)
func (s TaskStatus) String() string {
return string(s)
}
// Valid is function to control input/output data
func (s TaskStatus) Valid() error {
switch s {
case TaskStatusCreated,
TaskStatusRunning,
TaskStatusWaiting,
TaskStatusFinished,
TaskStatusFailed:
return nil
default:
return fmt.Errorf("invalid TaskStatus: %s", s)
}
}
// Validate is function to use callback to control input/output data
func (s TaskStatus) Validate(db *gorm.DB) {
if err := s.Valid(); err != nil {
db.AddError(err)
}
}
// Task is model to contain task information
// nolint:lll
type Task struct {
ID uint64 `form:"id" json:"id" validate:"min=0,numeric" gorm:"type:BIGINT;NOT NULL;PRIMARY_KEY;AUTO_INCREMENT"`
Status TaskStatus `form:"status" json:"status" validate:"valid,required" gorm:"type:TASK_STATUS;NOT NULL;default:'created'"`
Title string `form:"title" json:"title" validate:"required" gorm:"type:TEXT;NOT NULL;default:'untitled'"`
Input string `form:"input" json:"input" validate:"required" gorm:"type:TEXT;NOT NULL"`
Result string `form:"result" json:"result" validate:"omitempty" gorm:"type:TEXT;NOT NULL;default:''"`
FlowID uint64 `form:"flow_id" json:"flow_id" validate:"min=0,numeric,required" gorm:"type:BIGINT;NOT NULL"`
CreatedAt time.Time `form:"created_at,omitempty" json:"created_at,omitempty" validate:"omitempty" gorm:"type:TIMESTAMPTZ;default:CURRENT_TIMESTAMP"`
UpdatedAt time.Time `form:"updated_at,omitempty" json:"updated_at,omitempty" validate:"omitempty" gorm:"type:TIMESTAMPTZ;default:CURRENT_TIMESTAMP"`
}
// TableName returns the table name string to guaranty use correct table
func (t *Task) TableName() string {
return "tasks"
}
// Valid is function to control input/output data
func (t Task) Valid() error {
return validate.Struct(t)
}
// Validate is function to use callback to control input/output data
func (t Task) Validate(db *gorm.DB) {
if err := t.Valid(); err != nil {
db.AddError(err)
}
}
// TaskSubtasks is model to contain task and linked subtasks information
// nolint:lll
type TaskSubtasks struct {
Subtasks []Subtask `form:"subtasks" json:"subtasks" validate:"required" gorm:"foreignkey:TaskID;association_autoupdate:false;association_autocreate:false"`
Task `form:"" json:""`
}
// TableName returns the table name string to guaranty use correct table
func (ts *TaskSubtasks) TableName() string {
return "tasks"
}
// Valid is function to control input/output data
func (ts TaskSubtasks) Valid() error {
for i := range ts.Subtasks {
if err := ts.Subtasks[i].Valid(); err != nil {
return err
}
}
return ts.Task.Valid()
}
// Validate is function to use callback to control input/output data
func (ts TaskSubtasks) Validate(db *gorm.DB) {
if err := ts.Valid(); err != nil {
db.AddError(err)
}
}
+67
View File
@@ -0,0 +1,67 @@
package models
import (
"fmt"
"time"
"github.com/jinzhu/gorm"
)
type TermlogType string
const (
TermlogTypeStdin TermlogType = "stdin"
TermlogTypeStdout TermlogType = "stdout"
TermlogTypeStderr TermlogType = "stderr"
)
func (s TermlogType) String() string {
return string(s)
}
// Valid is function to control input/output data
func (s TermlogType) Valid() error {
switch s {
case TermlogTypeStdin, TermlogTypeStdout, TermlogTypeStderr:
return nil
default:
return fmt.Errorf("invalid TermlogType: %s", s)
}
}
// Validate is function to use callback to control input/output data
func (s TermlogType) Validate(db *gorm.DB) {
if err := s.Valid(); err != nil {
db.AddError(err)
}
}
// Termlog is model to contain termlog information
// nolint:lll
type Termlog struct {
ID uint64 `form:"id" json:"id" validate:"min=0,numeric" gorm:"type:BIGINT;NOT NULL;PRIMARY_KEY;AUTO_INCREMENT"`
Type TermlogType `form:"type" json:"type" validate:"valid,required" gorm:"type:TERMLOG_TYPE;NOT NULL"`
Text string `form:"text" json:"text" validate:"required" gorm:"type:TEXT;NOT NULL"`
ContainerID uint64 `form:"container_id" json:"container_id" validate:"min=0,numeric,required" gorm:"type:BIGINT;NOT NULL"`
FlowID uint64 `form:"flow_id" json:"flow_id" validate:"min=0,numeric,required" gorm:"type:BIGINT;NOT NULL"`
TaskID *uint64 `form:"task_id,omitempty" json:"task_id,omitempty" validate:"omitnil,min=0" gorm:"type:BIGINT"`
SubtaskID *uint64 `form:"subtask_id,omitempty" json:"subtask_id,omitempty" validate:"omitnil,min=0" gorm:"type:BIGINT"`
CreatedAt time.Time `form:"created_at,omitempty" json:"created_at,omitempty" validate:"omitempty" gorm:"type:TIMESTAMPTZ;default:CURRENT_TIMESTAMP"`
}
// TableName returns the table name string to guaranty use correct table
func (tl *Termlog) TableName() string {
return "termlogs"
}
// Valid is function to control input/output data
func (tl Termlog) Valid() error {
return validate.Struct(tl)
}
// Validate is function to use callback to control input/output data
func (tl Termlog) Validate(db *gorm.DB) {
if err := tl.Valid(); err != nil {
db.AddError(err)
}
}
+75
View File
@@ -0,0 +1,75 @@
package models
import (
"fmt"
"time"
"github.com/jinzhu/gorm"
)
type ToolcallStatus string
const (
ToolcallStatusReceived ToolcallStatus = "received"
ToolcallStatusRunning ToolcallStatus = "running"
ToolcallStatusFinished ToolcallStatus = "finished"
ToolcallStatusFailed ToolcallStatus = "failed"
)
func (s ToolcallStatus) String() string {
return string(s)
}
// Valid is function to control input/output data
func (s ToolcallStatus) Valid() error {
switch s {
case ToolcallStatusReceived,
ToolcallStatusRunning,
ToolcallStatusFinished,
ToolcallStatusFailed:
return nil
default:
return fmt.Errorf("invalid ToolcallStatus: %s", s)
}
}
// Validate is function to use callback to control input/output data
func (s ToolcallStatus) Validate(db *gorm.DB) {
if err := s.Valid(); err != nil {
db.AddError(err)
}
}
// Toolcall is model to contain tool call information
// nolint:lll
type Toolcall struct {
ID uint64 `form:"id" json:"id" validate:"min=0,numeric" gorm:"type:BIGINT;NOT NULL;PRIMARY_KEY;AUTO_INCREMENT"`
CallID string `form:"call_id" json:"call_id" validate:"required" gorm:"type:TEXT;NOT NULL"`
Status ToolcallStatus `form:"status" json:"status" validate:"valid,required" gorm:"type:TOOLCALL_STATUS;NOT NULL;default:'received'"`
Name string `form:"name" json:"name" validate:"required" gorm:"type:TEXT;NOT NULL"`
Args string `form:"args" json:"args" validate:"required" gorm:"type:JSON;NOT NULL"`
Result string `form:"result" json:"result" validate:"omitempty" gorm:"type:TEXT;NOT NULL;default:''"`
DurationSeconds float64 `form:"duration_seconds" json:"duration_seconds" validate:"min=0" gorm:"type:DOUBLE PRECISION;NOT NULL;default:0"`
FlowID uint64 `form:"flow_id" json:"flow_id" validate:"min=0,numeric,required" gorm:"type:BIGINT;NOT NULL"`
TaskID *uint64 `form:"task_id,omitempty" json:"task_id,omitempty" validate:"omitnil,min=0" gorm:"type:BIGINT"`
SubtaskID *uint64 `form:"subtask_id,omitempty" json:"subtask_id,omitempty" validate:"omitnil,min=0" gorm:"type:BIGINT"`
CreatedAt time.Time `form:"created_at,omitempty" json:"created_at,omitempty" validate:"omitempty" gorm:"type:TIMESTAMPTZ;default:CURRENT_TIMESTAMP"`
UpdatedAt time.Time `form:"updated_at,omitempty" json:"updated_at,omitempty" validate:"omitempty" gorm:"type:TIMESTAMPTZ;default:CURRENT_TIMESTAMP"`
}
// TableName returns the table name string to guaranty use correct table
func (tc *Toolcall) TableName() string {
return "toolcalls"
}
// Valid is function to control input/output data
func (tc Toolcall) Valid() error {
return validate.Struct(tc)
}
// Validate is function to use callback to control input/output data
func (tc Toolcall) Validate(db *gorm.DB) {
if err := tc.Valid(); err != nil {
db.AddError(err)
}
}
+304
View File
@@ -0,0 +1,304 @@
package models
import (
"database/sql/driver"
"encoding/json"
"fmt"
"time"
"github.com/jinzhu/gorm"
)
const RoleUser = 2
type UserStatus string
const (
UserStatusCreated UserStatus = "created"
UserStatusActive UserStatus = "active"
UserStatusBlocked UserStatus = "blocked"
)
func (s UserStatus) String() string {
return string(s)
}
// Valid is function to control input/output data
func (s UserStatus) Valid() error {
switch s {
case UserStatusCreated, UserStatusActive, UserStatusBlocked:
return nil
default:
return fmt.Errorf("invalid UserStatus: %s", s)
}
}
// Validate is function to use callback to control input/output data
func (s UserStatus) Validate(db *gorm.DB) {
if err := s.Valid(); err != nil {
db.AddError(err)
}
}
type UserType string
const (
UserTypeLocal UserType = "local"
UserTypeOAuth UserType = "oauth"
UserTypeAPI UserType = "api"
)
func (s UserType) String() string {
return string(s)
}
// Valid is function to control input/output data
func (s UserType) Valid() error {
switch s {
case UserTypeLocal, UserTypeOAuth, UserTypeAPI:
return nil
default:
return fmt.Errorf("invalid UserType: %s", s)
}
}
// Validate is function to use callback to control input/output data
func (s UserType) Validate(db *gorm.DB) {
if err := s.Valid(); err != nil {
db.AddError(err)
}
}
// User is model to contain user information
// nolint:lll
type User struct {
ID uint64 `form:"id" json:"id" validate:"min=0,numeric" gorm:"type:BIGINT;NOT NULL;PRIMARY_KEY;AUTO_INCREMENT"`
Hash string `form:"hash" json:"hash" validate:"len=32,hexadecimal,lowercase,omitempty" gorm:"type:TEXT;NOT NULL;UNIQUE_INDEX;default:MD5(RANDOM()::text)"`
Type UserType `form:"type" json:"type" validate:"valid,required" gorm:"type:USER_TYPE;NOT NULL;default:'local'"`
Mail string `form:"mail" json:"mail" validate:"max=50,vmail,required" gorm:"type:TEXT;NOT NULL;UNIQUE_INDEX"`
Name string `form:"name" json:"name" validate:"max=70,omitempty" gorm:"type:TEXT;NOT NULL;default:''"`
Status UserStatus `form:"status" json:"status" validate:"valid,required" gorm:"type:USER_STATUS;NOT NULL;default:'created'"`
RoleID uint64 `form:"role_id" json:"role_id" validate:"min=0,numeric,required" gorm:"type:BIGINT;NOT NULL;default:2"`
PasswordChangeRequired bool `form:"password_change_required" json:"password_change_required" gorm:"type:BOOL;NOT NULL;default:false"`
Provider *string `form:"provider,omitempty" json:"provider,omitempty" validate:"omitempty" gorm:"type:TEXT"`
CreatedAt time.Time `form:"created_at" json:"created_at" validate:"omitempty" gorm:"type:TIMESTAMPTZ;NOT NULL;default:CURRENT_TIMESTAMP"`
}
// TableName returns the table name string to guaranty use correct table
func (u *User) TableName() string {
return "users"
}
// Valid is function to control input/output data
func (u User) Valid() error {
return validate.Struct(u)
}
// Validate is function to use callback to control input/output data
func (u User) Validate(db *gorm.DB) {
if err := u.Valid(); err != nil {
db.AddError(err)
}
}
// UserPassword is model to contain user information
type UserPassword struct {
Password string `form:"password" json:"password" validate:"max=100,required" gorm:"column:password;type:TEXT"`
User `form:"" json:""`
}
// TableName returns the table name string to guaranty use correct table
func (up *UserPassword) TableName() string {
return "users"
}
// Valid is function to control input/output data
func (up UserPassword) Valid() error {
if err := up.User.Valid(); err != nil {
return err
}
return validate.Struct(up)
}
// Validate is function to use callback to control input/output data
func (up UserPassword) Validate(db *gorm.DB) {
if err := up.Valid(); err != nil {
db.AddError(err)
}
}
// Login is model to contain user information on Login procedure
// nolint:lll
type Login struct {
Mail string `form:"mail" json:"mail" validate:"max=50,required" gorm:"type:TEXT;NOT NULL;UNIQUE_INDEX"`
Password string `form:"password" json:"password" validate:"min=4,max=100,required" gorm:"type:TEXT"`
}
// TableName returns the table name string to guaranty use correct table
func (sin *Login) TableName() string {
return "users"
}
// Valid is function to control input/output data
func (sin Login) Valid() error {
return validate.Struct(sin)
}
// Validate is function to use callback to control input/output data
func (sin Login) Validate(db *gorm.DB) {
if err := sin.Valid(); err != nil {
db.AddError(err)
}
}
// AuthCallback is model to contain auth data information from external OAuth application
type AuthCallback struct {
Code string `form:"code" json:"code" validate:"required"`
IdToken string `form:"id_token" json:"id_token" validate:"required,jwt"`
Scope string `form:"scope" json:"scope" validate:"required,oauth_min_scope"`
State string `form:"state" json:"state" validate:"required"`
}
// Valid is function to control input/output data
func (au AuthCallback) Valid() error {
return validate.Struct(au)
}
// Password is model to contain user password to change it
// nolint:lll
type Password struct {
CurrentPassword string `form:"current_password" json:"current_password" validate:"nefield=Password,min=5,max=100,required" gorm:"-"`
Password string `form:"password" json:"password" validate:"stpass,max=100,required" gorm:"type:TEXT"`
ConfirmPassword string `form:"confirm_password" json:"confirm_password" validate:"eqfield=Password" gorm:"-"`
}
// TableName returns the table name string to guaranty use correct table
func (p *Password) TableName() string {
return "users"
}
// Valid is function to control input/output data
func (p Password) Valid() error {
return validate.Struct(p)
}
// Validate is function to use callback to control input/output data
func (p Password) Validate(db *gorm.DB) {
if err := p.Valid(); err != nil {
db.AddError(err)
}
}
// UserRole is model to contain user information linked with user role
// nolint:lll
type UserRole struct {
Role Role `form:"role,omitempty" json:"role,omitempty" gorm:"association_autoupdate:false;association_autocreate:false"`
User `form:"" json:""`
}
// Valid is function to control input/output data
func (ur UserRole) Valid() error {
if err := ur.Role.Valid(); err != nil {
return err
}
return ur.User.Valid()
}
// Validate is function to use callback to control input/output data
func (ur UserRole) Validate(db *gorm.DB) {
if err := ur.Valid(); err != nil {
db.AddError(err)
}
}
// UserRole is model to contain user information linked with user role
// nolint:lll
type UserRolePrivileges struct {
Role RolePrivileges `form:"role,omitempty" json:"role,omitempty" gorm:"association_autoupdate:false;association_autocreate:false"`
User `form:"" json:""`
}
// Valid is function to control input/output data
func (urp UserRolePrivileges) Valid() error {
if err := urp.Role.Valid(); err != nil {
return err
}
return urp.User.Valid()
}
// Validate is function to use callback to control input/output data
func (urp UserRolePrivileges) Validate(db *gorm.DB) {
if err := urp.Valid(); err != nil {
db.AddError(err)
}
}
// UserPreferencesOptions is model to contain user preferences as JSON
type UserPreferencesOptions struct {
FavoriteFlows []int64 `json:"favoriteFlows"`
}
// Value implements driver.Valuer interface for database write
func (upo UserPreferencesOptions) Value() (driver.Value, error) {
return json.Marshal(upo)
}
// Scan implements sql.Scanner interface for database read
func (upo *UserPreferencesOptions) Scan(value any) error {
if value == nil {
*upo = UserPreferencesOptions{FavoriteFlows: []int64{}}
return nil
}
bytes, ok := value.([]byte)
if !ok {
return fmt.Errorf("failed to scan UserPreferencesOptions: expected []byte, got %T", value)
}
return json.Unmarshal(bytes, upo)
}
// UserPreferences is model to contain user preferences information
type UserPreferences struct {
ID uint64 `json:"id" gorm:"type:BIGINT;NOT NULL;PRIMARY_KEY;AUTO_INCREMENT"`
UserID uint64 `json:"user_id" gorm:"type:BIGINT;NOT NULL;UNIQUE_INDEX"`
Preferences UserPreferencesOptions `json:"preferences" gorm:"type:JSONB;NOT NULL"`
CreatedAt time.Time `json:"created_at" gorm:"type:TIMESTAMPTZ;NOT NULL;default:CURRENT_TIMESTAMP"`
UpdatedAt time.Time `json:"updated_at" gorm:"type:TIMESTAMPTZ;NOT NULL;default:CURRENT_TIMESTAMP"`
}
// TableName returns the table name string to guaranty use correct table
func (up *UserPreferences) TableName() string {
return "user_preferences"
}
// Valid is function to control input/output data
func (up UserPreferences) Valid() error {
if up.UserID == 0 {
return fmt.Errorf("user_id is required")
}
return nil
}
// Validate is function to use callback to control input/output data
func (up UserPreferences) Validate(db *gorm.DB) {
if err := up.Valid(); err != nil {
db.AddError(err)
}
}
// NewUserPreferences creates a new UserPreferences with default values
func NewUserPreferences(userID uint64) *UserPreferences {
return &UserPreferences{
UserID: userID,
Preferences: UserPreferencesOptions{
FavoriteFlows: []int64{},
},
}
}
// UserWithPreferences is model to combine User and UserPreferences for transactional creation
type UserWithPreferences struct {
User User
Preferences UserPreferences
}
+562
View File
@@ -0,0 +1,562 @@
package models
import (
"testing"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
func TestUserStatusValid(t *testing.T) {
t.Parallel()
tests := []struct {
name string
status UserStatus
wantErr bool
}{
{"valid created", UserStatusCreated, false},
{"valid active", UserStatusActive, false},
{"valid blocked", UserStatusBlocked, false},
{"invalid empty", UserStatus(""), true},
{"invalid unknown", UserStatus("unknown"), true},
{"invalid suspended", UserStatus("suspended"), true},
}
for _, tt := range tests {
tt := tt
t.Run(tt.name, func(t *testing.T) {
t.Parallel()
err := tt.status.Valid()
if tt.wantErr {
assert.Error(t, err)
assert.Contains(t, err.Error(), "invalid UserStatus")
} else {
assert.NoError(t, err)
}
})
}
}
func TestUserStatusString(t *testing.T) {
t.Parallel()
assert.Equal(t, "created", UserStatusCreated.String())
assert.Equal(t, "active", UserStatusActive.String())
assert.Equal(t, "blocked", UserStatusBlocked.String())
}
func TestUserTypeValid(t *testing.T) {
t.Parallel()
tests := []struct {
name string
userType UserType
wantErr bool
}{
{"valid local", UserTypeLocal, false},
{"valid oauth", UserTypeOAuth, false},
{"valid api", UserTypeAPI, false},
{"invalid empty", UserType(""), true},
{"invalid unknown", UserType("unknown"), true},
{"invalid saml", UserType("saml"), true},
}
for _, tt := range tests {
tt := tt
t.Run(tt.name, func(t *testing.T) {
t.Parallel()
err := tt.userType.Valid()
if tt.wantErr {
assert.Error(t, err)
assert.Contains(t, err.Error(), "invalid UserType")
} else {
assert.NoError(t, err)
}
})
}
}
func TestUserTypeString(t *testing.T) {
t.Parallel()
assert.Equal(t, "local", UserTypeLocal.String())
assert.Equal(t, "oauth", UserTypeOAuth.String())
assert.Equal(t, "api", UserTypeAPI.String())
}
func TestLoginValid(t *testing.T) {
t.Parallel()
tests := []struct {
name string
login Login
wantErr bool
}{
{
name: "valid login",
login: Login{Mail: "test@example.com", Password: "password123"},
wantErr: false,
},
{
name: "valid admin mail",
login: Login{Mail: "admin", Password: "password123"},
wantErr: false,
},
{
name: "empty mail",
login: Login{Mail: "", Password: "password123"},
wantErr: true,
},
{
name: "empty password",
login: Login{Mail: "test@example.com", Password: ""},
wantErr: true,
},
{
name: "password too short",
login: Login{Mail: "test@example.com", Password: "ab"},
wantErr: true,
},
}
for _, tt := range tests {
tt := tt
t.Run(tt.name, func(t *testing.T) {
t.Parallel()
err := tt.login.Valid()
if tt.wantErr {
assert.Error(t, err)
} else {
assert.NoError(t, err)
}
})
}
}
func TestPasswordValid(t *testing.T) {
t.Parallel()
tests := []struct {
name string
pw Password
wantErr bool
}{
{
name: "valid strong password with special chars",
pw: Password{
CurrentPassword: "OldPass1!abc",
Password: "NewPass1!abc",
ConfirmPassword: "NewPass1!abc",
},
wantErr: false,
},
{
name: "valid long password over 15 chars",
pw: Password{
CurrentPassword: "oldpasswordvalue",
Password: "newpasswordvalue1",
ConfirmPassword: "newpasswordvalue1",
},
wantErr: false,
},
{
name: "confirm password mismatch",
pw: Password{
CurrentPassword: "OldPass1!abc",
Password: "NewPass1!abc",
ConfirmPassword: "DifferentPass1!",
},
wantErr: true,
},
{
name: "current equals new password",
pw: Password{
CurrentPassword: "SamePass1!abc",
Password: "SamePass1!abc",
ConfirmPassword: "SamePass1!abc",
},
wantErr: true,
},
{
name: "weak password no special chars",
pw: Password{
CurrentPassword: "OldPass1!abc",
Password: "newpass1",
ConfirmPassword: "newpass1",
},
wantErr: true,
},
{
name: "empty current password",
pw: Password{
CurrentPassword: "",
Password: "NewPass1!abc",
ConfirmPassword: "NewPass1!abc",
},
wantErr: true,
},
}
for _, tt := range tests {
tt := tt
t.Run(tt.name, func(t *testing.T) {
t.Parallel()
err := tt.pw.Valid()
if tt.wantErr {
assert.Error(t, err)
} else {
assert.NoError(t, err)
}
})
}
}
func TestUserValid(t *testing.T) {
t.Parallel()
validUser := User{
ID: 1,
Hash: "abcdef1234567890abcdef1234567890",
Type: UserTypeLocal,
Mail: "test@example.com",
Status: UserStatusActive,
RoleID: RoleUser,
}
t.Run("valid user", func(t *testing.T) {
t.Parallel()
assert.NoError(t, validUser.Valid())
})
t.Run("missing mail", func(t *testing.T) {
t.Parallel()
u := validUser
u.Mail = ""
assert.Error(t, u.Valid())
})
t.Run("invalid user type", func(t *testing.T) {
t.Parallel()
u := validUser
u.Type = UserType("invalid")
assert.Error(t, u.Valid())
})
t.Run("invalid user status", func(t *testing.T) {
t.Parallel()
u := validUser
u.Status = UserStatus("invalid")
assert.Error(t, u.Valid())
})
t.Run("invalid hash length", func(t *testing.T) {
t.Parallel()
u := validUser
u.Hash = "tooshort"
assert.Error(t, u.Valid())
})
}
func TestUserTableName(t *testing.T) {
t.Parallel()
u := &User{}
assert.Equal(t, "users", u.TableName())
}
func TestUserPasswordValid(t *testing.T) {
t.Parallel()
t.Run("valid user password", func(t *testing.T) {
t.Parallel()
up := UserPassword{
Password: "somepassword",
User: User{
ID: 1,
Hash: "abcdef1234567890abcdef1234567890",
Type: UserTypeLocal,
Mail: "test@example.com",
Status: UserStatusActive,
RoleID: RoleUser,
},
}
assert.NoError(t, up.Valid())
})
t.Run("empty password", func(t *testing.T) {
t.Parallel()
up := UserPassword{
Password: "",
User: User{
ID: 1,
Hash: "abcdef1234567890abcdef1234567890",
Type: UserTypeLocal,
Mail: "test@example.com",
Status: UserStatusActive,
RoleID: RoleUser,
},
}
assert.Error(t, up.Valid())
})
t.Run("invalid user in user password", func(t *testing.T) {
t.Parallel()
up := UserPassword{
Password: "somepassword",
User: User{
Mail: "",
},
}
assert.Error(t, up.Valid())
})
}
func TestUserPasswordTableName(t *testing.T) {
t.Parallel()
up := &UserPassword{}
assert.Equal(t, "users", up.TableName())
}
func TestLoginTableName(t *testing.T) {
t.Parallel()
l := &Login{}
assert.Equal(t, "users", l.TableName())
}
func TestPasswordTableName(t *testing.T) {
t.Parallel()
p := &Password{}
assert.Equal(t, "users", p.TableName())
}
func TestUserPreferencesOptionsValueScan(t *testing.T) {
t.Parallel()
t.Run("value and scan round trip", func(t *testing.T) {
t.Parallel()
original := UserPreferencesOptions{FavoriteFlows: []int64{1, 2, 3}}
val, err := original.Value()
require.NoError(t, err)
var scanned UserPreferencesOptions
switch v := val.(type) {
case string:
err = scanned.Scan([]byte(v))
case []byte:
err = scanned.Scan(v)
default:
t.Fatalf("unexpected Value() type: %T", val)
}
require.NoError(t, err)
assert.Equal(t, original.FavoriteFlows, scanned.FavoriteFlows)
})
t.Run("scan nil value", func(t *testing.T) {
t.Parallel()
var upo UserPreferencesOptions
err := upo.Scan(nil)
require.NoError(t, err)
assert.Equal(t, []int64{}, upo.FavoriteFlows)
})
t.Run("scan unsupported type", func(t *testing.T) {
t.Parallel()
var upo UserPreferencesOptions
err := upo.Scan(12345)
assert.Error(t, err)
assert.Contains(t, err.Error(), "expected []byte")
})
t.Run("scan invalid json", func(t *testing.T) {
t.Parallel()
var upo UserPreferencesOptions
err := upo.Scan([]byte("not json"))
assert.Error(t, err)
})
t.Run("value with empty flows", func(t *testing.T) {
t.Parallel()
upo := UserPreferencesOptions{FavoriteFlows: []int64{}}
val, err := upo.Value()
require.NoError(t, err)
var valStr string
switch v := val.(type) {
case string:
valStr = v
case []byte:
valStr = string(v)
}
assert.Contains(t, valStr, "favoriteFlows")
})
}
func TestUserPreferencesValid(t *testing.T) {
t.Parallel()
t.Run("valid preferences", func(t *testing.T) {
t.Parallel()
up := UserPreferences{UserID: 1}
assert.NoError(t, up.Valid())
})
t.Run("zero user id", func(t *testing.T) {
t.Parallel()
up := UserPreferences{UserID: 0}
err := up.Valid()
assert.Error(t, err)
assert.Contains(t, err.Error(), "user_id")
})
}
func TestUserPreferencesTableName(t *testing.T) {
t.Parallel()
up := &UserPreferences{}
assert.Equal(t, "user_preferences", up.TableName())
}
func TestNewUserPreferences(t *testing.T) {
t.Parallel()
up := NewUserPreferences(42)
assert.Equal(t, uint64(42), up.UserID)
assert.NotNil(t, up.Preferences.FavoriteFlows)
assert.Empty(t, up.Preferences.FavoriteFlows)
}
func TestAuthCallbackValid(t *testing.T) {
t.Parallel()
// JWT token with 3 dot-separated base64 segments (header.payload.signature)
validJWT := "eyJhbGciOiJSUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiIxMjM0NTY3ODkwIn0.POstGetfAytaZS82wHcjoTyoqhMyxXiWdR7Nn7A29DNSl0EiXLdwJ6xC6AfgZWF1bOsS_TuYI3OG85AmiExREkrS6tDfTQ2B3WXlrr-wp5AokiRbz3_oB4OxG-W9KcEEbDRcZc0nH3L7LzYptiy1PtAylQGxHTWZXtGz4ht0bAecBgmpdgXMguEIcoqPJ1n3pIWk_dUZegpqx0Lka21H6XxUTxiy8OcaarA8zdnPUnV6AmNP3ecFawIFYdvJB_cm-GvpCSbr8G8y_Mllj8f4x9nBH8pQux89_6gUY618iYv7tuPWBFfEbLxtF2pZS6YC1aSfLQxaOoaBSTNRg"
t.Run("valid callback", func(t *testing.T) {
t.Parallel()
ac := AuthCallback{
Code: "auth-code-123",
IdToken: validJWT,
Scope: "openid email profile",
State: "random-state-value",
}
assert.NoError(t, ac.Valid())
})
t.Run("missing code", func(t *testing.T) {
t.Parallel()
ac := AuthCallback{
Code: "",
IdToken: validJWT,
Scope: "openid email",
State: "state123",
}
assert.Error(t, ac.Valid())
})
t.Run("scope missing openid", func(t *testing.T) {
t.Parallel()
ac := AuthCallback{
Code: "code123",
IdToken: validJWT,
Scope: "email profile",
State: "state123",
}
assert.Error(t, ac.Valid())
})
t.Run("invalid id token not jwt", func(t *testing.T) {
t.Parallel()
ac := AuthCallback{
Code: "code123",
IdToken: "not-a-jwt",
Scope: "openid email",
State: "state123",
}
assert.Error(t, ac.Valid())
})
}
func TestUserRoleValid(t *testing.T) {
t.Parallel()
validUserForRole := User{
Hash: "abcdef1234567890abcdef1234567890",
Type: UserTypeLocal,
Mail: "test@example.com",
Status: UserStatusActive,
RoleID: RoleUser,
}
t.Run("valid user role", func(t *testing.T) {
t.Parallel()
ur := UserRole{
Role: Role{ID: 1, Name: "admin"},
User: validUserForRole,
}
assert.NoError(t, ur.Valid())
})
t.Run("invalid role", func(t *testing.T) {
t.Parallel()
ur := UserRole{
Role: Role{Name: ""},
User: validUserForRole,
}
assert.Error(t, ur.Valid())
})
t.Run("invalid user", func(t *testing.T) {
t.Parallel()
ur := UserRole{
Role: Role{ID: 1, Name: "admin"},
User: User{Mail: ""},
}
assert.Error(t, ur.Valid())
})
}
func TestUserRolePrivilegesValid(t *testing.T) {
t.Parallel()
validUserForRole := User{
Hash: "abcdef1234567890abcdef1234567890",
Type: UserTypeLocal,
Mail: "test@example.com",
Status: UserStatusActive,
RoleID: RoleUser,
}
t.Run("valid user role privileges", func(t *testing.T) {
t.Parallel()
urp := UserRolePrivileges{
Role: RolePrivileges{
Privileges: []Privilege{{Name: "read"}},
Role: Role{ID: 1, Name: "admin"},
},
User: validUserForRole,
}
assert.NoError(t, urp.Valid())
})
t.Run("invalid role privileges", func(t *testing.T) {
t.Parallel()
urp := UserRolePrivileges{
Role: RolePrivileges{
Privileges: []Privilege{{Name: ""}},
Role: Role{ID: 1, Name: "admin"},
},
User: validUserForRole,
}
assert.Error(t, urp.Valid())
})
t.Run("invalid user", func(t *testing.T) {
t.Parallel()
urp := UserRolePrivileges{
Role: RolePrivileges{
Privileges: []Privilege{{Name: "read"}},
Role: Role{ID: 1, Name: "admin"},
},
User: User{Mail: ""},
}
assert.Error(t, urp.Valid())
})
}
+69
View File
@@ -0,0 +1,69 @@
package models
import (
"fmt"
"time"
"github.com/jinzhu/gorm"
)
type VecstoreActionType string
const (
VecstoreActionTypeRetrieve VecstoreActionType = "retrieve"
VecstoreActionTypeStore VecstoreActionType = "store"
)
func (s VecstoreActionType) String() string {
return string(s)
}
// Valid is function to control input/output data
func (s VecstoreActionType) Valid() error {
switch s {
case VecstoreActionTypeRetrieve, VecstoreActionTypeStore:
return nil
default:
return fmt.Errorf("invalid VecstoreActionType: %s", s)
}
}
// Validate is function to use callback to control input/output data
func (s VecstoreActionType) Validate(db *gorm.DB) {
if err := s.Valid(); err != nil {
db.AddError(err)
}
}
// Vecstorelog is model to contain vecstore action information
// nolint:lll
type Vecstorelog struct {
ID uint64 `form:"id" json:"id" validate:"min=0,numeric" gorm:"type:BIGINT;NOT NULL;PRIMARY_KEY;AUTO_INCREMENT"`
Initiator MsgchainType `json:"initiator" validate:"valid,required" gorm:"type:MSGCHAIN_TYPE;NOT NULL"`
Executor MsgchainType `json:"executor" validate:"valid,required" gorm:"type:MSGCHAIN_TYPE;NOT NULL"`
Filter string `json:"filter" validate:"required" gorm:"type:JSON;NOT NULL"`
Query string `json:"query" validate:"required" gorm:"type:TEXT;NOT NULL"`
Action VecstoreActionType `json:"action" validate:"valid,required" gorm:"type:VECSTORE_ACTION_TYPE;NOT NULL"`
Result string `json:"result" validate:"omitempty" gorm:"type:TEXT;NOT NULL"`
FlowID uint64 `form:"flow_id" json:"flow_id" validate:"min=0,numeric,required" gorm:"type:BIGINT;NOT NULL"`
TaskID *uint64 `form:"task_id,omitempty" json:"task_id,omitempty" validate:"omitnil,min=0" gorm:"type:BIGINT;NOT NULL"`
SubtaskID *uint64 `form:"subtask_id,omitempty" json:"subtask_id,omitempty" validate:"omitnil,min=0" gorm:"type:BIGINT;NOT NULL"`
CreatedAt time.Time `form:"created_at,omitempty" json:"created_at,omitempty" validate:"omitempty" gorm:"type:TIMESTAMPTZ;default:CURRENT_TIMESTAMP"`
}
// TableName returns the table name string to guaranty use correct table
func (ml *Vecstorelog) TableName() string {
return "vecstorelogs"
}
// Valid is function to control input/output data
func (ml Vecstorelog) Valid() error {
return validate.Struct(ml)
}
// Validate is function to use callback to control input/output data
func (ml Vecstorelog) Validate(db *gorm.DB) {
if err := ml.Valid(); err != nil {
db.AddError(err)
}
}
+64
View File
@@ -0,0 +1,64 @@
package oauth
import (
"context"
"fmt"
"golang.org/x/oauth2"
)
type OAuthEmailResolver func(ctx context.Context, nonce string, token *oauth2.Token) (string, error)
type OAuthClient interface {
ProviderName() string
ResolveEmail(ctx context.Context, nonce string, token *oauth2.Token) (string, error)
TokenSource(ctx context.Context, token *oauth2.Token) oauth2.TokenSource
Exchange(ctx context.Context, code string, opts ...oauth2.AuthCodeOption) (*oauth2.Token, error)
RefreshToken(ctx context.Context, token string) (*oauth2.Token, error)
AuthCodeURL(state string, opts ...oauth2.AuthCodeOption) string
}
type oauthClient struct {
name string
verifier string
conf *oauth2.Config
emailResolver OAuthEmailResolver
}
func NewOAuthClient(name string, conf *oauth2.Config, emailResolver OAuthEmailResolver) OAuthClient {
return &oauthClient{
name: name,
verifier: oauth2.GenerateVerifier(),
conf: conf,
emailResolver: emailResolver,
}
}
func (o *oauthClient) ProviderName() string {
return o.name
}
func (o *oauthClient) ResolveEmail(ctx context.Context, nonce string, token *oauth2.Token) (string, error) {
if o.emailResolver == nil {
return "", fmt.Errorf("email resolver is not set")
}
return o.emailResolver(ctx, nonce, token)
}
func (o *oauthClient) TokenSource(ctx context.Context, token *oauth2.Token) oauth2.TokenSource {
return o.conf.TokenSource(ctx, token)
}
func (o *oauthClient) Exchange(ctx context.Context, code string, opts ...oauth2.AuthCodeOption) (*oauth2.Token, error) {
opts = append(opts, oauth2.VerifierOption(o.verifier))
return o.conf.Exchange(ctx, code, opts...)
}
func (o *oauthClient) RefreshToken(ctx context.Context, token string) (*oauth2.Token, error) {
return o.conf.TokenSource(ctx, &oauth2.Token{RefreshToken: token}).Token()
}
func (o *oauthClient) AuthCodeURL(state string, opts ...oauth2.AuthCodeOption) string {
opts = append(opts, oauth2.S256ChallengeOption(o.verifier))
return o.conf.AuthCodeURL(state, opts...)
}
+70
View File
@@ -0,0 +1,70 @@
package oauth
import (
"context"
"encoding/json"
"fmt"
"io"
"net/http"
"golang.org/x/oauth2"
"golang.org/x/oauth2/github"
)
type githubEmail struct {
Email string `json:"email"`
Primary bool `json:"primary"`
Verified bool `json:"verified"`
Visibility string `json:"visibility"`
}
func githubEmailResolver(ctx context.Context, nonce string, token *oauth2.Token) (string, error) {
req, err := http.NewRequestWithContext(ctx, http.MethodGet, "https://api.github.com/user/emails", nil)
if err != nil {
return "", err
}
req.Header.Set("Authorization", fmt.Sprintf("token %s", token.AccessToken))
resp, err := http.DefaultClient.Do(req)
if err != nil {
return "", err
}
defer resp.Body.Close()
body, err := io.ReadAll(resp.Body)
if err != nil {
return "", err
}
emails := []githubEmail{}
if err := json.Unmarshal(body, &emails); err != nil {
return "", err
}
for _, email := range emails {
if email.Verified && email.Primary {
return email.Email, nil
}
}
for _, email := range emails {
if email.Verified {
return email.Email, nil
}
}
return "", fmt.Errorf("no verified primary email found")
}
func NewGithubOAuthClient(clientID, clientSecret, redirectURL string) OAuthClient {
return NewOAuthClient("github", &oauth2.Config{
ClientID: clientID,
ClientSecret: clientSecret,
RedirectURL: redirectURL,
Scopes: []string{
"user:email",
"openid",
},
Endpoint: github.Endpoint,
}, githubEmailResolver)
}
+76
View File
@@ -0,0 +1,76 @@
package oauth
import (
"context"
"fmt"
"github.com/coreos/go-oidc/v3/oidc"
"golang.org/x/oauth2"
"golang.org/x/oauth2/google"
)
type googleTokenClaims struct {
Nonce string `json:"nonce"`
Email string `json:"email"`
EmailVerified bool `json:"email_verified"`
}
func newGoogleEmailResolver(clientID string) OAuthEmailResolver {
return func(ctx context.Context, nonce string, token *oauth2.Token) (string, error) {
provider, err := oidc.NewProvider(ctx, "https://accounts.google.com")
if err != nil {
return "", fmt.Errorf("could not create Google OpenID client: %w", err)
}
oidToken, ok := token.Extra("id_token").(string)
if !ok {
return "", fmt.Errorf("id_token is not present in the token")
}
verifier := provider.Verifier(&oidc.Config{ClientID: clientID})
idToken, err := verifier.Verify(ctx, oidToken)
if err != nil {
return "", fmt.Errorf("could not verify Google ID Token: %w", err)
}
if idToken.Nonce != nonce {
return "", fmt.Errorf("nonce mismatch in Google ID Token")
}
if err = idToken.VerifyAccessToken(token.AccessToken); err != nil {
return "", fmt.Errorf("failed to verify Google Access Token: %w", err)
}
claims := googleTokenClaims{}
if err := idToken.Claims(&claims); err != nil {
return "", fmt.Errorf("failed to parse Google ID Token claims: %w", err)
}
if claims.Nonce != nonce {
return "", fmt.Errorf("nonce mismatch in Google ID Token claims")
}
if !claims.EmailVerified {
return "", fmt.Errorf("email not verified in Google ID Token claims")
}
if claims.Email == "" {
return "", fmt.Errorf("email is empty in Google ID Token claims")
}
return claims.Email, nil
}
}
func NewGoogleOAuthClient(clientID, clientSecret, redirectURL string) OAuthClient {
return NewOAuthClient("google", &oauth2.Config{
ClientID: clientID,
ClientSecret: clientSecret,
RedirectURL: redirectURL,
Scopes: []string{
"https://www.googleapis.com/auth/userinfo.email",
"openid",
},
Endpoint: google.Endpoint,
}, newGoogleEmailResolver(clientID))
}
+370
View File
@@ -0,0 +1,370 @@
package rdb
import (
"crypto/md5" //nolint:gosec
"encoding/hex"
"errors"
"slices"
"strconv"
"strings"
"time"
"github.com/google/uuid"
"github.com/jinzhu/gorm"
"golang.org/x/crypto/bcrypt"
)
// TableFilter is auxiliary struct to contain method of filtering
//
//nolint:lll
type TableFilter struct {
Value any `form:"value" json:"value,omitempty" binding:"required" swaggertype:"object"`
Field string `form:"field" json:"field,omitempty" binding:"required"`
Operator string `form:"operator" json:"operator,omitempty" binding:"oneof='<' '<=' '>=' '>' '=' '!=' 'like' 'not like' 'in',omitempty" default:"like" enums:"<,<=,>=,>,=,!=,like,not like,in"`
}
// TableSort is auxiliary struct to contain method of sorting
type TableSort struct {
Prop string `form:"prop" json:"prop,omitempty" binding:"omitempty"`
Order string `form:"order" json:"order,omitempty" binding:"oneof=ascending descending,required_with=Prop,omitempty" enums:"ascending,descending"`
}
// TableQuery is main struct to contain input params
//
//nolint:lll
type TableQuery struct {
// Number of page (since 1)
Page int `form:"page" json:"page" binding:"min=1,required" default:"1" minimum:"1"`
// Amount items per page (min -1, max 1000, -1 means unlimited)
Size int `form:"pageSize" json:"pageSize" binding:"min=-1,max=1000" default:"5" minimum:"-1" maximum:"1000"`
// Type of request
Type string `form:"type" json:"type" binding:"oneof=sort filter init page size,required" default:"init" enums:"sort,filter,init,page,size"`
// Sorting result on server e.g. {"prop":"...","order":"..."}
// field order is "ascending" or "descending" value
// order is required if prop is not empty
Sort []TableSort `form:"sort[]" json:"sort[],omitempty" binding:"omitempty,dive" swaggertype:"array,string"`
// Filtering result on server e.g. {"value":[...],"field":"...","operator":"..."}
// field is the unique identifier of the table column, different for each endpoint
// value should be integer or string or array type, "value":123 or "value":"string" or "value":[123,456]
// operator value should be one of <,<=,>=,>,=,!=,like,not like,in
// default operator value is 'like' or '=' if field is 'id' or '*_id' or '*_at'
Filters []TableFilter `form:"filters[]" json:"filters[],omitempty" binding:"omitempty,dive" swaggertype:"array,string"`
// Field to group results by
Group string `form:"group" json:"group,omitempty" binding:"omitempty" swaggertype:"string"`
// non input arguments
table string `form:"-" json:"-"`
groupField string `form:"-" json:"-"`
sqlMappers map[string]any `form:"-" json:"-"`
sqlFind func(out any) func(*gorm.DB) *gorm.DB `form:"-" json:"-"`
sqlFilters []func(*gorm.DB) *gorm.DB `form:"-" json:"-"`
sqlOrders []func(*gorm.DB) *gorm.DB `form:"-" json:"-"`
}
// Init is function to set table name and sql mapping to data columns
func (q *TableQuery) Init(table string, sqlMappers map[string]any) error {
q.table = table
q.sqlFind = func(out any) func(db *gorm.DB) *gorm.DB {
return func(db *gorm.DB) *gorm.DB {
return db.Find(out)
}
}
q.sqlMappers = make(map[string]any)
q.sqlOrders = append(q.sqlOrders, func(db *gorm.DB) *gorm.DB {
return db.Order("id DESC")
})
for k, v := range sqlMappers {
switch t := v.(type) {
case string:
t = q.DoConditionFormat(t)
if isNumbericField(k) {
q.sqlMappers[k] = t
} else {
q.sqlMappers[k] = "LOWER(" + t + "::text)"
}
case func(q *TableQuery, db *gorm.DB, value any) *gorm.DB:
q.sqlMappers[k] = t
default:
continue
}
}
if q.Group != "" {
var ok bool
q.groupField, ok = q.sqlMappers[q.Group].(string)
if !ok {
return errors.New("wrong field for grouping")
}
}
return nil
}
// DoConditionFormat is auxiliary function to prepare condition to the table
func (q *TableQuery) DoConditionFormat(cond string) string {
cond = strings.ReplaceAll(cond, "{{type}}", q.Type)
cond = strings.ReplaceAll(cond, "{{table}}", q.table)
cond = strings.ReplaceAll(cond, "{{page}}", strconv.Itoa(q.Page))
cond = strings.ReplaceAll(cond, "{{size}}", strconv.Itoa(q.Size))
return cond
}
// SetFilters is function to set custom filters to build target SQL query
func (q *TableQuery) SetFilters(sqlFilters []func(*gorm.DB) *gorm.DB) {
q.sqlFilters = sqlFilters
}
// SetFind is function to set custom find function to build target SQL query
func (q *TableQuery) SetFind(find func(out any) func(*gorm.DB) *gorm.DB) {
q.sqlFind = find
}
// SetOrders is function to set custom ordering to build target SQL query
func (q *TableQuery) SetOrders(sqlOrders []func(*gorm.DB) *gorm.DB) {
q.sqlOrders = sqlOrders
}
// Mappers is getter for private field (SQL find funcction to use it in custom query)
func (q *TableQuery) Find(out any) func(*gorm.DB) *gorm.DB {
return q.sqlFind(out)
}
// Mappers is getter for private field (SQL mappers fields to table ones)
func (q *TableQuery) Mappers() map[string]any {
return q.sqlMappers
}
// Table is getter for private field (table name)
func (q *TableQuery) Table() string {
return q.table
}
// Ordering is function to get order of data rows according with input params
func (q *TableQuery) Ordering() func(db *gorm.DB) *gorm.DB {
var sortItems []TableSort
for _, sort := range q.Sort {
var t TableSort
switch sort.Order {
case "ascending":
t.Order = "ASC"
case "descending":
t.Order = "DESC"
}
if v, ok := q.sqlMappers[sort.Prop]; ok {
if s, ok := v.(string); ok {
t.Prop = s
}
}
if t.Prop != "" && t.Order != "" {
sortItems = append(sortItems, t)
}
}
return func(db *gorm.DB) *gorm.DB {
for _, sort := range sortItems {
// sort.Prop comes from server-side whitelist (q.sqlMappers)
// sort.Order is validated to be only "ASC" or "DESC"
db = db.Order(sort.Prop + " " + sort.Order)
}
for _, order := range q.sqlOrders {
db = order(db)
}
return db
}
}
// Paginate is function to navigate between pages according with input params
func (q *TableQuery) Paginate() func(db *gorm.DB) *gorm.DB {
return func(db *gorm.DB) *gorm.DB {
if q.Page <= 0 && q.Size >= 0 {
return db.Limit(q.Size)
} else if q.Page > 0 && q.Size >= 0 {
offset := (q.Page - 1) * q.Size
return db.Offset(offset).Limit(q.Size)
}
return db
}
}
// GroupBy is function to group results by some field
func (q *TableQuery) GroupBy(total *uint64, result any) func(db *gorm.DB) *gorm.DB {
return func(db *gorm.DB) *gorm.DB {
return db.Group(q.groupField).Where(q.groupField+" IS NOT NULL").Count(total).Pluck(q.groupField, result)
}
}
// DataFilter is function to build main data filter from filters input params
func (q *TableQuery) DataFilter() func(db *gorm.DB) *gorm.DB {
type item struct {
op string
v any
}
fl := make(map[string][]item)
setFilter := func(field, operator string, value any) {
if operator == "" {
operator = "like" // nolint:goconst
}
fvalue := []item{}
if fv, ok := fl[field]; ok {
fvalue = fv
}
switch tvalue := value.(type) {
case string, float64, bool:
fl[field] = append(fvalue, item{operator, tvalue})
case []any:
fl[field] = append(fvalue, item{operator, tvalue})
}
}
patchOperator := func(f *TableFilter) {
switch f.Operator {
case "<", "<=", ">=", ">", "=", "!=", "in":
case "not like":
if isNumbericField(f.Field) {
f.Operator = "!="
}
default:
f.Operator = "like"
if isNumbericField(f.Field) {
f.Operator = "="
}
}
}
for _, f := range q.Filters {
f := f
patchOperator(&f)
if _, ok := q.sqlMappers[f.Field]; ok {
if v, ok := f.Value.(string); ok && v != "" {
vs := v
if slices.Contains([]string{"like", "not like"}, f.Operator) {
vs = "%" + strings.ToLower(vs) + "%"
}
setFilter(f.Field, f.Operator, vs)
}
if v, ok := f.Value.(float64); ok {
setFilter(f.Field, f.Operator, v)
}
if v, ok := f.Value.(bool); ok {
setFilter(f.Field, f.Operator, v)
}
if v, ok := f.Value.([]any); ok && len(v) != 0 {
var vi []any
for _, ti := range v {
if ts, ok := ti.(string); ok {
vi = append(vi, strings.ToLower(ts))
}
if ts, ok := ti.(float64); ok {
vi = append(vi, ts)
}
if ts, ok := ti.(bool); ok {
vi = append(vi, ts)
}
}
if len(vi) != 0 {
setFilter(f.Field, "in", vi)
}
}
}
}
return func(db *gorm.DB) *gorm.DB {
doFilter := func(db *gorm.DB, k, s string, v any) *gorm.DB {
switch t := q.sqlMappers[k].(type) {
case string:
return db.Where(t+s, v)
case func(q *TableQuery, db *gorm.DB, value any) *gorm.DB:
return t(q, db, v)
default:
return db
}
}
for k, f := range fl {
for _, it := range f {
if _, ok := it.v.([]any); ok {
db = doFilter(db, k, " "+it.op+" (?)", it.v)
} else {
db = doFilter(db, k, " "+it.op+" ?", it.v)
}
}
}
for _, filter := range q.sqlFilters {
db = filter(db)
}
return db
}
}
// Query is function to retrieve table data according with input params
func (q *TableQuery) Query(db *gorm.DB, result any,
funcs ...func(*gorm.DB) *gorm.DB) (uint64, error) {
var total uint64
err := ApplyToChainDB(
ApplyToChainDB(db.Table(q.Table()), funcs...).Scopes(q.DataFilter()).Count(&total),
q.Ordering(),
q.Paginate(),
q.Find(result),
).Error
return uint64(total), err
}
// QueryGrouped is function to retrieve grouped data according with input params
func (q *TableQuery) QueryGrouped(db *gorm.DB, result any,
funcs ...func(*gorm.DB) *gorm.DB) (uint64, error) {
if _, ok := q.sqlMappers[q.Group]; !ok {
return 0, errors.New("group field not found")
}
var total uint64
err := ApplyToChainDB(
ApplyToChainDB(db.Table(q.Table()), funcs...).Scopes(q.DataFilter()),
q.GroupBy(&total, result),
).Error
return uint64(total), err
}
// ApplyToChainDB is function to extend gorm method chaining by custom functions
func ApplyToChainDB(db *gorm.DB, funcs ...func(*gorm.DB) *gorm.DB) (tx *gorm.DB) {
for _, f := range funcs {
db = f(db)
}
return db
}
// EncryptPassword is function to prepare user data as a password
func EncryptPassword(password string) (hpass []byte, err error) {
hpass, err = bcrypt.GenerateFromPassword([]byte(password), bcrypt.DefaultCost)
return
}
// MakeMD5Hash is function to generate common hash by value
func MakeMD5Hash(value, salt string) string {
currentTime := time.Now().Format("2006-01-02 15:04:05.000000000")
hash := md5.Sum([]byte(currentTime + value + salt)) // nolint:gosec
return hex.EncodeToString(hash[:])
}
// MakeUserHash is function to generate user hash from name
func MakeUserHash(name string) string {
currentTime := time.Now().Format("2006-01-02 15:04:05.000000000")
return MakeMD5Hash(name+currentTime, "248a8bd896595be1319e65c308a903c568afdb9b")
}
// MakeUuidStrFromHash is function to convert format view from hash to UUID
func MakeUuidStrFromHash(hash string) (string, error) {
hashBytes, err := hex.DecodeString(hash)
if err != nil {
return "", err
}
userIdUuid, err := uuid.FromBytes(hashBytes)
if err != nil {
return "", err
}
return userIdUuid.String(), nil
}
func isNumbericField(field string) bool {
return strings.HasSuffix(field, "_id") || strings.HasSuffix(field, "_at") || field == "id"
}
+171
View File
@@ -0,0 +1,171 @@
package response
// general
var ErrInternal = NewHttpError(500, "Internal", "internal server error")
var ErrInternalDBNotFound = NewHttpError(500, "Internal.DBNotFound", "db not found")
var ErrInternalServiceNotFound = NewHttpError(500, "Internal.ServiceNotFound", "service not found")
var ErrInternalDBEncryptorNotFound = NewHttpError(500, "Internal.DBEncryptorNotFound", "DBEncryptor not found")
var ErrNotPermitted = NewHttpError(403, "NotPermitted", "action not permitted")
var ErrAuthRequired = NewHttpError(403, "AuthRequired", "auth required")
var ErrLocalUserRequired = NewHttpError(403, "LocalUserRequired", "local user required")
var ErrPrivilegesRequired = NewHttpError(403, "PrivilegesRequired", "some privileges required")
var ErrAdminRequired = NewHttpError(403, "AdminRequired", "admin required")
var ErrSuperRequired = NewHttpError(403, "SuperRequired", "super admin required")
// auth
var ErrAuthInvalidLoginRequest = NewHttpError(400, "Auth.InvalidLoginRequest", "invalid login data")
var ErrAuthInvalidAuthorizeQuery = NewHttpError(400, "Auth.InvalidAuthorizeQuery", "invalid authorize query")
var ErrAuthInvalidLoginCallbackRequest = NewHttpError(400, "Auth.InvalidLoginCallbackRequest", "invalid login callback data")
var ErrAuthInvalidAuthorizationState = NewHttpError(400, "Auth.InvalidAuthorizationState", "invalid authorization state data")
var ErrAuthInvalidSwitchServiceHash = NewHttpError(400, "Auth.InvalidSwitchServiceHash", "invalid switch service hash input data")
var ErrAuthInvalidAuthorizationNonce = NewHttpError(400, "Auth.InvalidAuthorizationNonce", "invalid authorization nonce data")
var ErrAuthInvalidCredentials = NewHttpError(401, "Auth.InvalidCredentials", "invalid login or password")
var ErrAuthInvalidUserData = NewHttpError(500, "Auth.InvalidUserData", "invalid user data")
var ErrAuthInactiveUser = NewHttpError(403, "Auth.InactiveUser", "user is inactive")
var ErrAuthExchangeTokenFail = NewHttpError(403, "Auth.ExchangeTokenFail", "error on exchanging token")
var ErrAuthTokenExpired = NewHttpError(403, "Auth.TokenExpired", "token is expired")
var ErrAuthVerificationTokenFail = NewHttpError(403, "Auth.VerificationTokenFail", "error on verifying token")
var ErrAuthInvalidServiceData = NewHttpError(500, "Auth.InvalidServiceData", "invalid service data")
var ErrAuthInvalidTenantData = NewHttpError(500, "Auth.InvalidTenantData", "invalid tenant data")
// info
var ErrInfoUserNotFound = NewHttpError(404, "Info.UserNotFound", "user not found")
var ErrInfoInvalidUserData = NewHttpError(500, "Info.InvalidUserData", "invalid user data")
var ErrInfoInvalidServiceData = NewHttpError(500, "Info.InvalidServiceData", "invalid service data")
// users
var ErrUsersNotFound = NewHttpError(404, "Users.NotFound", "user not found")
var ErrUsersInvalidData = NewHttpError(500, "Users.InvalidData", "invalid user data")
var ErrUsersInvalidRequest = NewHttpError(400, "Users.InvalidRequest", "invalid user request data")
var ErrChangePasswordCurrentUserInvalidPassword = NewHttpError(400, "Users.ChangePasswordCurrentUser.InvalidPassword", "failed to validate user password")
var ErrChangePasswordCurrentUserInvalidCurrentPassword = NewHttpError(403, "Users.ChangePasswordCurrentUser.InvalidCurrentPassword", "invalid current password")
var ErrChangePasswordCurrentUserInvalidNewPassword = NewHttpError(400, "Users.ChangePasswordCurrentUser.InvalidNewPassword", "invalid new password form data")
var ErrGetUserModelsNotFound = NewHttpError(404, "Users.GetUser.ModelsNotFound", "user linked models not found")
var ErrCreateUserInvalidUser = NewHttpError(400, "Users.CreateUser.InvalidUser", "failed to validate user")
var ErrPatchUserModelsNotFound = NewHttpError(404, "Users.PatchUser.ModelsNotFound", "user linked models not found")
var ErrDeleteUserModelsNotFound = NewHttpError(404, "Users.DeleteUser.ModelsNotFound", "user linked models not found")
// roles
var ErrRolesInvalidRequest = NewHttpError(400, "Roles.InvalidRequest", "invalid role request data")
var ErrRolesInvalidData = NewHttpError(500, "Roles.InvalidData", "invalid role data")
var ErrRolesNotFound = NewHttpError(404, "Roles.NotFound", "role not found")
// prompts
var ErrPromptsInvalidRequest = NewHttpError(400, "Prompts.InvalidRequest", "invalid prompt request data")
var ErrPromptsInvalidData = NewHttpError(500, "Prompts.InvalidData", "invalid prompt data")
var ErrPromptsNotFound = NewHttpError(404, "Prompts.NotFound", "prompt not found")
// screenshots
var ErrScreenshotsInvalidRequest = NewHttpError(400, "Screenshots.InvalidRequest", "invalid screenshot request data")
var ErrScreenshotsNotFound = NewHttpError(404, "Screenshots.NotFound", "screenshot not found")
var ErrScreenshotsInvalidData = NewHttpError(500, "Screenshots.InvalidData", "invalid screenshot data")
// containers
var ErrContainersInvalidRequest = NewHttpError(400, "Containers.InvalidRequest", "invalid container request data")
var ErrContainersNotFound = NewHttpError(404, "Containers.NotFound", "container not found")
var ErrContainersInvalidData = NewHttpError(500, "Containers.InvalidData", "invalid container data")
// agentlogs
var ErrAgentlogsInvalidRequest = NewHttpError(400, "Agentlogs.InvalidRequest", "invalid agentlog request data")
var ErrAgentlogsInvalidData = NewHttpError(500, "Agentlogs.InvalidData", "invalid agentlog data")
// assistantlogs
var ErrAssistantlogsInvalidRequest = NewHttpError(400, "Assistantlogs.InvalidRequest", "invalid assistantlog request data")
var ErrAssistantlogsInvalidData = NewHttpError(500, "Assistantlogs.InvalidData", "invalid assistantlog data")
// msglogs
var ErrMsglogsInvalidRequest = NewHttpError(400, "Msglogs.InvalidRequest", "invalid msglog request data")
var ErrMsglogsInvalidData = NewHttpError(500, "Msglogs.InvalidData", "invalid msglog data")
// searchlogs
var ErrSearchlogsInvalidRequest = NewHttpError(400, "Searchlogs.InvalidRequest", "invalid searchlog request data")
var ErrSearchlogsInvalidData = NewHttpError(500, "Searchlogs.InvalidData", "invalid searchlog data")
// termlogs
var ErrTermlogsInvalidRequest = NewHttpError(400, "Termlogs.InvalidRequest", "invalid termlog request data")
var ErrTermlogsInvalidData = NewHttpError(500, "Termlogs.InvalidData", "invalid termlog data")
// vecstorelogs
var ErrVecstorelogsInvalidRequest = NewHttpError(400, "Vecstorelogs.InvalidRequest", "invalid vecstorelog request data")
var ErrVecstorelogsInvalidData = NewHttpError(500, "Vecstorelogs.InvalidData", "invalid vecstorelog data")
// flows
var ErrFlowsInvalidRequest = NewHttpError(400, "Flows.InvalidRequest", "invalid flow request data")
var ErrFlowsNotFound = NewHttpError(404, "Flows.NotFound", "flow not found")
var ErrFlowsInvalidData = NewHttpError(500, "Flows.InvalidData", "invalid flow data")
// flow files
var ErrFlowFilesInvalidRequest = NewHttpError(400, "FlowFiles.InvalidRequest", "invalid flow file request data")
var ErrFlowFilesNotFound = NewHttpError(404, "FlowFiles.NotFound", "flow file not found")
var ErrFlowFilesInvalidData = NewHttpError(400, "FlowFiles.InvalidData", "invalid flow file data")
var ErrFlowFilesAlreadyExists = NewHttpError(409, "FlowFiles.AlreadyExists", "flow file already exists")
var ErrFlowFilesContainerNotRunning = NewHttpError(400, "FlowFiles.ContainerNotRunning", "container is not running")
// tasks
var ErrTasksInvalidRequest = NewHttpError(400, "Tasks.InvalidRequest", "invalid task request data")
var ErrTasksNotFound = NewHttpError(404, "Tasks.NotFound", "task not found")
var ErrTasksInvalidData = NewHttpError(500, "Tasks.InvalidData", "invalid task data")
// subtasks
var ErrSubtasksInvalidRequest = NewHttpError(400, "Subtasks.InvalidRequest", "invalid subtask request data")
var ErrSubtasksNotFound = NewHttpError(404, "Subtasks.NotFound", "subtask not found")
var ErrSubtasksInvalidData = NewHttpError(500, "Subtasks.InvalidData", "invalid subtask data")
// assistants
var ErrAssistantsInvalidRequest = NewHttpError(400, "Assistants.InvalidRequest", "invalid assistant request data")
var ErrAssistantsNotFound = NewHttpError(404, "Assistants.NotFound", "assistant not found")
var ErrAssistantsInvalidData = NewHttpError(500, "Assistants.InvalidData", "invalid assistant data")
// resources
var ErrResourcesInvalidRequest = NewHttpError(400, "Resources.InvalidRequest", "invalid resource request data")
var ErrResourcesNotFound = NewHttpError(404, "Resources.NotFound", "resource not found")
var ErrResourcesAlreadyExists = NewHttpError(409, "Resources.AlreadyExists", "resource already exists")
var ErrResourcesInvalidData = NewHttpError(400, "Resources.InvalidData", "invalid resource data")
var ErrResourcesConflict = NewHttpError(409, "Resources.Conflict", "resource conflict: use force=true to merge")
// knowledge
var ErrKnowledgeInvalidRequest = NewHttpError(400, "Knowledge.InvalidRequest", "invalid knowledge request data")
var ErrKnowledgeNotFound = NewHttpError(404, "Knowledge.NotFound", "knowledge document not found")
var ErrKnowledgeUnauthorized = NewHttpError(403, "Knowledge.Unauthorized", "not authorized to manage this knowledge document")
var ErrKnowledgeStoreUnavail = NewHttpError(503, "Knowledge.StoreUnavailable", "knowledge store (embedding provider) is not configured")
var ErrKnowledgeInvalidData = NewHttpError(500, "Knowledge.InvalidData", "invalid knowledge document data")
// toolcalls
var ErrToolcallsInvalidRequest = NewHttpError(400, "Toolcalls.InvalidRequest", "invalid toolcall request data")
var ErrToolcallsNotFound = NewHttpError(404, "Toolcalls.NotFound", "toolcall not found")
var ErrToolcallsInvalidData = NewHttpError(500, "Toolcalls.InvalidData", "invalid toolcall data")
// anonymize
var ErrAnonymizeInvalidRequest = NewHttpError(400, "Anonymize.InvalidRequest", "invalid anonymize request data")
var ErrAnonymizeUnavailable = NewHttpError(503, "Anonymize.Unavailable", "anonymizer is not configured")
// tokens
var ErrTokenCreationDisabled = NewHttpError(400, "Token.CreationDisabled", "token creation is disabled with default configuration")
var ErrTokenNotFound = NewHttpError(404, "Token.NotFound", "token not found")
var ErrTokenUnauthorized = NewHttpError(403, "Token.Unauthorized", "not authorized to manage this token")
var ErrTokenInvalidRequest = NewHttpError(400, "Token.InvalidRequest", "invalid token request data")
var ErrTokenInvalidData = NewHttpError(500, "Token.InvalidData", "invalid token data")
+75
View File
@@ -0,0 +1,75 @@
package response
import (
"fmt"
"pentagi/pkg/server/logger"
"pentagi/pkg/version"
"github.com/gin-gonic/gin"
"github.com/sirupsen/logrus"
)
type HttpError struct {
message string
code string
httpCode int
}
func (h *HttpError) Code() string {
return h.code
}
func (h *HttpError) HttpCode() int {
return h.httpCode
}
func (h *HttpError) Msg() string {
return h.message
}
func NewHttpError(httpCode int, code, message string) *HttpError {
return &HttpError{httpCode: httpCode, message: message, code: code}
}
func (h *HttpError) Error() string {
return fmt.Sprintf("%s: %s", h.code, h.message)
}
func Error(c *gin.Context, err *HttpError, original error) {
body := gin.H{
"status": "error",
"code": err.Code(),
"msg": err.Msg(),
}
if version.IsDevelopMode() && original != nil {
body["error"] = original.Error()
}
fields := logrus.Fields{
"code": err.HttpCode(),
"message": err.Msg(),
}
logger.FromContext(c).WithFields(fields).WithError(original).Error("api error")
c.AbortWithStatusJSON(err.HttpCode(), body)
}
func Success(c *gin.Context, code int, data any) {
c.JSON(code, gin.H{"status": "success", "data": data})
}
//lint:ignore U1000 successResp
type successResp struct {
Status string `json:"status" example:"success"`
Data any `json:"data" swaggertype:"object"`
} // @name SuccessResponse
//lint:ignore U1000 errorResp
type errorResp struct {
Status string `json:"status" example:"error"`
Code string `json:"code" example:"Internal"`
Msg string `json:"msg,omitempty" example:"internal server error"`
Error string `json:"error,omitempty" example:"original server error message"`
} // @name ErrorResponse
+434
View File
@@ -0,0 +1,434 @@
package response
import (
"encoding/json"
"errors"
"net/http"
"net/http/httptest"
"testing"
"pentagi/pkg/version"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
func init() {
gin.SetMode(gin.TestMode)
}
func TestNewHttpError(t *testing.T) {
t.Parallel()
err := NewHttpError(404, "NotFound", "resource not found")
assert.Equal(t, 404, err.HttpCode())
assert.Equal(t, "NotFound", err.Code())
assert.Equal(t, "resource not found", err.Msg())
}
func TestHttpError_Error(t *testing.T) {
t.Parallel()
err := NewHttpError(500, "Internal", "something broke")
assert.Equal(t, "Internal: something broke", err.Error())
}
func TestHttpError_ImplementsError(t *testing.T) {
t.Parallel()
var err error = NewHttpError(400, "Bad", "bad request")
assert.Error(t, err)
assert.Contains(t, err.Error(), "Bad")
}
func TestPredefinedErrors(t *testing.T) {
t.Parallel()
tests := []struct {
name string
err *HttpError
httpCode int
code string
}{
// General errors
{"ErrInternal", ErrInternal, 500, "Internal"},
{"ErrInternalDBNotFound", ErrInternalDBNotFound, 500, "Internal.DBNotFound"},
{"ErrInternalServiceNotFound", ErrInternalServiceNotFound, 500, "Internal.ServiceNotFound"},
{"ErrInternalDBEncryptorNotFound", ErrInternalDBEncryptorNotFound, 500, "Internal.DBEncryptorNotFound"},
{"ErrNotPermitted", ErrNotPermitted, 403, "NotPermitted"},
{"ErrAuthRequired", ErrAuthRequired, 403, "AuthRequired"},
{"ErrLocalUserRequired", ErrLocalUserRequired, 403, "LocalUserRequired"},
{"ErrPrivilegesRequired", ErrPrivilegesRequired, 403, "PrivilegesRequired"},
{"ErrAdminRequired", ErrAdminRequired, 403, "AdminRequired"},
{"ErrSuperRequired", ErrSuperRequired, 403, "SuperRequired"},
// Auth errors
{"ErrAuthInvalidLoginRequest", ErrAuthInvalidLoginRequest, 400, "Auth.InvalidLoginRequest"},
{"ErrAuthInvalidAuthorizeQuery", ErrAuthInvalidAuthorizeQuery, 400, "Auth.InvalidAuthorizeQuery"},
{"ErrAuthInvalidLoginCallbackRequest", ErrAuthInvalidLoginCallbackRequest, 400, "Auth.InvalidLoginCallbackRequest"},
{"ErrAuthInvalidAuthorizationState", ErrAuthInvalidAuthorizationState, 400, "Auth.InvalidAuthorizationState"},
{"ErrAuthInvalidSwitchServiceHash", ErrAuthInvalidSwitchServiceHash, 400, "Auth.InvalidSwitchServiceHash"},
{"ErrAuthInvalidAuthorizationNonce", ErrAuthInvalidAuthorizationNonce, 400, "Auth.InvalidAuthorizationNonce"},
{"ErrAuthInvalidCredentials", ErrAuthInvalidCredentials, 401, "Auth.InvalidCredentials"},
{"ErrAuthInvalidUserData", ErrAuthInvalidUserData, 500, "Auth.InvalidUserData"},
{"ErrAuthInactiveUser", ErrAuthInactiveUser, 403, "Auth.InactiveUser"},
{"ErrAuthExchangeTokenFail", ErrAuthExchangeTokenFail, 403, "Auth.ExchangeTokenFail"},
{"ErrAuthTokenExpired", ErrAuthTokenExpired, 403, "Auth.TokenExpired"},
{"ErrAuthVerificationTokenFail", ErrAuthVerificationTokenFail, 403, "Auth.VerificationTokenFail"},
{"ErrAuthInvalidServiceData", ErrAuthInvalidServiceData, 500, "Auth.InvalidServiceData"},
{"ErrAuthInvalidTenantData", ErrAuthInvalidTenantData, 500, "Auth.InvalidTenantData"},
// Info errors
{"ErrInfoUserNotFound", ErrInfoUserNotFound, 404, "Info.UserNotFound"},
{"ErrInfoInvalidUserData", ErrInfoInvalidUserData, 500, "Info.InvalidUserData"},
{"ErrInfoInvalidServiceData", ErrInfoInvalidServiceData, 500, "Info.InvalidServiceData"},
// Users errors
{"ErrUsersNotFound", ErrUsersNotFound, 404, "Users.NotFound"},
{"ErrUsersInvalidData", ErrUsersInvalidData, 500, "Users.InvalidData"},
{"ErrUsersInvalidRequest", ErrUsersInvalidRequest, 400, "Users.InvalidRequest"},
{"ErrChangePasswordCurrentUserInvalidPassword", ErrChangePasswordCurrentUserInvalidPassword, 400, "Users.ChangePasswordCurrentUser.InvalidPassword"},
{"ErrChangePasswordCurrentUserInvalidCurrentPassword", ErrChangePasswordCurrentUserInvalidCurrentPassword, 403, "Users.ChangePasswordCurrentUser.InvalidCurrentPassword"},
{"ErrChangePasswordCurrentUserInvalidNewPassword", ErrChangePasswordCurrentUserInvalidNewPassword, 400, "Users.ChangePasswordCurrentUser.InvalidNewPassword"},
{"ErrGetUserModelsNotFound", ErrGetUserModelsNotFound, 404, "Users.GetUser.ModelsNotFound"},
{"ErrCreateUserInvalidUser", ErrCreateUserInvalidUser, 400, "Users.CreateUser.InvalidUser"},
{"ErrPatchUserModelsNotFound", ErrPatchUserModelsNotFound, 404, "Users.PatchUser.ModelsNotFound"},
{"ErrDeleteUserModelsNotFound", ErrDeleteUserModelsNotFound, 404, "Users.DeleteUser.ModelsNotFound"},
// Roles errors
{"ErrRolesInvalidRequest", ErrRolesInvalidRequest, 400, "Roles.InvalidRequest"},
{"ErrRolesInvalidData", ErrRolesInvalidData, 500, "Roles.InvalidData"},
{"ErrRolesNotFound", ErrRolesNotFound, 404, "Roles.NotFound"},
// Prompts errors
{"ErrPromptsInvalidRequest", ErrPromptsInvalidRequest, 400, "Prompts.InvalidRequest"},
{"ErrPromptsInvalidData", ErrPromptsInvalidData, 500, "Prompts.InvalidData"},
{"ErrPromptsNotFound", ErrPromptsNotFound, 404, "Prompts.NotFound"},
// Screenshots errors
{"ErrScreenshotsInvalidRequest", ErrScreenshotsInvalidRequest, 400, "Screenshots.InvalidRequest"},
{"ErrScreenshotsNotFound", ErrScreenshotsNotFound, 404, "Screenshots.NotFound"},
{"ErrScreenshotsInvalidData", ErrScreenshotsInvalidData, 500, "Screenshots.InvalidData"},
// Containers errors
{"ErrContainersInvalidRequest", ErrContainersInvalidRequest, 400, "Containers.InvalidRequest"},
{"ErrContainersNotFound", ErrContainersNotFound, 404, "Containers.NotFound"},
{"ErrContainersInvalidData", ErrContainersInvalidData, 500, "Containers.InvalidData"},
// Agentlogs errors
{"ErrAgentlogsInvalidRequest", ErrAgentlogsInvalidRequest, 400, "Agentlogs.InvalidRequest"},
{"ErrAgentlogsInvalidData", ErrAgentlogsInvalidData, 500, "Agentlogs.InvalidData"},
// Assistantlogs errors
{"ErrAssistantlogsInvalidRequest", ErrAssistantlogsInvalidRequest, 400, "Assistantlogs.InvalidRequest"},
{"ErrAssistantlogsInvalidData", ErrAssistantlogsInvalidData, 500, "Assistantlogs.InvalidData"},
// Msglogs errors
{"ErrMsglogsInvalidRequest", ErrMsglogsInvalidRequest, 400, "Msglogs.InvalidRequest"},
{"ErrMsglogsInvalidData", ErrMsglogsInvalidData, 500, "Msglogs.InvalidData"},
// Searchlogs errors
{"ErrSearchlogsInvalidRequest", ErrSearchlogsInvalidRequest, 400, "Searchlogs.InvalidRequest"},
{"ErrSearchlogsInvalidData", ErrSearchlogsInvalidData, 500, "Searchlogs.InvalidData"},
// Termlogs errors
{"ErrTermlogsInvalidRequest", ErrTermlogsInvalidRequest, 400, "Termlogs.InvalidRequest"},
{"ErrTermlogsInvalidData", ErrTermlogsInvalidData, 500, "Termlogs.InvalidData"},
// Vecstorelogs errors
{"ErrVecstorelogsInvalidRequest", ErrVecstorelogsInvalidRequest, 400, "Vecstorelogs.InvalidRequest"},
{"ErrVecstorelogsInvalidData", ErrVecstorelogsInvalidData, 500, "Vecstorelogs.InvalidData"},
// Flows errors
{"ErrFlowsInvalidRequest", ErrFlowsInvalidRequest, 400, "Flows.InvalidRequest"},
{"ErrFlowsNotFound", ErrFlowsNotFound, 404, "Flows.NotFound"},
{"ErrFlowsInvalidData", ErrFlowsInvalidData, 500, "Flows.InvalidData"},
// Tasks errors
{"ErrTasksInvalidRequest", ErrTasksInvalidRequest, 400, "Tasks.InvalidRequest"},
{"ErrTasksNotFound", ErrTasksNotFound, 404, "Tasks.NotFound"},
{"ErrTasksInvalidData", ErrTasksInvalidData, 500, "Tasks.InvalidData"},
// Subtasks errors
{"ErrSubtasksInvalidRequest", ErrSubtasksInvalidRequest, 400, "Subtasks.InvalidRequest"},
{"ErrSubtasksNotFound", ErrSubtasksNotFound, 404, "Subtasks.NotFound"},
{"ErrSubtasksInvalidData", ErrSubtasksInvalidData, 500, "Subtasks.InvalidData"},
// Assistants errors
{"ErrAssistantsInvalidRequest", ErrAssistantsInvalidRequest, 400, "Assistants.InvalidRequest"},
{"ErrAssistantsNotFound", ErrAssistantsNotFound, 404, "Assistants.NotFound"},
{"ErrAssistantsInvalidData", ErrAssistantsInvalidData, 500, "Assistants.InvalidData"},
// Tokens errors
{"ErrTokenCreationDisabled", ErrTokenCreationDisabled, 400, "Token.CreationDisabled"},
{"ErrTokenNotFound", ErrTokenNotFound, 404, "Token.NotFound"},
{"ErrTokenUnauthorized", ErrTokenUnauthorized, 403, "Token.Unauthorized"},
{"ErrTokenInvalidRequest", ErrTokenInvalidRequest, 400, "Token.InvalidRequest"},
{"ErrTokenInvalidData", ErrTokenInvalidData, 500, "Token.InvalidData"},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
t.Parallel()
assert.Equal(t, tt.httpCode, tt.err.HttpCode())
assert.Equal(t, tt.code, tt.err.Code())
assert.NotEmpty(t, tt.err.Msg())
assert.NotEmpty(t, tt.err.Error())
})
}
}
func TestSuccessResponse(t *testing.T) {
t.Parallel()
w := httptest.NewRecorder()
c, _ := gin.CreateTestContext(w)
data := map[string]string{"id": "123"}
Success(c, http.StatusOK, data)
assert.Equal(t, http.StatusOK, w.Code)
var body map[string]any
err := json.Unmarshal(w.Body.Bytes(), &body)
require.NoError(t, err)
assert.Equal(t, "success", body["status"])
assert.NotNil(t, body["data"])
}
func TestSuccessResponse_Created(t *testing.T) {
t.Parallel()
w := httptest.NewRecorder()
c, _ := gin.CreateTestContext(w)
Success(c, http.StatusCreated, gin.H{"name": "test"})
assert.Equal(t, http.StatusCreated, w.Code)
}
func TestErrorResponse(t *testing.T) {
t.Parallel()
w := httptest.NewRecorder()
c, _ := gin.CreateTestContext(w)
c.Request = httptest.NewRequest(http.MethodGet, "/test", nil)
Error(c, ErrInternal, errors.New("db connection failed"))
assert.Equal(t, http.StatusInternalServerError, w.Code)
var body map[string]any
err := json.Unmarshal(w.Body.Bytes(), &body)
require.NoError(t, err)
assert.Equal(t, "error", body["status"])
assert.Equal(t, "Internal", body["code"])
assert.Equal(t, "internal server error", body["msg"])
}
func TestErrorResponse_DevMode(t *testing.T) {
// Save original version and restore after test
oldVer := version.PackageVer
defer func() { version.PackageVer = oldVer }()
// Enable dev mode
version.PackageVer = ""
w := httptest.NewRecorder()
c, _ := gin.CreateTestContext(w)
c.Request = httptest.NewRequest(http.MethodGet, "/test", nil)
originalErr := errors.New("detailed error info")
Error(c, ErrInternal, originalErr)
var body map[string]any
err := json.Unmarshal(w.Body.Bytes(), &body)
require.NoError(t, err)
// In dev mode, original error should be included
assert.Equal(t, "detailed error info", body["error"])
}
func TestErrorResponse_ProductionMode(t *testing.T) {
// Save original version and restore after test
oldVer := version.PackageVer
defer func() { version.PackageVer = oldVer }()
// Set production mode
version.PackageVer = "1.0.0"
w := httptest.NewRecorder()
c, _ := gin.CreateTestContext(w)
c.Request = httptest.NewRequest(http.MethodGet, "/test", nil)
Error(c, ErrInternal, errors.New("should not appear"))
var body map[string]any
err := json.Unmarshal(w.Body.Bytes(), &body)
require.NoError(t, err)
// In production mode, original error should NOT be included
_, hasError := body["error"]
assert.False(t, hasError)
}
func TestErrorResponse_NilOriginalError(t *testing.T) {
t.Parallel()
w := httptest.NewRecorder()
c, _ := gin.CreateTestContext(w)
c.Request = httptest.NewRequest(http.MethodGet, "/test", nil)
Error(c, ErrNotPermitted, nil)
assert.Equal(t, http.StatusForbidden, w.Code)
var body map[string]any
err := json.Unmarshal(w.Body.Bytes(), &body)
require.NoError(t, err)
assert.Equal(t, "NotPermitted", body["code"])
}
func TestHttpError_MultipleInstancesIndependent(t *testing.T) {
t.Parallel()
err1 := NewHttpError(404, "NotFound", "resource 1 not found")
err2 := NewHttpError(404, "NotFound", "resource 2 not found")
// Verify they are independent instances
assert.NotEqual(t, err1.Msg(), err2.Msg())
assert.Equal(t, err1.Code(), err2.Code())
assert.Equal(t, err1.HttpCode(), err2.HttpCode())
}
func TestSuccessResponse_EmptyData(t *testing.T) {
t.Parallel()
w := httptest.NewRecorder()
c, _ := gin.CreateTestContext(w)
Success(c, http.StatusOK, nil)
assert.Equal(t, http.StatusOK, w.Code)
var body map[string]any
err := json.Unmarshal(w.Body.Bytes(), &body)
require.NoError(t, err)
assert.Equal(t, "success", body["status"])
assert.Nil(t, body["data"])
}
func TestSuccessResponse_ComplexData(t *testing.T) {
t.Parallel()
w := httptest.NewRecorder()
c, _ := gin.CreateTestContext(w)
data := gin.H{
"users": []gin.H{
{"id": 1, "name": "Alice"},
{"id": 2, "name": "Bob"},
},
"count": 2,
"meta": gin.H{
"page": 1,
"total": 100,
},
}
Success(c, http.StatusOK, data)
assert.Equal(t, http.StatusOK, w.Code)
var body map[string]any
err := json.Unmarshal(w.Body.Bytes(), &body)
require.NoError(t, err)
assert.Equal(t, "success", body["status"])
responseData, ok := body["data"].(map[string]any)
require.True(t, ok)
assert.Equal(t, float64(2), responseData["count"])
}
func TestErrorResponse_DifferentHttpCodes(t *testing.T) {
t.Parallel()
tests := []struct {
name string
err *HttpError
expected int
}{
{"400 Bad Request", ErrPromptsInvalidRequest, http.StatusBadRequest},
{"401 Unauthorized", ErrAuthInvalidCredentials, http.StatusUnauthorized},
{"403 Forbidden", ErrNotPermitted, http.StatusForbidden},
{"404 Not Found", ErrUsersNotFound, http.StatusNotFound},
{"500 Internal", ErrInternal, http.StatusInternalServerError},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
t.Parallel()
w := httptest.NewRecorder()
c, _ := gin.CreateTestContext(w)
c.Request = httptest.NewRequest(http.MethodGet, "/test", nil)
Error(c, tt.err, nil)
assert.Equal(t, tt.expected, w.Code)
})
}
}
func TestErrorResponse_ResponseStructure(t *testing.T) {
t.Parallel()
w := httptest.NewRecorder()
c, _ := gin.CreateTestContext(w)
c.Request = httptest.NewRequest(http.MethodGet, "/test", nil)
Error(c, ErrUsersNotFound, nil)
var body map[string]any
err := json.Unmarshal(w.Body.Bytes(), &body)
require.NoError(t, err)
// Verify required fields
assert.Equal(t, "error", body["status"])
assert.Equal(t, "Users.NotFound", body["code"])
assert.Equal(t, "user not found", body["msg"])
// Verify error field is not present in non-dev mode
_, hasError := body["error"]
assert.False(t, hasError)
}
func TestSuccessResponse_StatusCodes(t *testing.T) {
t.Parallel()
tests := []struct {
name string
statusCode int
expectedCode int
}{
{"200 OK", http.StatusOK, 200},
{"201 Created", http.StatusCreated, 201},
{"202 Accepted", http.StatusAccepted, 202},
{"204 No Content", http.StatusNoContent, 204},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
t.Parallel()
w := httptest.NewRecorder()
c, _ := gin.CreateTestContext(w)
Success(c, tt.statusCode, gin.H{"test": "data"})
assert.Equal(t, tt.expectedCode, w.Code)
})
}
}
+695
View File
@@ -0,0 +1,695 @@
package router
import (
"context"
"encoding/gob"
"net"
"net/http"
"net/http/httputil"
"net/url"
"os"
"path"
"path/filepath"
"slices"
"strings"
"time"
"pentagi/pkg/config"
"pentagi/pkg/controller"
"pentagi/pkg/database"
"pentagi/pkg/database/knowledge"
"pentagi/pkg/docker"
"pentagi/pkg/graph/subscriptions"
"pentagi/pkg/providers"
"pentagi/pkg/server/auth"
"pentagi/pkg/server/logger"
"pentagi/pkg/server/oauth"
"pentagi/pkg/server/services"
_ "pentagi/pkg/server/docs" // swagger docs
"github.com/gin-contrib/cors"
"github.com/gin-contrib/sessions"
"github.com/gin-contrib/sessions/cookie"
"github.com/gin-contrib/static"
"github.com/gin-gonic/gin"
"github.com/jinzhu/gorm"
"github.com/sirupsen/logrus"
ginSwagger "github.com/swaggo/gin-swagger"
"github.com/swaggo/gin-swagger/swaggerFiles"
"github.com/vxcontrol/cloud/anonymizer"
"github.com/vxcontrol/cloud/anonymizer/patterns"
"github.com/vxcontrol/langchaingo/vectorstores/pgvector"
)
const baseURL = "/api/v1"
const corsAllowGoogleOAuth = "https://accounts.google.com"
// frontendRoutes defines the list of URI prefixes that should be handled by the frontend SPA.
// Add new frontend base routes here if they are added in the frontend router (e.g., in App.tsx).
var frontendRoutes = []string{
"/chat",
"/oauth",
"/login",
"/flows",
"/settings",
"/templates",
"/resources",
"/knowledges",
"/dashboard",
}
// @title PentAGI Swagger API
// @version 1.0
// @description Swagger API for Penetration Testing Advanced General Intelligence PentAGI.
// @termsOfService http://swagger.io/terms/
// @contact.url https://pentagi.com
// @contact.name PentAGI Development Team
// @contact.email team@pentagi.com
// @license.name MIT
// @license.url https://opensource.org/license/mit
// @query.collection.format multi
// @securityDefinitions.apikey BearerAuth
// @in header
// @name Authorization
// @description Type "Bearer" followed by a space and JWT token.
// @BasePath /api/v1
func NewRouter(
db *database.Queries,
orm *gorm.DB,
cfg *config.Config,
providers providers.ProviderController,
controller controller.FlowController,
subscriptions subscriptions.SubscriptionsController,
dockerClient docker.DockerClient,
) *gin.Engine {
gin.SetMode(gin.ReleaseMode)
if cfg.Debug {
gin.SetMode(gin.DebugMode)
}
gob.Register([]string{})
tokenCache := auth.NewTokenCache(orm)
userCache := auth.NewUserCache(orm)
authMiddleware := auth.NewAuthMiddleware(baseURL, cfg.CookieSigningSalt, tokenCache, userCache)
oauthClients := make(map[string]oauth.OAuthClient)
oauthLoginCallbackURL := "/auth/login-callback"
publicURL, err := url.Parse(cfg.PublicURL)
if err == nil {
publicURL.Path = path.Join(baseURL, oauthLoginCallbackURL)
}
if publicURL != nil && cfg.OAuthGoogleClientID != "" && cfg.OAuthGoogleClientSecret != "" {
googleClient := oauth.NewGoogleOAuthClient(
cfg.OAuthGoogleClientID,
cfg.OAuthGoogleClientSecret,
publicURL.String(),
)
oauthClients[googleClient.ProviderName()] = googleClient
}
if publicURL != nil && cfg.OAuthGithubClientID != "" && cfg.OAuthGithubClientSecret != "" {
githubClient := oauth.NewGithubOAuthClient(
cfg.OAuthGithubClientID,
cfg.OAuthGithubClientSecret,
publicURL.String(),
)
oauthClients[githubClient.ProviderName()] = githubClient
}
// ---- Knowledge (pgvector) store -----------------------------------------
// Shared by both the GraphQL and REST layers.
// Store and embedder are nil when no embedding provider is configured;
// the knowledge store handles that gracefully (embedding-dependent ops error).
embedder := providers.Embedder()
var pgStore *pgvector.Store
if embedder.IsAvailable() {
opts := []pgvector.Option{
pgvector.WithEmbedder(embedder),
pgvector.WithCollectionName("langchain"),
}
if cfg.PgxPool != nil {
opts = append(opts, pgvector.WithConn(cfg.PgxPool))
} else {
opts = append(opts, pgvector.WithConnectionURL(cfg.DatabaseURL))
}
if s, err := pgvector.New(context.Background(), opts...); err == nil {
pgStore = &s
} else {
logrus.WithError(err).Warn("failed to initialise pgvector store for knowledge API; embedding operations will be unavailable")
}
}
var knowledgeStore knowledge.KnowledgeStore
knowledgeStore = knowledge.NewKnowledgeStore(db, pgStore, embedder, subscriptions.NewKnowledgePublisher, cfg.EmbeddingMaxTextBytes)
// ---- Anonymizer replacer ------------------------------------------------
// Shared singleton used by the GraphQL anonymizeText mutation.
// Falls back to a no-op nil replacer on failure so the rest of the server still starts correctly.
var textReplacer anonymizer.Replacer
if allPatterns, err := patterns.LoadPatterns(patterns.PatternListTypeAll); err != nil {
logrus.WithError(err).Warn("failed to load anonymizer patterns; anonymizeText mutation will be unavailable")
} else {
allPatterns.Patterns = append(allPatterns.Patterns, cfg.GetSecretPatterns()...)
if r, err := anonymizer.NewReplacer(allPatterns.Regexes(), allPatterns.Names()); err != nil {
logrus.WithError(err).Warn("failed to create anonymizer replacer; anonymizeText mutation will be unavailable")
} else {
textReplacer = r
}
}
// services
authService := services.NewAuthService(
services.AuthServiceConfig{
BaseURL: baseURL,
LoginCallbackURL: oauthLoginCallbackURL,
SessionTimeout: 4 * 60 * 60, // 4 hours
},
orm,
oauthClients,
)
userService := services.NewUserService(orm, userCache)
roleService := services.NewRoleService(orm)
providerService := services.NewProviderService(providers)
settingsService := services.NewSettingsService(cfg)
flowService := services.NewFlowService(orm, providers, controller, subscriptions)
flowFileService := services.NewFlowFileService(orm, cfg.DataDir, dockerClient, subscriptions)
resourceService := services.NewResourceService(orm, cfg.DataDir, subscriptions)
taskService := services.NewTaskService(orm)
subtaskService := services.NewSubtaskService(orm)
containerService := services.NewContainerService(orm)
toolcallService := services.NewToolcallService(orm)
assistantService := services.NewAssistantService(orm, providers, controller, subscriptions)
agentlogService := services.NewAgentlogService(orm)
assistantlogService := services.NewAssistantlogService(orm)
msglogService := services.NewMsglogService(orm)
searchlogService := services.NewSearchlogService(orm)
vecstorelogService := services.NewVecstorelogService(orm)
termlogService := services.NewTermlogService(orm)
screenshotService := services.NewScreenshotService(orm, cfg.DataDir)
promptService := services.NewPromptService(orm)
analyticsService := services.NewAnalyticsService(orm)
tokenService := services.NewTokenService(orm, cfg.CookieSigningSalt, tokenCache, subscriptions)
knowledgeService := services.NewKnowledgeService(orm, knowledgeStore)
anonymizerService := services.NewAnonymizerService(textReplacer)
graphqlService := services.NewGraphqlService(
db, cfg, baseURL, cfg.CorsOrigins, tokenCache, providers, controller, subscriptions, knowledgeStore, textReplacer,
)
router := gin.Default()
// Setup Cross-Origin Resource Sharing policy
config := cors.DefaultConfig()
if !slices.Contains(cfg.CorsOrigins, "*") {
config.AllowCredentials = true
}
config.AllowWildcard = true
config.AllowWebSockets = true
config.AllowPrivateNetwork = true
// Add OAuth provider origins to CORS allowed origins
allowedOrigins := make([]string, len(cfg.CorsOrigins))
copy(allowedOrigins, cfg.CorsOrigins)
// Google OAuth uses POST callback from accounts.google.com
if cfg.OAuthGoogleClientID != "" && cfg.OAuthGoogleClientSecret != "" {
if !slices.Contains(allowedOrigins, corsAllowGoogleOAuth) && !slices.Contains(cfg.CorsOrigins, "*") {
allowedOrigins = append(allowedOrigins, corsAllowGoogleOAuth)
logrus.Infof("Added %s to CORS allowed origins for Google OAuth", corsAllowGoogleOAuth)
}
}
config.AllowOrigins = allowedOrigins
config.AllowMethods = []string{"GET", "POST", "PUT", "DELETE", "OPTIONS"}
if err := config.Validate(); err != nil {
logrus.WithError(err).Error("failed to validate cors config")
} else {
router.Use(cors.New(config))
}
router.Use(gin.Recovery())
router.Use(logger.WithGinLogger("pentagi-api"))
cookieStore := cookie.NewStore(auth.MakeCookieStoreKey(cfg.CookieSigningSalt)...)
router.Use(sessions.Sessions("auth", cookieStore))
api := router.Group(baseURL)
api.Use(noCacheMiddleware())
// Special case for local user own password change
changePasswordGroup := api.Group("/user")
changePasswordGroup.Use(authMiddleware.AuthUserRequired)
changePasswordGroup.Use(localUserRequired())
changePasswordGroup.PUT("/password", userService.ChangePasswordCurrentUser)
publicGroup := api.Group("/")
publicGroup.Use(authMiddleware.TryAuth)
{
publicGroup.GET("/info", authService.Info)
developerGroup := publicGroup.Group("/")
{
developerGroup.GET("/graphql/playground", graphqlService.ServeGraphqlPlayground)
developerGroup.GET("/swagger/*any", ginSwagger.WrapHandler(swaggerFiles.Handler))
}
authGroup := publicGroup.Group("/auth")
{
authGroup.POST("/login", authService.AuthLogin)
authGroup.GET("/logout", authService.AuthLogout)
authGroup.GET("/authorize", authService.AuthAuthorize)
authGroup.GET("/login-callback", authService.AuthLoginGetCallback)
authGroup.POST("/login-callback", authService.AuthLoginPostCallback)
authGroup.POST("/logout-callback", authService.AuthLogoutCallback)
}
}
privateGroup := api.Group("/")
privateGroup.Use(authMiddleware.AuthTokenRequired)
{
setGraphqlGroup(privateGroup, graphqlService)
setKnowledgeGroup(privateGroup, knowledgeService)
setProvidersGroup(privateGroup, providerService)
setSettingsGroup(privateGroup, settingsService)
setFlowsGroup(privateGroup, flowService)
setFlowFilesGroup(privateGroup, flowFileService)
setResourcesGroup(privateGroup, resourceService)
setTasksGroup(privateGroup, taskService)
setSubtasksGroup(privateGroup, subtaskService)
setContainersGroup(privateGroup, containerService)
setToolcallsGroup(privateGroup, toolcallService)
setAssistantsGroup(privateGroup, assistantService)
setAgentlogsGroup(privateGroup, agentlogService)
setAssistantlogsGroup(privateGroup, assistantlogService)
setMsglogsGroup(privateGroup, msglogService)
setTermlogsGroup(privateGroup, termlogService)
setSearchlogsGroup(privateGroup, searchlogService)
setVecstorelogsGroup(privateGroup, vecstorelogService)
setScreenshotsGroup(privateGroup, screenshotService)
setPromptsGroup(privateGroup, promptService)
setAnonymizeGroup(privateGroup, anonymizerService)
setAnalyticsGroup(privateGroup, analyticsService)
}
privateUserGroup := api.Group("/")
privateUserGroup.Use(authMiddleware.AuthUserRequired)
{
setRolesGroup(privateGroup, roleService)
setUsersGroup(privateGroup, userService)
setTokensGroup(privateGroup, tokenService)
}
if cfg.StaticURL != nil && cfg.StaticURL.Scheme != "" && cfg.StaticURL.Host != "" {
router.NoRoute(func() gin.HandlerFunc {
return func(c *gin.Context) {
director := func(req *http.Request) {
*req = *c.Request
req.URL.Scheme = cfg.StaticURL.Scheme
req.URL.Host = cfg.StaticURL.Host
}
dialer := &net.Dialer{
Timeout: 30 * time.Second,
KeepAlive: 30 * time.Second,
}
httpTransport := &http.Transport{
DialContext: dialer.DialContext,
ForceAttemptHTTP2: true,
MaxIdleConns: 20,
IdleConnTimeout: 60 * time.Second,
TLSHandshakeTimeout: 10 * time.Second,
ExpectContinueTimeout: 1 * time.Second,
}
proxy := &httputil.ReverseProxy{
Director: director,
Transport: httpTransport,
}
proxy.ServeHTTP(c.Writer, c.Request)
}
}())
} else {
router.Use(static.Serve("/", static.LocalFile(cfg.StaticDir, true)))
indexExists := true
indexPath := filepath.Join(cfg.StaticDir, "index.html")
if _, err := os.Stat(indexPath); err != nil {
indexExists = false
}
router.NoRoute(func(c *gin.Context) {
if c.Request.Method == "GET" && !strings.HasPrefix(c.Request.URL.Path, baseURL) {
isFrontendRoute := false
path := c.Request.URL.Path
for _, prefix := range frontendRoutes {
if path == prefix || strings.HasPrefix(path, prefix+"/") {
isFrontendRoute = true
break
}
}
if isFrontendRoute && indexExists {
c.File(indexPath)
return
}
}
c.Redirect(http.StatusMovedPermanently, "/")
})
}
return router
}
func setKnowledgeGroup(parent *gin.RouterGroup, svc *services.KnowledgeService) {
kg := parent.Group("/knowledge")
{
kg.GET("/", svc.ListDocuments)
kg.GET("/:id", svc.GetDocument)
kg.POST("/", svc.CreateDocument)
kg.POST("/search", svc.SearchDocuments)
kg.PUT("/:id", svc.UpdateDocument)
kg.DELETE("/:id", svc.DeleteDocument)
}
}
func setProvidersGroup(parent *gin.RouterGroup, svc *services.ProviderService) {
providersGroup := parent.Group("/providers")
{
providersGroup.GET("/", svc.GetProviders)
}
}
func setSettingsGroup(parent *gin.RouterGroup, svc *services.SettingsService) {
settingsGroup := parent.Group("/settings")
{
settingsGroup.GET("/", svc.GetSettings)
}
}
func setGraphqlGroup(parent *gin.RouterGroup, svc *services.GraphqlService) {
graphqlGroup := parent.Group("/")
{
graphqlGroup.Any("/graphql", svc.ServeGraphql)
}
}
func setSubtasksGroup(parent *gin.RouterGroup, svc *services.SubtaskService) {
flowSubtasksViewGroup := parent.Group("/flows/:flowID/subtasks")
{
flowSubtasksViewGroup.GET("/", svc.GetFlowSubtasks)
}
flowTaskSubtasksViewGroup := parent.Group("/flows/:flowID/tasks/:taskID/subtasks")
{
flowTaskSubtasksViewGroup.GET("/", svc.GetFlowTaskSubtasks)
flowTaskSubtasksViewGroup.GET("/:subtaskID", svc.GetFlowTaskSubtask)
}
}
func setTasksGroup(parent *gin.RouterGroup, svc *services.TaskService) {
flowTaskViewGroup := parent.Group("/flows/:flowID/tasks")
{
flowTaskViewGroup.GET("/", svc.GetFlowTasks)
flowTaskViewGroup.GET("/:taskID", svc.GetFlowTask)
flowTaskViewGroup.GET("/:taskID/graph", svc.GetFlowTaskGraph)
}
}
func setFlowsGroup(parent *gin.RouterGroup, svc *services.FlowService) {
flowCreateGroup := parent.Group("/flows")
{
flowCreateGroup.POST("/", svc.CreateFlow)
}
flowDeleteGroup := parent.Group("/flows")
{
flowDeleteGroup.DELETE("/:flowID", svc.DeleteFlow)
}
flowEditGroup := parent.Group("/flows")
{
flowEditGroup.PUT("/:flowID", svc.PatchFlow)
}
flowsViewGroup := parent.Group("/flows")
{
flowsViewGroup.GET("/", svc.GetFlows)
flowsViewGroup.GET("/:flowID", svc.GetFlow)
flowsViewGroup.GET("/:flowID/graph", svc.GetFlowGraph)
}
}
func setFlowFilesGroup(parent *gin.RouterGroup, svc *services.FlowFileService) {
flowFilesGroup := parent.Group("/flows/:flowID/files")
{
flowFilesGroup.GET("/", svc.GetFlowFiles)
flowFilesGroup.GET("/container", svc.GetFlowContainerFiles)
flowFilesGroup.POST("/", svc.UploadFlowFiles)
flowFilesGroup.DELETE("/", svc.DeleteFlowFile)
flowFilesGroup.GET("/download", svc.DownloadFlowFile)
flowFilesGroup.POST("/pull", svc.PullFlowFiles)
flowFilesGroup.POST("/resources", svc.AddResourcesToFlow)
flowFilesGroup.POST("/to-resources", svc.AddResourceFromFlow)
}
}
func setResourcesGroup(parent *gin.RouterGroup, svc *services.ResourceService) {
rg := parent.Group("/resources")
{
rg.GET("/", svc.ListResources)
rg.POST("/", svc.UploadResources)
rg.POST("/mkdir", svc.MkdirResource)
rg.PUT("/move", svc.MoveResource)
rg.POST("/copy", svc.CopyResource)
rg.DELETE("/", svc.DeleteResource)
rg.GET("/download", svc.DownloadResource)
}
}
func setContainersGroup(parent *gin.RouterGroup, svc *services.ContainerService) {
containersViewGroup := parent.Group("/containers")
{
containersViewGroup.GET("/", svc.GetContainers)
}
flowContainersViewGroup := parent.Group("/flows/:flowID/containers")
{
flowContainersViewGroup.GET("/", svc.GetFlowContainers)
flowContainersViewGroup.GET("/:containerID", svc.GetFlowContainer)
}
}
func setToolcallsGroup(parent *gin.RouterGroup, svc *services.ToolcallService) {
toolcallsViewGroup := parent.Group("/toolcalls")
{
toolcallsViewGroup.GET("/", svc.GetToolcalls)
}
flowToolcallsViewGroup := parent.Group("/flows/:flowID/toolcalls")
{
flowToolcallsViewGroup.GET("/", svc.GetFlowToolcalls)
flowToolcallsViewGroup.GET("/:toolcallID", svc.GetFlowToolcall)
}
}
func setAssistantsGroup(parent *gin.RouterGroup, svc *services.AssistantService) {
flowCreateGroup := parent.Group("/flows/:flowID/assistants")
{
flowCreateGroup.POST("/", svc.CreateFlowAssistant)
}
flowDeleteGroup := parent.Group("/flows/:flowID/assistants")
{
flowDeleteGroup.DELETE("/:assistantID", svc.DeleteAssistant)
}
flowEditGroup := parent.Group("/flows/:flowID/assistants")
{
flowEditGroup.PUT("/:assistantID", svc.PatchAssistant)
}
flowsViewGroup := parent.Group("/flows/:flowID/assistants")
{
flowsViewGroup.GET("/", svc.GetFlowAssistants)
flowsViewGroup.GET("/:assistantID", svc.GetFlowAssistant)
}
}
func setAgentlogsGroup(parent *gin.RouterGroup, svc *services.AgentlogService) {
agentlogsViewGroup := parent.Group("/agentlogs")
{
agentlogsViewGroup.GET("/", svc.GetAgentlogs)
}
flowAgentlogsViewGroup := parent.Group("/flows/:flowID/agentlogs")
{
flowAgentlogsViewGroup.GET("/", svc.GetFlowAgentlogs)
}
}
func setAssistantlogsGroup(parent *gin.RouterGroup, svc *services.AssistantlogService) {
assistantlogsViewGroup := parent.Group("/assistantlogs")
{
assistantlogsViewGroup.GET("/", svc.GetAssistantlogs)
}
flowAssistantlogsViewGroup := parent.Group("/flows/:flowID/assistantlogs")
{
flowAssistantlogsViewGroup.GET("/", svc.GetFlowAssistantlogs)
}
}
func setMsglogsGroup(parent *gin.RouterGroup, svc *services.MsglogService) {
msglogsViewGroup := parent.Group("/msglogs")
{
msglogsViewGroup.GET("/", svc.GetMsglogs)
}
flowMsglogsViewGroup := parent.Group("/flows/:flowID/msglogs")
{
flowMsglogsViewGroup.GET("/", svc.GetFlowMsglogs)
}
}
func setSearchlogsGroup(parent *gin.RouterGroup, svc *services.SearchlogService) {
searchlogsViewGroup := parent.Group("/searchlogs")
{
searchlogsViewGroup.GET("/", svc.GetSearchlogs)
}
flowSearchlogsViewGroup := parent.Group("/flows/:flowID/searchlogs")
{
flowSearchlogsViewGroup.GET("/", svc.GetFlowSearchlogs)
}
}
func setTermlogsGroup(parent *gin.RouterGroup, svc *services.TermlogService) {
termlogsViewGroup := parent.Group("/termlogs")
{
termlogsViewGroup.GET("/", svc.GetTermlogs)
}
flowTermlogsViewGroup := parent.Group("/flows/:flowID/termlogs")
{
flowTermlogsViewGroup.GET("/", svc.GetFlowTermlogs)
}
}
func setVecstorelogsGroup(parent *gin.RouterGroup, svc *services.VecstorelogService) {
vecstorelogsViewGroup := parent.Group("/vecstorelogs")
{
vecstorelogsViewGroup.GET("/", svc.GetVecstorelogs)
}
flowVecstorelogsViewGroup := parent.Group("/flows/:flowID/vecstorelogs")
{
flowVecstorelogsViewGroup.GET("/", svc.GetFlowVecstorelogs)
}
}
func setScreenshotsGroup(parent *gin.RouterGroup, svc *services.ScreenshotService) {
screenshotsViewGroup := parent.Group("/screenshots")
{
screenshotsViewGroup.GET("/", svc.GetScreenshots)
}
flowScreenshotsViewGroup := parent.Group("/flows/:flowID/screenshots")
{
flowScreenshotsViewGroup.GET("/", svc.GetFlowScreenshots)
flowScreenshotsViewGroup.GET("/:screenshotID", svc.GetFlowScreenshot)
flowScreenshotsViewGroup.GET("/:screenshotID/file", svc.GetFlowScreenshotFile)
}
}
func setAnonymizeGroup(parent *gin.RouterGroup, svc *services.AnonymizerService) {
group := parent.Group("/anonymize")
{
group.POST("/text", svc.AnonymizeText)
}
}
func setPromptsGroup(parent *gin.RouterGroup, svc *services.PromptService) {
promptsViewGroup := parent.Group("/prompts")
{
promptsViewGroup.GET("/", svc.GetPrompts)
promptsViewGroup.GET("/:promptType", svc.GetPrompt)
}
promptsEditGroup := parent.Group("/prompts")
{
promptsEditGroup.PUT("/:promptType", svc.PatchPrompt)
promptsEditGroup.POST("/:promptType/default", svc.ResetPrompt)
promptsEditGroup.DELETE("/:promptType", svc.DeletePrompt)
}
}
func setRolesGroup(parent *gin.RouterGroup, svc *services.RoleService) {
rolesViewGroup := parent.Group("/roles")
{
rolesViewGroup.GET("/", svc.GetRoles)
rolesViewGroup.GET("/:roleID", svc.GetRole)
}
}
func setUsersGroup(parent *gin.RouterGroup, svc *services.UserService) {
usersCreateGroup := parent.Group("/users")
{
usersCreateGroup.POST("/", svc.CreateUser)
}
usersDeleteGroup := parent.Group("/users")
{
usersDeleteGroup.DELETE("/:hash", svc.DeleteUser)
}
usersEditGroup := parent.Group("/users")
{
usersEditGroup.PUT("/:hash", svc.PatchUser)
}
usersViewGroup := parent.Group("/users")
{
usersViewGroup.GET("/", svc.GetUsers)
usersViewGroup.GET("/:hash", svc.GetUser)
}
userViewGroup := parent.Group("/user")
{
userViewGroup.GET("/", svc.GetCurrentUser)
}
}
func setAnalyticsGroup(parent *gin.RouterGroup, svc *services.AnalyticsService) {
// System-wide analytics
usageViewGroup := parent.Group("/usage")
{
usageViewGroup.GET("/", svc.GetSystemUsage)
usageViewGroup.GET("/:period", svc.GetPeriodUsage)
}
// Flow-specific analytics
flowUsageViewGroup := parent.Group("/flows/:flowID/usage")
{
flowUsageViewGroup.GET("/", svc.GetFlowUsage)
}
}
func setTokensGroup(parent *gin.RouterGroup, svc *services.TokenService) {
tokensGroup := parent.Group("/tokens")
{
tokensGroup.POST("/", svc.CreateToken)
tokensGroup.GET("/", svc.ListTokens)
tokensGroup.GET("/:tokenID", svc.GetToken)
tokensGroup.PUT("/:tokenID", svc.UpdateToken)
tokensGroup.DELETE("/:tokenID", svc.DeleteToken)
}
}
+220
View File
@@ -0,0 +1,220 @@
package services
import (
"errors"
"net/http"
"slices"
"strconv"
"pentagi/pkg/server/logger"
"pentagi/pkg/server/models"
"pentagi/pkg/server/rdb"
"pentagi/pkg/server/response"
"github.com/gin-gonic/gin"
"github.com/jinzhu/gorm"
)
type agentlogs struct {
AgentLogs []models.Agentlog `json:"agentlogs"`
Total uint64 `json:"total"`
}
type agentlogsGrouped struct {
Grouped []string `json:"grouped"`
Total uint64 `json:"total"`
}
var agentlogsSQLMappers = map[string]any{
"id": "{{table}}.id",
"initiator": "{{table}}.initiator",
"executor": "{{table}}.executor",
"task": "{{table}}.task",
"result": "{{table}}.result",
"flow_id": "{{table}}.flow_id",
"task_id": "{{table}}.task_id",
"subtask_id": "{{table}}.subtask_id",
"created_at": "{{table}}.created_at",
"data": "({{table}}.task || ' ' || {{table}}.result)",
}
type AgentlogService struct {
db *gorm.DB
}
func NewAgentlogService(db *gorm.DB) *AgentlogService {
return &AgentlogService{
db: db,
}
}
// GetAgentlogs is a function to return agentlogs list
// @Summary Retrieve agentlogs list
// @Tags Agentlogs
// @Produce json
// @Security BearerAuth
// @Param request query rdb.TableQuery true "query table params"
// @Success 200 {object} response.successResp{data=agentlogs} "agentlogs list received successful"
// @Failure 400 {object} response.errorResp "invalid query request data"
// @Failure 403 {object} response.errorResp "getting agentlogs not permitted"
// @Failure 500 {object} response.errorResp "internal error on getting agentlogs"
// @Router /agentlogs/ [get]
func (s *AgentlogService) GetAgentlogs(c *gin.Context) {
var (
err error
query rdb.TableQuery
resp agentlogs
)
if err = c.ShouldBindQuery(&query); err != nil {
logger.FromContext(c).WithError(err).Errorf("error binding query")
response.Error(c, response.ErrAgentlogsInvalidRequest, err)
return
}
uid := c.GetUint64("uid")
privs := c.GetStringSlice("prm")
var scope func(db *gorm.DB) *gorm.DB
if slices.Contains(privs, "agentlogs.admin") {
scope = func(db *gorm.DB) *gorm.DB {
return db.
Joins("INNER JOIN flows f ON f.id = flow_id")
}
} else if slices.Contains(privs, "agentlogs.view") {
scope = func(db *gorm.DB) *gorm.DB {
return db.
Joins("INNER JOIN flows f ON f.id = flow_id").
Where("f.user_id = ?", uid)
}
} else {
logger.FromContext(c).Errorf("error filtering user role permissions: permission not found")
response.Error(c, response.ErrNotPermitted, nil)
return
}
query.Init("agentlogs", agentlogsSQLMappers)
if query.Group != "" {
if _, ok := agentlogsSQLMappers[query.Group]; !ok {
logger.FromContext(c).Errorf("error finding agentlogs grouped: group field not found")
response.Error(c, response.ErrAgentlogsInvalidRequest, errors.New("group field not found"))
return
}
var respGrouped agentlogsGrouped
if respGrouped.Total, err = query.QueryGrouped(s.db, &respGrouped.Grouped, scope); err != nil {
logger.FromContext(c).WithError(err).Errorf("error finding agentlogs grouped")
response.Error(c, response.ErrInternal, err)
return
}
response.Success(c, http.StatusOK, respGrouped)
return
}
if resp.Total, err = query.Query(s.db, &resp.AgentLogs, scope); err != nil {
logger.FromContext(c).WithError(err).Errorf("error finding agentlogs")
response.Error(c, response.ErrInternal, err)
return
}
for i := 0; i < len(resp.AgentLogs); i++ {
if err = resp.AgentLogs[i].Valid(); err != nil {
logger.FromContext(c).WithError(err).Errorf("error validating agentlog data '%d'", resp.AgentLogs[i].ID)
response.Error(c, response.ErrAgentlogsInvalidData, err)
return
}
}
response.Success(c, http.StatusOK, resp)
}
// GetFlowAgentlogs is a function to return agentlogs list by flow id
// @Summary Retrieve agentlogs list by flow id
// @Tags Agentlogs
// @Produce json
// @Security BearerAuth
// @Param flowID path int true "flow id" minimum(0)
// @Param request query rdb.TableQuery true "query table params"
// @Success 200 {object} response.successResp{data=agentlogs} "agentlogs list received successful"
// @Failure 400 {object} response.errorResp "invalid query request data"
// @Failure 403 {object} response.errorResp "getting agentlogs not permitted"
// @Failure 500 {object} response.errorResp "internal error on getting agentlogs"
// @Router /flows/{flowID}/agentlogs/ [get]
func (s *AgentlogService) GetFlowAgentlogs(c *gin.Context) {
var (
err error
flowID uint64
query rdb.TableQuery
resp agentlogs
)
if flowID, err = strconv.ParseUint(c.Param("flowID"), 10, 64); err != nil {
logger.FromContext(c).WithError(err).Errorf("error parsing flow id")
response.Error(c, response.ErrAgentlogsInvalidRequest, err)
return
}
if err = c.ShouldBindQuery(&query); err != nil {
logger.FromContext(c).WithError(err).Errorf("error binding query")
response.Error(c, response.ErrAgentlogsInvalidRequest, err)
return
}
uid := c.GetUint64("uid")
privs := c.GetStringSlice("prm")
var scope func(db *gorm.DB) *gorm.DB
if slices.Contains(privs, "agentlogs.admin") {
scope = func(db *gorm.DB) *gorm.DB {
return db.
Joins("INNER JOIN flows f ON f.id = flow_id").
Where("f.id = ?", flowID)
}
} else if slices.Contains(privs, "agentlogs.view") {
scope = func(db *gorm.DB) *gorm.DB {
return db.
Joins("INNER JOIN flows f ON f.id = flow_id").
Where("f.id = ? AND f.user_id = ?", flowID, uid)
}
} else {
logger.FromContext(c).Errorf("error filtering user role permissions: permission not found")
response.Error(c, response.ErrNotPermitted, nil)
return
}
query.Init("agentlogs", agentlogsSQLMappers)
if query.Group != "" {
if _, ok := agentlogsSQLMappers[query.Group]; !ok {
logger.FromContext(c).Errorf("error finding agentlogs grouped: group field not found")
response.Error(c, response.ErrAgentlogsInvalidRequest, errors.New("group field not found"))
return
}
var respGrouped agentlogsGrouped
if respGrouped.Total, err = query.QueryGrouped(s.db, &respGrouped.Grouped, scope); err != nil {
logger.FromContext(c).WithError(err).Errorf("error finding agentlogs grouped")
response.Error(c, response.ErrInternal, err)
return
}
response.Success(c, http.StatusOK, respGrouped)
return
}
if resp.Total, err = query.Query(s.db, &resp.AgentLogs, scope); err != nil {
logger.FromContext(c).WithError(err).Errorf("error finding agentlogs")
response.Error(c, response.ErrInternal, err)
return
}
for i := 0; i < len(resp.AgentLogs); i++ {
if err = resp.AgentLogs[i].Valid(); err != nil {
logger.FromContext(c).WithError(err).Errorf("error validating agentlog data '%d'", resp.AgentLogs[i].ID)
response.Error(c, response.ErrAgentlogsInvalidData, err)
return
}
}
response.Success(c, http.StatusOK, resp)
}
File diff suppressed because it is too large Load Diff
+67
View File
@@ -0,0 +1,67 @@
package services
import (
"net/http"
"slices"
"pentagi/pkg/server/logger"
"pentagi/pkg/server/response"
"github.com/gin-gonic/gin"
"github.com/vxcontrol/cloud/anonymizer"
)
type anonymizeTextRequest struct {
Text string `json:"text" binding:"required"`
}
type anonymizeTextResponse struct {
Text string `json:"text"`
}
// AnonymizerService handles text anonymization via the REST API.
type AnonymizerService struct {
replacer anonymizer.Replacer
}
func NewAnonymizerService(replacer anonymizer.Replacer) *AnonymizerService {
return &AnonymizerService{replacer: replacer}
}
// AnonymizeText replaces sensitive data patterns in the provided text.
// @Summary Anonymize text
// @Description Replace sensitive data patterns (credentials, keys, PII, etc.) in the provided text with safe placeholders.
// @Tags Anonymize
// @Accept json
// @Produce json
// @Security BearerAuth
// @Param json body anonymizeTextRequest true "Text to anonymize"
// @Success 200 {object} response.successResp{data=anonymizeTextResponse} "anonymized text returned successfully"
// @Failure 400 {object} response.errorResp "invalid request body"
// @Failure 403 {object} response.errorResp "anonymize.call permission required"
// @Failure 503 {object} response.errorResp "anonymizer is not configured"
// @Router /anonymize/text [post]
func (s *AnonymizerService) AnonymizeText(c *gin.Context) {
privs := c.GetStringSlice("prm")
if !slices.Contains(privs, "anonymize.call") {
logger.FromContext(c).Errorf("error filtering user role permissions: permission not found")
response.Error(c, response.ErrNotPermitted, nil)
return
}
if s.replacer == nil {
response.Error(c, response.ErrAnonymizeUnavailable, nil)
return
}
var req anonymizeTextRequest
if err := c.ShouldBindJSON(&req); err != nil {
logger.FromContext(c).WithError(err).Error("error binding anonymize request")
response.Error(c, response.ErrAnonymizeInvalidRequest, err)
return
}
response.Success(c, http.StatusOK, anonymizeTextResponse{
Text: s.replacer.ReplaceString(req.Text),
})
}
+424
View File
@@ -0,0 +1,424 @@
package services
import (
"errors"
"net/http"
"time"
"pentagi/pkg/database"
"pentagi/pkg/graph/subscriptions"
"pentagi/pkg/server/auth"
"pentagi/pkg/server/logger"
"pentagi/pkg/server/models"
"pentagi/pkg/server/response"
"github.com/gin-gonic/gin"
"github.com/jinzhu/gorm"
)
type tokens struct {
Tokens []models.APIToken `json:"tokens"`
Total uint64 `json:"total"`
}
// TokenService handles API token management
type TokenService struct {
db *gorm.DB
globalSalt string
tokenCache *auth.TokenCache
ss subscriptions.SubscriptionsController
}
// NewTokenService creates a new TokenService instance
func NewTokenService(
db *gorm.DB,
globalSalt string,
tokenCache *auth.TokenCache,
ss subscriptions.SubscriptionsController,
) *TokenService {
return &TokenService{
db: db,
globalSalt: globalSalt,
tokenCache: tokenCache,
ss: ss,
}
}
// CreateToken creates a new API token
// @Summary Create new API token for automation
// @Tags Tokens
// @Accept json
// @Produce json
// @Param json body models.CreateAPITokenRequest true "Token creation request"
// @Success 201 {object} response.successResp{data=models.APITokenWithSecret} "token created successful"
// @Failure 400 {object} response.errorResp "invalid token request or default salt"
// @Failure 403 {object} response.errorResp "creating token not permitted"
// @Failure 500 {object} response.errorResp "internal error on creating token"
// @Router /tokens [post]
func (s *TokenService) CreateToken(c *gin.Context) {
// check for default salt
if s.globalSalt == "" || s.globalSalt == "salt" {
logger.FromContext(c).Errorf("token creation attempted with default salt")
response.Error(c, response.ErrTokenCreationDisabled, errors.New("token creation is disabled with default salt"))
return
}
uid := c.GetUint64("uid")
rid := c.GetUint64("rid")
uhash := c.GetString("uhash")
var req models.CreateAPITokenRequest
if err := c.ShouldBindJSON(&req); err != nil {
logger.FromContext(c).WithError(err).Errorf("error binding JSON")
response.Error(c, response.ErrTokenInvalidRequest, err)
return
}
if err := req.Valid(); err != nil {
logger.FromContext(c).WithError(err).Errorf("error validating JSON")
response.Error(c, response.ErrTokenInvalidRequest, err)
return
}
// check if name is unique for this user (if provided)
if req.Name != nil && *req.Name != "" {
var existing models.APIToken
err := s.db.
Where("user_id = ? AND name = ? AND deleted_at IS NULL", uid, *req.Name).
First(&existing).
Error
if err == nil {
logger.FromContext(c).Errorf("token with name '%s' already exists for user %d", *req.Name, uid)
response.Error(c, response.ErrTokenInvalidRequest, errors.New("token with this name already exists"))
return
}
}
// generate token_id
tokenID, err := auth.GenerateTokenID()
if err != nil {
logger.FromContext(c).WithError(err).Errorf("error generating token ID")
response.Error(c, response.ErrInternal, err)
return
}
// create JWT claims
claims := auth.MakeAPITokenClaims(tokenID, uhash, uid, rid, req.TTL)
// sign token
token, err := auth.MakeAPIToken(s.globalSalt, claims)
if err != nil {
logger.FromContext(c).WithError(err).Errorf("error signing token")
response.Error(c, response.ErrInternal, err)
return
}
// save to database
apiToken := models.APIToken{
TokenID: tokenID,
UserID: uid,
RoleID: rid,
Name: req.Name,
TTL: req.TTL,
Status: models.TokenStatusActive,
}
if err := s.db.Create(&apiToken).Error; err != nil {
logger.FromContext(c).WithError(err).Errorf("error creating token in database")
response.Error(c, response.ErrInternal, err)
return
}
result := models.APITokenWithSecret{
APIToken: apiToken,
Token: token,
}
// invalidate cache for negative caching results
s.tokenCache.Invalidate(apiToken.TokenID)
s.tokenCache.InvalidateUser(apiToken.UserID)
if s.ss != nil {
publisher := s.ss.NewAPITokenPublisher(int64(apiToken.UserID))
publisher.APITokenCreated(c, database.APITokenWithSecret{
ApiToken: convertAPITokenToDatabase(apiToken),
Token: token,
})
}
response.Success(c, http.StatusCreated, result)
}
// ListTokens returns a list of tokens (user sees only their own, admin sees all)
// @Summary List API tokens
// @Tags Tokens
// @Produce json
// @Success 200 {object} response.successResp{data=tokens} "tokens retrieved successful"
// @Failure 403 {object} response.errorResp "listing tokens not permitted"
// @Failure 500 {object} response.errorResp "internal error on listing tokens"
// @Router /tokens [get]
func (s *TokenService) ListTokens(c *gin.Context) {
uid := c.GetUint64("uid")
prms := c.GetStringSlice("prm")
query := s.db.Where("deleted_at IS NULL")
// check if user has admin privilege
hasAdmin := auth.LookupPerm(prms, "settings.tokens.admin")
if !hasAdmin {
// regular user sees only their own tokens
query = query.Where("user_id = ?", uid)
}
var tokenList []models.APIToken
var total uint64
if err := query.Order("created_at DESC").Find(&tokenList).Count(&total).Error; err != nil {
logger.FromContext(c).WithError(err).Errorf("error finding tokens")
response.Error(c, response.ErrInternal, err)
return
}
for i := range tokenList {
token := &tokenList[i]
isExpired := token.CreatedAt.Add(time.Duration(token.TTL) * time.Second).Before(time.Now())
if token.Status == models.TokenStatusActive && isExpired {
token.Status = models.TokenStatusExpired
}
}
result := tokens{
Tokens: tokenList,
Total: total,
}
response.Success(c, http.StatusOK, result)
}
// GetToken returns information about a specific token
// @Summary Get API token details
// @Tags Tokens
// @Produce json
// @Param tokenID path string true "Token ID"
// @Success 200 {object} response.successResp{data=models.APIToken} "token retrieved successful"
// @Failure 403 {object} response.errorResp "accessing token not permitted"
// @Failure 404 {object} response.errorResp "token not found"
// @Failure 500 {object} response.errorResp "internal error on getting token"
// @Router /tokens/{tokenID} [get]
func (s *TokenService) GetToken(c *gin.Context) {
uid := c.GetUint64("uid")
prms := c.GetStringSlice("prm")
tokenID := c.Param("tokenID")
var token models.APIToken
if err := s.db.Where("token_id = ? AND deleted_at IS NULL", tokenID).First(&token).Error; err != nil {
logger.FromContext(c).WithError(err).Errorf("error finding token")
if errors.Is(err, gorm.ErrRecordNotFound) {
response.Error(c, response.ErrTokenNotFound, err)
} else {
response.Error(c, response.ErrInternal, err)
}
return
}
// check authorization
hasAdmin := auth.LookupPerm(prms, "settings.tokens.admin")
if !hasAdmin && token.UserID != uid {
logger.FromContext(c).Errorf("user %d attempted to access token of user %d", uid, token.UserID)
response.Error(c, response.ErrTokenUnauthorized, errors.New("not authorized to access this token"))
return
}
isExpired := token.CreatedAt.Add(time.Duration(token.TTL) * time.Second).Before(time.Now())
if token.Status == models.TokenStatusActive && isExpired {
token.Status = models.TokenStatusExpired
}
if err := token.Valid(); err != nil {
logger.FromContext(c).WithError(err).Errorf("error validating token data")
response.Error(c, response.ErrTokenInvalidData, err)
return
}
response.Success(c, http.StatusOK, token)
}
// UpdateToken updates name and/or status of a token
// @Summary Update API token
// @Tags Tokens
// @Accept json
// @Produce json
// @Param tokenID path string true "Token ID"
// @Param json body models.UpdateAPITokenRequest true "Token update request"
// @Success 200 {object} response.successResp{data=models.APIToken} "token updated successful"
// @Failure 400 {object} response.errorResp "invalid update request"
// @Failure 403 {object} response.errorResp "updating token not permitted"
// @Failure 404 {object} response.errorResp "token not found"
// @Failure 500 {object} response.errorResp "internal error on updating token"
// @Router /tokens/{tokenID} [put]
func (s *TokenService) UpdateToken(c *gin.Context) {
uid := c.GetUint64("uid")
prms := c.GetStringSlice("prm")
tokenID := c.Param("tokenID")
var req models.UpdateAPITokenRequest
if err := c.ShouldBindJSON(&req); err != nil {
logger.FromContext(c).WithError(err).Errorf("error binding JSON")
response.Error(c, response.ErrTokenInvalidRequest, err)
return
}
if err := req.Valid(); err != nil {
logger.FromContext(c).WithError(err).Errorf("error validating JSON")
response.Error(c, response.ErrTokenInvalidRequest, err)
return
}
var token models.APIToken
if err := s.db.Where("token_id = ? AND deleted_at IS NULL", tokenID).First(&token).Error; err != nil {
logger.FromContext(c).WithError(err).Errorf("error finding token")
if errors.Is(err, gorm.ErrRecordNotFound) {
response.Error(c, response.ErrTokenNotFound, err)
} else {
response.Error(c, response.ErrInternal, err)
}
return
}
// check authorization
hasAdmin := auth.LookupPerm(prms, "settings.tokens.admin")
if !hasAdmin && token.UserID != uid {
logger.FromContext(c).Errorf("user %d attempted to update token of user %d", uid, token.UserID)
response.Error(c, response.ErrTokenUnauthorized, errors.New("not authorized to update this token"))
return
}
// update fields
updates := make(map[string]any)
if req.Name != nil {
// check uniqueness if name is changing
if token.Name == nil || *token.Name != *req.Name {
if *req.Name != "" {
var existing models.APIToken
err := s.db.
Where("user_id = ? AND name = ? AND token_id != ? AND deleted_at IS NULL", token.UserID, *req.Name, tokenID).
First(&existing).
Error
if err == nil {
logger.FromContext(c).Errorf("token with name '%s' already exists for user %d", *req.Name, token.UserID)
response.Error(c, response.ErrTokenInvalidRequest, errors.New("token with this name already exists"))
return
}
}
}
updates["name"] = req.Name
}
switch req.Status {
case models.TokenStatusActive:
updates["status"] = models.TokenStatusActive
case models.TokenStatusRevoked:
updates["status"] = models.TokenStatusRevoked
case models.TokenStatusExpired:
updates["status"] = models.TokenStatusRevoked
}
if len(updates) > 0 {
if err := s.db.Model(&token).Updates(updates).Error; err != nil {
logger.FromContext(c).WithError(err).Errorf("error updating token")
response.Error(c, response.ErrInternal, err)
return
}
// invalidate cache if status changed
if req.Status != "" {
s.tokenCache.Invalidate(tokenID)
// also invalidate all tokens for this user (in case of role change or security event)
s.tokenCache.InvalidateUser(token.UserID)
}
// reload token
if err := s.db.Where("token_id = ?", tokenID).First(&token).Error; err != nil {
logger.FromContext(c).WithError(err).Errorf("error reloading token")
response.Error(c, response.ErrInternal, err)
return
}
}
isExpired := token.CreatedAt.Add(time.Duration(token.TTL) * time.Second).Before(time.Now())
if token.Status == models.TokenStatusActive && isExpired {
token.Status = models.TokenStatusExpired
}
if s.ss != nil {
publisher := s.ss.NewAPITokenPublisher(int64(token.UserID))
publisher.APITokenUpdated(c, convertAPITokenToDatabase(token))
}
response.Success(c, http.StatusOK, token)
}
// DeleteToken performs soft delete of a token
// @Summary Delete API token
// @Tags Tokens
// @Produce json
// @Param tokenID path string true "Token ID"
// @Success 200 {object} response.successResp "token deleted successful"
// @Failure 403 {object} response.errorResp "deleting token not permitted"
// @Failure 404 {object} response.errorResp "token not found"
// @Failure 500 {object} response.errorResp "internal error on deleting token"
// @Router /tokens/{tokenID} [delete]
func (s *TokenService) DeleteToken(c *gin.Context) {
uid := c.GetUint64("uid")
prms := c.GetStringSlice("prm")
tokenID := c.Param("tokenID")
var token models.APIToken
if err := s.db.Where("token_id = ? AND deleted_at IS NULL", tokenID).First(&token).Error; err != nil {
logger.FromContext(c).WithError(err).Errorf("error finding token")
if errors.Is(err, gorm.ErrRecordNotFound) {
response.Error(c, response.ErrTokenNotFound, err)
} else {
response.Error(c, response.ErrInternal, err)
}
return
}
// check authorization
hasAdmin := auth.LookupPerm(prms, "settings.tokens.admin")
if !hasAdmin && token.UserID != uid {
logger.FromContext(c).Errorf("user %d attempted to delete token of user %d", uid, token.UserID)
response.Error(c, response.ErrTokenUnauthorized, errors.New("not authorized to delete this token"))
return
}
// soft delete
if err := s.db.Delete(&token).Error; err != nil {
logger.FromContext(c).WithError(err).Errorf("error deleting token")
response.Error(c, response.ErrInternal, err)
return
}
// invalidate cache for this token and all user's tokens
s.tokenCache.Invalidate(tokenID)
s.tokenCache.InvalidateUser(token.UserID)
if s.ss != nil {
publisher := s.ss.NewAPITokenPublisher(int64(token.UserID))
publisher.APITokenDeleted(c, convertAPITokenToDatabase(token))
}
response.Success(c, http.StatusOK, gin.H{"message": "token deleted successfully"})
}
func convertAPITokenToDatabase(apiToken models.APIToken) database.ApiToken {
return database.ApiToken{
ID: int64(apiToken.ID),
TokenID: apiToken.TokenID,
UserID: int64(apiToken.UserID),
RoleID: int64(apiToken.RoleID),
Name: database.StringToNullString(*apiToken.Name),
Ttl: int64(apiToken.TTL),
Status: database.TokenStatus(apiToken.Status),
CreatedAt: database.TimeToNullTime(apiToken.CreatedAt),
UpdatedAt: database.TimeToNullTime(apiToken.UpdatedAt),
DeletedAt: database.PtrTimeToNullTime(apiToken.DeletedAt),
}
}
File diff suppressed because it is too large Load Diff
@@ -0,0 +1,219 @@
package services
import (
"errors"
"net/http"
"slices"
"strconv"
"pentagi/pkg/server/logger"
"pentagi/pkg/server/models"
"pentagi/pkg/server/rdb"
"pentagi/pkg/server/response"
"github.com/gin-gonic/gin"
"github.com/jinzhu/gorm"
)
type assistantlogs struct {
AssistantLogs []models.Assistantlog `json:"assistantlogs"`
Total uint64 `json:"total"`
}
type assistantlogsGrouped struct {
Grouped []string `json:"grouped"`
Total uint64 `json:"total"`
}
var assistantlogsSQLMappers = map[string]any{
"id": "{{table}}.id",
"type": "{{table}}.type",
"message": "{{table}}.message",
"result": "{{table}}.result",
"result_format": "{{table}}.result_format",
"flow_id": "{{table}}.flow_id",
"assistant_id": "{{table}}.assistant_id",
"created_at": "{{table}}.created_at",
"data": "({{table}}.type || ' ' || {{table}}.message || ' ' || {{table}}.result)",
}
type AssistantlogService struct {
db *gorm.DB
}
func NewAssistantlogService(db *gorm.DB) *AssistantlogService {
return &AssistantlogService{
db: db,
}
}
// GetAssistantlogs is a function to return assistantlogs list
// @Summary Retrieve assistantlogs list
// @Tags Assistantlogs
// @Produce json
// @Security BearerAuth
// @Param request query rdb.TableQuery true "query table params"
// @Success 200 {object} response.successResp{data=assistantlogs} "assistantlogs list received successful"
// @Failure 400 {object} response.errorResp "invalid query request data"
// @Failure 403 {object} response.errorResp "getting assistantlogs not permitted"
// @Failure 500 {object} response.errorResp "internal error on getting assistantlogs"
// @Router /assistantlogs/ [get]
func (s *AssistantlogService) GetAssistantlogs(c *gin.Context) {
var (
err error
query rdb.TableQuery
resp assistantlogs
)
if err = c.ShouldBindQuery(&query); err != nil {
logger.FromContext(c).WithError(err).Errorf("error binding query")
response.Error(c, response.ErrAssistantlogsInvalidRequest, err)
return
}
uid := c.GetUint64("uid")
privs := c.GetStringSlice("prm")
var scope func(db *gorm.DB) *gorm.DB
if slices.Contains(privs, "assistantlogs.admin") {
scope = func(db *gorm.DB) *gorm.DB {
return db.
Joins("INNER JOIN flows f ON f.id = flow_id")
}
} else if slices.Contains(privs, "assistantlogs.view") {
scope = func(db *gorm.DB) *gorm.DB {
return db.
Joins("INNER JOIN flows f ON f.id = flow_id").
Where("f.user_id = ?", uid)
}
} else {
logger.FromContext(c).Errorf("error filtering user role permissions: permission not found")
response.Error(c, response.ErrNotPermitted, nil)
return
}
query.Init("assistantlogs", assistantlogsSQLMappers)
if query.Group != "" {
if _, ok := assistantlogsSQLMappers[query.Group]; !ok {
logger.FromContext(c).Errorf("error finding assistantlogs grouped: group field not found")
response.Error(c, response.ErrAssistantlogsInvalidRequest, errors.New("group field not found"))
return
}
var respGrouped assistantlogsGrouped
if respGrouped.Total, err = query.QueryGrouped(s.db, &respGrouped.Grouped, scope); err != nil {
logger.FromContext(c).WithError(err).Errorf("error finding assistantlogs grouped")
response.Error(c, response.ErrInternal, err)
return
}
response.Success(c, http.StatusOK, respGrouped)
return
}
if resp.Total, err = query.Query(s.db, &resp.AssistantLogs, scope); err != nil {
logger.FromContext(c).WithError(err).Errorf("error finding assistantlogs")
response.Error(c, response.ErrInternal, err)
return
}
for i := 0; i < len(resp.AssistantLogs); i++ {
if err = resp.AssistantLogs[i].Valid(); err != nil {
logger.FromContext(c).WithError(err).Errorf("error validating assistantlog data '%d'", resp.AssistantLogs[i].ID)
response.Error(c, response.ErrAssistantlogsInvalidData, err)
return
}
}
response.Success(c, http.StatusOK, resp)
}
// GetFlowAssistantlogs is a function to return assistantlogs list by flow id
// @Summary Retrieve assistantlogs list by flow id
// @Tags Assistantlogs
// @Produce json
// @Security BearerAuth
// @Param flowID path int true "flow id" minimum(0)
// @Param request query rdb.TableQuery true "query table params"
// @Success 200 {object} response.successResp{data=assistantlogs} "assistantlogs list received successful"
// @Failure 400 {object} response.errorResp "invalid query request data"
// @Failure 403 {object} response.errorResp "getting assistantlogs not permitted"
// @Failure 500 {object} response.errorResp "internal error on getting assistantlogs"
// @Router /flows/{flowID}/assistantlogs/ [get]
func (s *AssistantlogService) GetFlowAssistantlogs(c *gin.Context) {
var (
err error
flowID uint64
query rdb.TableQuery
resp assistantlogs
)
if flowID, err = strconv.ParseUint(c.Param("flowID"), 10, 64); err != nil {
logger.FromContext(c).WithError(err).Errorf("error parsing flow id")
response.Error(c, response.ErrAssistantlogsInvalidRequest, err)
return
}
if err = c.ShouldBindQuery(&query); err != nil {
logger.FromContext(c).WithError(err).Errorf("error binding query")
response.Error(c, response.ErrAssistantlogsInvalidRequest, err)
return
}
uid := c.GetUint64("uid")
privs := c.GetStringSlice("prm")
var scope func(db *gorm.DB) *gorm.DB
if slices.Contains(privs, "assistantlogs.admin") {
scope = func(db *gorm.DB) *gorm.DB {
return db.
Joins("INNER JOIN flows f ON f.id = flow_id").
Where("f.id = ?", flowID)
}
} else if slices.Contains(privs, "assistantlogs.view") {
scope = func(db *gorm.DB) *gorm.DB {
return db.
Joins("INNER JOIN flows f ON f.id = flow_id").
Where("f.id = ? AND f.user_id = ?", flowID, uid)
}
} else {
logger.FromContext(c).Errorf("error filtering user role permissions: permission not found")
response.Error(c, response.ErrNotPermitted, nil)
return
}
query.Init("assistantlogs", assistantlogsSQLMappers)
if query.Group != "" {
if _, ok := assistantlogsSQLMappers[query.Group]; !ok {
logger.FromContext(c).Errorf("error finding assistantlogs grouped: group field not found")
response.Error(c, response.ErrAssistantlogsInvalidRequest, errors.New("group field not found"))
return
}
var respGrouped assistantlogsGrouped
if respGrouped.Total, err = query.QueryGrouped(s.db, &respGrouped.Grouped, scope); err != nil {
logger.FromContext(c).WithError(err).Errorf("error finding assistantlogs grouped")
response.Error(c, response.ErrInternal, err)
return
}
response.Success(c, http.StatusOK, respGrouped)
return
}
if resp.Total, err = query.Query(s.db, &resp.AssistantLogs, scope); err != nil {
logger.FromContext(c).WithError(err).Errorf("error finding assistantlogs")
response.Error(c, response.ErrInternal, err)
return
}
for i := 0; i < len(resp.AssistantLogs); i++ {
if err = resp.AssistantLogs[i].Valid(); err != nil {
logger.FromContext(c).WithError(err).Errorf("error validating assistantlog data '%d'", resp.AssistantLogs[i].ID)
response.Error(c, response.ErrAssistantlogsInvalidData, err)
return
}
}
response.Success(c, http.StatusOK, resp)
}
+684
View File
@@ -0,0 +1,684 @@
package services
import (
"encoding/json"
"errors"
"fmt"
"net/http"
"slices"
"strconv"
"pentagi/pkg/controller"
"pentagi/pkg/database"
"pentagi/pkg/graph/subscriptions"
"pentagi/pkg/providers"
"pentagi/pkg/providers/provider"
"pentagi/pkg/server/logger"
"pentagi/pkg/server/models"
"pentagi/pkg/server/rdb"
"pentagi/pkg/server/response"
"github.com/gin-gonic/gin"
"github.com/jinzhu/gorm"
)
type assistants struct {
Assistants []models.Assistant `json:"assistants"`
Total uint64 `json:"total"`
}
type assistantsGrouped struct {
Grouped []string `json:"grouped"`
Total uint64 `json:"total"`
}
var assistantsSQLMappers = map[string]any{
"id": "{{table}}.id",
"status": "{{table}}.status",
"title": "{{table}}.title",
"model": "{{table}}.model",
"model_provider_name": "{{table}}.model_provider_name",
"model_provider_type": "{{table}}.model_provider_type",
"language": "{{table}}.language",
"flow_id": "{{table}}.flow_id",
"msgchain_id": "{{table}}.msgchain_id",
"created_at": "{{table}}.created_at",
"updated_at": "{{table}}.updated_at",
"data": "({{table}}.status || ' ' || {{table}}.title || ' ' || {{table}}.flow_id)",
}
type AssistantService struct {
db *gorm.DB
pc providers.ProviderController
fc controller.FlowController
ss subscriptions.SubscriptionsController
}
func NewAssistantService(
db *gorm.DB,
pc providers.ProviderController,
fc controller.FlowController,
ss subscriptions.SubscriptionsController,
) *AssistantService {
return &AssistantService{
db: db,
pc: pc,
fc: fc,
ss: ss,
}
}
// GetAssistants is a function to return assistants list
// @Summary Retrieve assistants list
// @Tags Assistants
// @Produce json
// @Security BearerAuth
// @Param flowID path int true "flow id" minimum(0)
// @Param request query rdb.TableQuery true "query table params"
// @Success 200 {object} response.successResp{data=assistants} "assistants list received successful"
// @Failure 400 {object} response.errorResp "invalid query request data"
// @Failure 403 {object} response.errorResp "getting assistants not permitted"
// @Failure 500 {object} response.errorResp "internal error on getting assistants"
// @Router /flows/{flowID}/assistants/ [get]
func (s *AssistantService) GetFlowAssistants(c *gin.Context) {
var (
err error
flowID uint64
query rdb.TableQuery
resp assistants
)
if err = c.ShouldBindQuery(&query); err != nil {
logger.FromContext(c).WithError(err).Errorf("error binding query")
response.Error(c, response.ErrAssistantsInvalidRequest, err)
return
}
if flowID, err = strconv.ParseUint(c.Param("flowID"), 10, 64); err != nil {
logger.FromContext(c).WithError(err).Errorf("error parsing flow id")
response.Error(c, response.ErrAssistantsInvalidRequest, err)
return
}
uid := c.GetUint64("uid")
privs := c.GetStringSlice("prm")
var scope func(db *gorm.DB) *gorm.DB
if slices.Contains(privs, "assistants.admin") {
scope = func(db *gorm.DB) *gorm.DB {
return db.
Joins("INNER JOIN flows f ON f.id = assistants.flow_id").
Where("f.id = ?", flowID)
}
} else if slices.Contains(privs, "assistants.view") {
scope = func(db *gorm.DB) *gorm.DB {
return db.
Joins("INNER JOIN flows f ON f.id = assistants.flow_id").
Where("f.id = ? AND f.user_id = ?", flowID, uid)
}
} else {
logger.FromContext(c).Errorf("error filtering user role permissions: permission not found")
response.Error(c, response.ErrNotPermitted, nil)
return
}
query.Init("assistants", assistantsSQLMappers)
if query.Group != "" {
if _, ok := assistantsSQLMappers[query.Group]; !ok {
logger.FromContext(c).Errorf("error finding assistants grouped: group field not found")
response.Error(c, response.ErrAssistantsInvalidRequest, errors.New("group field not found"))
return
}
var respGrouped assistantsGrouped
if respGrouped.Total, err = query.QueryGrouped(s.db, &respGrouped.Grouped, scope); err != nil {
logger.FromContext(c).WithError(err).Errorf("error finding assistants grouped")
response.Error(c, response.ErrInternal, err)
return
}
response.Success(c, http.StatusOK, respGrouped)
return
}
if resp.Total, err = query.Query(s.db, &resp.Assistants, scope); err != nil {
logger.FromContext(c).WithError(err).Errorf("error finding assistants")
response.Error(c, response.ErrInternal, err)
return
}
for i := 0; i < len(resp.Assistants); i++ {
if err = resp.Assistants[i].Valid(); err != nil {
logger.FromContext(c).WithError(err).Errorf("error validating assistant data '%d'", resp.Assistants[i].ID)
response.Error(c, response.ErrAssistantsInvalidData, err)
return
}
}
response.Success(c, http.StatusOK, resp)
}
// GetFlowAssistant is a function to return flow assistant by id
// @Summary Retrieve flow assistant by id
// @Tags Assistants
// @Produce json
// @Security BearerAuth
// @Param flowID path int true "flow id" minimum(0)
// @Param assistantID path int true "assistant id" minimum(0)
// @Success 200 {object} response.successResp{data=models.Assistant} "flow assistant received successful"
// @Failure 403 {object} response.errorResp "getting flow assistant not permitted"
// @Failure 404 {object} response.errorResp "flow assistant not found"
// @Failure 500 {object} response.errorResp "internal error on getting flow assistant"
// @Router /flows/{flowID}/assistants/{assistantID} [get]
func (s *AssistantService) GetFlowAssistant(c *gin.Context) {
var (
err error
flowID uint64
assistantID uint64
resp models.Assistant
)
if flowID, err = strconv.ParseUint(c.Param("flowID"), 10, 64); err != nil {
logger.FromContext(c).WithError(err).Errorf("error parsing flow id")
response.Error(c, response.ErrAssistantsInvalidRequest, err)
return
}
if assistantID, err = strconv.ParseUint(c.Param("assistantID"), 10, 64); err != nil {
logger.FromContext(c).WithError(err).Errorf("error parsing assistant id")
response.Error(c, response.ErrAssistantsInvalidRequest, err)
return
}
uid := c.GetUint64("uid")
privs := c.GetStringSlice("prm")
var scope func(db *gorm.DB) *gorm.DB
if slices.Contains(privs, "assistants.admin") {
scope = func(db *gorm.DB) *gorm.DB {
return db.
Joins("INNER JOIN flows f ON f.id = assistants.flow_id").
Where("f.id = ?", flowID)
}
} else if slices.Contains(privs, "assistants.view") {
scope = func(db *gorm.DB) *gorm.DB {
return db.
Joins("INNER JOIN flows f ON f.id = assistants.flow_id").
Where("f.id = ? AND f.user_id = ?", flowID, uid)
}
} else {
logger.FromContext(c).Errorf("error filtering user role permissions: permission not found")
response.Error(c, response.ErrNotPermitted, nil)
return
}
err = s.db.Model(&resp).
Scopes(scope).
Where("assistants.id = ?", assistantID).
Take(&resp).Error
if err != nil {
logger.FromContext(c).WithError(err).Errorf("error on getting flow assistant by id")
if gorm.IsRecordNotFoundError(err) {
response.Error(c, response.ErrAssistantsNotFound, err)
} else {
response.Error(c, response.ErrInternal, err)
}
return
}
response.Success(c, http.StatusOK, resp)
}
// CreateFlowAssistant is a function to create new assistant with custom functions
// @Summary Create new assistant with custom functions
// @Tags Assistants
// @Accept json
// @Produce json
// @Security BearerAuth
// @Param flowID path int true "flow id (use 0 to create a new flow together with the assistant)" minimum(0)
// @Param json body models.CreateAssistant true "assistant model to create"
// @Success 201 {object} response.successResp{data=models.AssistantFlow} "assistant created successful"
// @Failure 400 {object} response.errorResp "invalid assistant request data"
// @Failure 403 {object} response.errorResp "creating assistant not permitted"
// @Failure 500 {object} response.errorResp "internal error on creating assistant"
// @Router /flows/{flowID}/assistants/ [post]
func (s *AssistantService) CreateFlowAssistant(c *gin.Context) {
var (
err error
flowID uint64
resp models.AssistantFlow
createAssistant models.CreateAssistant
)
if err := c.ShouldBindJSON(&createAssistant); err != nil {
logger.FromContext(c).WithError(err).Errorf("error binding JSON")
response.Error(c, response.ErrAssistantsInvalidRequest, err)
return
}
if err := createAssistant.Valid(); err != nil {
logger.FromContext(c).WithError(err).Errorf("error validating assistant data")
response.Error(c, response.ErrAssistantsInvalidData, err)
return
}
if flowID, err = strconv.ParseUint(c.Param("flowID"), 10, 64); err != nil {
logger.FromContext(c).WithError(err).Errorf("error parsing flow id")
response.Error(c, response.ErrAssistantsInvalidRequest, err)
return
}
uid := c.GetUint64("uid")
privs := c.GetStringSlice("prm")
if flowID == 0 {
// Creating a new flow together with the assistant requires both permissions.
if !slices.Contains(privs, "assistants.create") || !slices.Contains(privs, "flows.create") {
logger.FromContext(c).Errorf("error filtering user role permissions: permission not found")
response.Error(c, response.ErrNotPermitted, nil)
return
}
} else {
if !slices.Contains(privs, "assistants.create") {
logger.FromContext(c).Errorf("error filtering user role permissions: permission not found")
response.Error(c, response.ErrNotPermitted, nil)
return
}
// Verify the flow exists and belongs to the user (unless admin).
var flow models.Flow
var flowScope func(db *gorm.DB) *gorm.DB
if slices.Contains(privs, "flows.admin") {
flowScope = func(db *gorm.DB) *gorm.DB { return db.Where("id = ?", flowID) }
} else {
flowScope = func(db *gorm.DB) *gorm.DB { return db.Where("id = ? AND user_id = ?", flowID, uid) }
}
if err = s.db.Model(&flow).Scopes(flowScope).Take(&flow).Error; err != nil {
logger.FromContext(c).WithError(err).Errorf("error getting flow by id")
if gorm.IsRecordNotFoundError(err) {
response.Error(c, response.ErrFlowsNotFound, err)
} else {
response.Error(c, response.ErrInternal, err)
}
return
}
}
prvname := provider.ProviderName(createAssistant.Provider)
prv, err := s.pc.GetProvider(c, prvname, int64(uid))
if err != nil {
logger.FromContext(c).WithError(err).Errorf("error getting provider: not found")
response.Error(c, response.ErrInternal, err)
return
}
prvtype := prv.Type()
dbResources, err := validateServiceResources(s.db, uid, privs, createAssistant.ResourceIDs)
if err != nil {
logger.FromContext(c).WithError(err).Errorf("error validating resource ids")
response.Error(c, response.ErrAssistantsInvalidRequest, err)
return
}
aw, err := s.fc.CreateAssistant(
c,
int64(uid),
int64(flowID),
createAssistant.Input,
createAssistant.UseAgents,
prvname,
prvtype,
createAssistant.Functions,
dbResources,
)
if err != nil {
logger.FromContext(c).WithError(err).Errorf("error creating assistant")
response.Error(c, response.ErrInternal, err)
return
}
if err = s.db.Model(&resp.Assistant).Where("id = ?", aw.GetAssistantID()).Take(&resp.Assistant).Error; err != nil {
logger.FromContext(c).WithError(err).Errorf("error getting assistant by id")
response.Error(c, response.ErrInternal, err)
return
}
if err = s.db.Model(&resp.Flow).Where("id = ?", resp.Assistant.FlowID).Take(&resp.Flow).Error; err != nil {
logger.FromContext(c).WithError(err).Errorf("error getting flow by id")
response.Error(c, response.ErrInternal, err)
return
}
response.Success(c, http.StatusCreated, resp)
}
// PatchAssistant is a function to patch assistant
// @Summary Patch assistant
// @Tags Assistants
// @Accept json
// @Produce json
// @Security BearerAuth
// @Param flowID path int true "flow id" minimum(0)
// @Param assistantID path int true "assistant id" minimum(0)
// @Param json body models.PatchAssistant true "assistant model to patch"
// @Success 200 {object} response.successResp{data=models.AssistantFlow} "assistant patched successful"
// @Failure 400 {object} response.errorResp "invalid assistant request data"
// @Failure 403 {object} response.errorResp "patching assistant not permitted"
// @Failure 500 {object} response.errorResp "internal error on patching assistant"
// @Router /flows/{flowID}/assistants/{assistantID} [put]
func (s *AssistantService) PatchAssistant(c *gin.Context) {
var (
err error
flowID uint64
assistant models.AssistantFlow
assistantID uint64
patchAssistant models.PatchAssistant
)
if err := c.ShouldBindJSON(&patchAssistant); err != nil {
logger.FromContext(c).WithError(err).Errorf("error binding JSON")
response.Error(c, response.ErrAssistantsInvalidRequest, err)
return
}
if err := patchAssistant.Valid(); err != nil {
logger.FromContext(c).WithError(err).Errorf("error validating assistant data")
response.Error(c, response.ErrAssistantsInvalidData, err)
return
}
flowID, err = strconv.ParseUint(c.Param("flowID"), 10, 64)
if err != nil {
logger.FromContext(c).WithError(err).Errorf("error parsing flow id")
response.Error(c, response.ErrAssistantsInvalidRequest, err)
return
}
assistantID, err = strconv.ParseUint(c.Param("assistantID"), 10, 64)
if err != nil {
logger.FromContext(c).WithError(err).Errorf("error parsing assistant id")
response.Error(c, response.ErrAssistantsInvalidRequest, err)
return
}
uid := c.GetUint64("uid")
privs := c.GetStringSlice("prm")
var scope func(db *gorm.DB) *gorm.DB
if slices.Contains(privs, "assistants.admin") {
scope = func(db *gorm.DB) *gorm.DB {
return db.Where("id = ?", assistantID)
}
} else if slices.Contains(privs, "assistants.edit") {
scope = func(db *gorm.DB) *gorm.DB {
return db.Where("assistants.id = ? AND assistants.flow_id = ?", assistantID, flowID).
Joins("INNER JOIN flows f ON f.id = assistants.flow_id").
Where("f.user_id = ?", uid)
}
} else {
logger.FromContext(c).Errorf("error filtering user role permissions: permission not found")
response.Error(c, response.ErrNotPermitted, nil)
return
}
if err = s.db.Model(&assistant).Scopes(scope).Take(&assistant).Error; err != nil {
logger.FromContext(c).WithError(err).Errorf("error getting assistant by id")
if gorm.IsRecordNotFoundError(err) {
response.Error(c, response.ErrAssistantsNotFound, err)
} else {
response.Error(c, response.ErrInternal, err)
}
return
}
fw, err := s.fc.GetFlow(c, int64(flowID))
if err != nil {
logger.FromContext(c).WithError(err).Errorf("error getting flow by id in flow controller")
response.Error(c, response.ErrInternal, err)
return
}
aw, err := fw.GetAssistant(c, int64(assistantID))
if err != nil {
logger.FromContext(c).WithError(err).Errorf("error getting assistant by id in flow controller")
response.Error(c, response.ErrInternal, err)
return
}
switch patchAssistant.Action {
case "stop":
if err := aw.Stop(c); err != nil {
logger.FromContext(c).WithError(err).Errorf("error stopping assistant")
response.Error(c, response.ErrInternal, err)
return
}
case "input":
if patchAssistant.Input == nil || *patchAssistant.Input == "" {
logger.FromContext(c).Errorf("error sending input to assistant: input is empty")
response.Error(c, response.ErrAssistantsInvalidRequest, nil)
return
}
dbResources, err := validateServiceResources(s.db, uid, privs, patchAssistant.ResourceIDs)
if err != nil {
logger.FromContext(c).WithError(err).Errorf("error validating resource ids")
response.Error(c, response.ErrAssistantsInvalidRequest, err)
return
}
if err := aw.PutInput(c, *patchAssistant.Input, patchAssistant.UseAgents, dbResources); err != nil {
logger.FromContext(c).WithError(err).Errorf("error sending input to assistant")
response.Error(c, response.ErrInternal, err)
return
}
default:
logger.FromContext(c).Errorf("error filtering assistant action")
response.Error(c, response.ErrAssistantsInvalidRequest, nil)
return
}
if err = s.db.Model(&assistant).Scopes(scope).Take(&assistant).Error; err != nil {
logger.FromContext(c).WithError(err).Errorf("error getting assistant by id")
if gorm.IsRecordNotFoundError(err) {
response.Error(c, response.ErrAssistantsNotFound, err)
} else {
response.Error(c, response.ErrInternal, err)
}
return
}
assistantDB, err := convertAssistantToDatabase(assistant.Assistant)
if err != nil {
logger.FromContext(c).WithError(err).Errorf("error converting assistant to database")
response.Error(c, response.ErrInternal, err)
return
}
if s.ss != nil {
publisher := s.ss.NewFlowPublisher(int64(assistant.Flow.UserID), int64(assistant.FlowID))
publisher.AssistantUpdated(c, assistantDB)
}
response.Success(c, http.StatusOK, assistant)
}
// DeleteAssistant is a function to delete assistant by id
// @Summary Delete assistant by id
// @Tags Assistants
// @Security BearerAuth
// @Param flowID path int true "flow id" minimum(0)
// @Param assistantID path int true "assistant id" minimum(0)
// @Success 200 {object} response.successResp{data=models.AssistantFlow} "assistant deleted successful"
// @Failure 403 {object} response.errorResp "deleting assistant not permitted"
// @Failure 404 {object} response.errorResp "assistant not found"
// @Failure 500 {object} response.errorResp "internal error on deleting assistant"
// @Router /flows/{flowID}/assistants/{assistantID} [delete]
func (s *AssistantService) DeleteAssistant(c *gin.Context) {
var (
err error
flowID uint64
assistant models.AssistantFlow
assistantID uint64
)
flowID, err = strconv.ParseUint(c.Param("flowID"), 10, 64)
if err != nil {
logger.FromContext(c).WithError(err).Errorf("error parsing flow id")
response.Error(c, response.ErrAssistantsInvalidRequest, err)
return
}
assistantID, err = strconv.ParseUint(c.Param("assistantID"), 10, 64)
if err != nil {
logger.FromContext(c).WithError(err).Errorf("error parsing assistant id")
response.Error(c, response.ErrAssistantsInvalidRequest, err)
return
}
uid := c.GetUint64("uid")
privs := c.GetStringSlice("prm")
var scope func(db *gorm.DB) *gorm.DB
if slices.Contains(privs, "assistants.admin") {
scope = func(db *gorm.DB) *gorm.DB {
return db.Where("id = ?", assistantID)
}
} else if slices.Contains(privs, "assistants.delete") {
scope = func(db *gorm.DB) *gorm.DB {
return db.Where("assistants.id = ? AND assistants.flow_id = ?", assistantID, flowID).
Joins("INNER JOIN flows f ON f.id = assistants.flow_id").
Where("f.user_id = ?", uid)
}
} else {
logger.FromContext(c).Errorf("error filtering user role permissions: permission not found")
response.Error(c, response.ErrNotPermitted, nil)
return
}
if err = s.db.Model(&assistant).Scopes(scope).Take(&assistant).Error; err != nil {
logger.FromContext(c).WithError(err).Errorf("error getting assistant by id")
if gorm.IsRecordNotFoundError(err) {
response.Error(c, response.ErrAssistantsNotFound, err)
} else {
response.Error(c, response.ErrInternal, err)
}
return
}
fw, err := s.fc.GetFlow(c, int64(flowID))
if err != nil {
logger.FromContext(c).WithError(err).Errorf("error getting flow by id in flow controller")
response.Error(c, response.ErrInternal, err)
return
}
aw, err := fw.GetAssistant(c, int64(assistantID))
if err != nil {
logger.FromContext(c).WithError(err).Errorf("error getting assistant by id in flow controller")
response.Error(c, response.ErrInternal, err)
return
}
if err := aw.Finish(c); err != nil {
logger.FromContext(c).WithError(err).Errorf("error stopping assistant")
response.Error(c, response.ErrInternal, err)
return
}
if err = s.db.Scopes(scope).Delete(&assistant).Error; err != nil {
logger.FromContext(c).WithError(err).Errorf("error deleting assistant by id")
if gorm.IsRecordNotFoundError(err) {
response.Error(c, response.ErrAssistantsNotFound, err)
} else {
response.Error(c, response.ErrInternal, err)
}
return
}
assistantDB, err := convertAssistantToDatabase(assistant.Assistant)
if err != nil {
logger.FromContext(c).WithError(err).Errorf("error converting assistant to database")
response.Error(c, response.ErrInternal, err)
return
}
if s.ss != nil {
publisher := s.ss.NewFlowPublisher(int64(assistant.Flow.UserID), int64(assistant.FlowID))
publisher.AssistantDeleted(c, assistantDB)
}
response.Success(c, http.StatusOK, assistant)
}
func convertAssistantToDatabase(assistant models.Assistant) (database.Assistant, error) {
functions, err := json.Marshal(assistant.Functions)
if err != nil {
return database.Assistant{}, err
}
return database.Assistant{
ID: int64(assistant.ID),
Status: database.AssistantStatus(assistant.Status),
Title: assistant.Title,
Model: assistant.Model,
ModelProviderName: assistant.ModelProviderName,
Language: assistant.Language,
Functions: functions,
TraceID: database.PtrStringToNullString(assistant.TraceID),
FlowID: int64(assistant.FlowID),
UseAgents: assistant.UseAgents,
MsgchainID: database.Uint64ToNullInt64(assistant.MsgchainID),
CreatedAt: database.TimeToNullTime(assistant.CreatedAt),
UpdatedAt: database.TimeToNullTime(assistant.UpdatedAt),
DeletedAt: database.PtrTimeToNullTime(assistant.DeletedAt),
ModelProviderType: database.ProviderType(assistant.ModelProviderType),
ToolCallIDTemplate: assistant.ToolCallIDTemplate,
}, nil
}
// validateServiceResources fetches and validates user resource ownership for REST handlers.
// It mirrors the logic in validateUserResources from the graph package but works with *gorm.DB.
// privs must contain at least "resources.view" or "resources.admin"; otherwise permission is denied.
// "resources.admin" bypasses the user_id ownership check.
func validateServiceResources(db *gorm.DB, uid uint64, privs []string, ids []uint64) ([]database.UserResource, error) {
if len(ids) == 0 {
return nil, nil
}
isAdmin := slices.Contains(privs, "resources.admin")
if !isAdmin && !slices.Contains(privs, "resources.view") && len(ids) > 0 {
return nil, fmt.Errorf("permission 'resources.view' required to use resource IDs")
}
var recs []models.UserResource
if err := db.Model(&models.UserResource{}).Where("id IN (?)", ids).Find(&recs).Error; err != nil {
return nil, fmt.Errorf("failed to fetch resources: %w", err)
}
found := make(map[uint64]models.UserResource, len(recs))
for _, r := range recs {
found[r.ID] = r
}
result := make([]database.UserResource, 0, len(ids))
for _, id := range ids {
r, ok := found[id]
if !ok {
return nil, fmt.Errorf("resource %d not found", id)
}
if !isAdmin && r.UserID != uid {
return nil, fmt.Errorf("resource %d not accessible", id)
}
result = append(result, database.UserResource{
ID: int64(r.ID),
UserID: int64(r.UserID),
Hash: r.Hash,
Name: r.Name,
Path: r.Path,
Size: r.Size,
IsDir: r.IsDir,
CreatedAt: database.TimeToNullTime(r.CreatedAt),
UpdatedAt: database.TimeToNullTime(r.UpdatedAt),
})
}
return result, nil
}
+889
View File
@@ -0,0 +1,889 @@
package services
import (
"crypto/hmac"
"crypto/rand"
"crypto/sha256"
"encoding/base64"
"encoding/json"
"errors"
"fmt"
"io"
"net/http"
"net/url"
"path"
"slices"
"strconv"
"strings"
"time"
"pentagi/pkg/server/logger"
"pentagi/pkg/server/models"
"pentagi/pkg/server/oauth"
"pentagi/pkg/server/rdb"
"pentagi/pkg/server/response"
"pentagi/pkg/version"
"github.com/gin-contrib/sessions"
"github.com/gin-gonic/gin"
"github.com/jinzhu/gorm"
"github.com/sirupsen/logrus"
"golang.org/x/crypto/bcrypt"
"golang.org/x/oauth2"
)
const (
authStateCookieName = "state"
authNonceCookieName = "nonce"
authStateRequestTTL = 5 * time.Minute
)
type AuthServiceConfig struct {
BaseURL string
LoginCallbackURL string
SessionTimeout int // in seconds
}
type AuthService struct {
cfg AuthServiceConfig
db *gorm.DB
key []byte
oauth map[string]oauth.OAuthClient
}
func NewAuthService(
cfg AuthServiceConfig,
db *gorm.DB,
oauth map[string]oauth.OAuthClient,
) *AuthService {
var count int
err := db.Model(&models.User{}).Where("type = 'local'").Count(&count).Error
if err != nil {
logrus.WithError(err).Errorf("error getting local users count")
}
key, err := randBytes(32)
if err != nil {
logrus.WithError(err).Errorf("error generating key")
}
return &AuthService{
cfg: cfg,
db: db,
key: key,
oauth: oauth,
}
}
// AuthLogin is function to login user in the system
// @Summary Login user into system
// @Tags Public
// @Accept json
// @Produce json
// @Param json body models.Login true "Login form JSON data"
// @Success 200 {object} response.successResp "login successful"
// @Failure 400 {object} response.errorResp "invalid login data"
// @Failure 401 {object} response.errorResp "invalid login or password"
// @Failure 403 {object} response.errorResp "login not permitted"
// @Failure 500 {object} response.errorResp "internal error on login"
// @Router /auth/login [post]
func (s *AuthService) AuthLogin(c *gin.Context) {
var data models.Login
if err := c.ShouldBindJSON(&data); err != nil || data.Valid() != nil {
if err == nil {
err = data.Valid()
}
logger.FromContext(c).WithError(err).Errorf("error validating request data")
response.Error(c, response.ErrAuthInvalidLoginRequest, err)
return
}
var user models.UserPassword
if err := s.db.Take(&user, "mail = ? AND password IS NOT NULL", data.Mail).Error; err != nil {
logrus.WithError(err).Errorf("error getting user by mail '%s'", data.Mail)
response.Error(c, response.ErrAuthInvalidCredentials, err)
return
} else if err = user.Valid(); err != nil {
logger.FromContext(c).WithError(err).Errorf("error validating user data '%s'", user.Hash)
response.Error(c, response.ErrAuthInvalidUserData, err)
return
} else if user.RoleID == 100 {
logger.FromContext(c).WithError(err).Errorf("can't authorize external user '%s'", user.Hash)
response.Error(c, response.ErrAuthInvalidUserData, fmt.Errorf("user is external"))
return
}
if err := bcrypt.CompareHashAndPassword([]byte(user.Password), []byte(data.Password)); err != nil {
logger.FromContext(c).Errorf("error matching user input password")
response.Error(c, response.ErrAuthInvalidCredentials, err)
return
}
if user.Status != "active" {
logger.FromContext(c).Errorf("error checking active state for user '%s'", user.Status)
response.Error(c, response.ErrAuthInactiveUser, fmt.Errorf("user is inactive"))
return
}
var privs []string
err := s.db.Table("privileges").
Where("role_id = ?", user.RoleID).
Pluck("name", &privs).Error
if err != nil {
logger.FromContext(c).WithError(err).Errorf("error getting user privileges list '%s'", user.Hash)
response.Error(c, response.ErrAuthInvalidServiceData, err)
return
}
uuid, err := rdb.MakeUuidStrFromHash(user.Hash)
if err != nil {
logger.FromContext(c).WithError(err).Errorf("error validating user data '%s'", user.Hash)
response.Error(c, response.ErrAuthInvalidUserData, err)
return
}
expires := s.cfg.SessionTimeout
session := sessions.Default(c)
session.Set("uid", user.ID)
session.Set("uhash", user.Hash)
session.Set("rid", user.RoleID)
session.Set("tid", models.UserTypeLocal.String())
session.Set("prm", privs)
session.Set("gtm", time.Now().Unix())
session.Set("exp", time.Now().Add(time.Duration(expires)*time.Second).Unix())
session.Set("uuid", uuid)
session.Set("uname", user.Name)
session.Options(sessions.Options{
HttpOnly: true,
Secure: c.Request.TLS != nil,
SameSite: http.SameSiteLaxMode,
Path: s.cfg.BaseURL,
MaxAge: int(expires),
})
if err := session.Save(); err != nil {
logger.FromContext(c).WithError(err).Errorf("error saving session")
response.Error(c, response.ErrInternal, err)
return
}
logger.FromContext(c).
WithFields(logrus.Fields{
"age": expires,
"uid": user.ID,
"uhash": user.Hash,
"rid": user.RoleID,
"tid": session.Get("tid"),
"gtm": session.Get("gtm"),
"exp": session.Get("exp"),
"prm": session.Get("prm"),
}).
Infof("user made successful local login for '%s'", data.Mail)
response.Success(c, http.StatusOK, struct{}{})
}
func (s *AuthService) refreshCookie(c *gin.Context, resp *info, privs []string) error {
session := sessions.Default(c)
expires := int(s.cfg.SessionTimeout)
session.Set("prm", privs)
session.Set("gtm", time.Now().Unix())
session.Set("exp", time.Now().Add(time.Duration(expires)*time.Second).Unix())
resp.Privs = privs
session.Set("uid", resp.User.ID)
session.Set("uhash", resp.User.Hash)
session.Set("rid", resp.User.RoleID)
session.Set("tid", resp.User.Type.String())
session.Options(sessions.Options{
HttpOnly: true,
Secure: c.Request.TLS != nil,
SameSite: http.SameSiteLaxMode,
Path: s.cfg.BaseURL,
MaxAge: expires,
})
if err := session.Save(); err != nil {
logger.FromContext(c).WithError(err).Errorf("error saving session")
return err
}
logger.FromContext(c).
WithFields(logrus.Fields{
"age": expires,
"uid": resp.User.ID,
"uhash": resp.User.Hash,
"rid": resp.User.RoleID,
"tid": session.Get("tid"),
"gtm": session.Get("gtm"),
"exp": session.Get("exp"),
"prm": session.Get("prm"),
}).
Infof("session was refreshed for '%s' '%s'", resp.User.Mail, resp.User.Name)
return nil
}
// AuthAuthorize is function to login user in OAuth2 external system
// @Summary Login user into OAuth2 external system via HTTP redirect
// @Tags Public
// @Produce json
// @Param return_uri query string false "URI to redirect user there after login" default(/)
// @Param provider query string false "OAuth provider name (google, github, etc.)" default(google) enums:"google,github"
// @Success 307 "redirect to SSO login page"
// @Failure 400 {object} response.errorResp "invalid autorizarion query"
// @Failure 403 {object} response.errorResp "authorize not permitted"
// @Failure 500 {object} response.errorResp "internal error on autorizarion"
// @Router /auth/authorize [get]
func (s *AuthService) AuthAuthorize(c *gin.Context) {
stateData := map[string]string{
"exp": strconv.FormatInt(time.Now().Add(authStateRequestTTL).Unix(), 10),
}
queryReturnURI := c.Query("return_uri")
if queryReturnURI != "" {
returnURL, err := url.Parse(queryReturnURI)
if err != nil {
logger.FromContext(c).WithError(err).Errorf("failed to parse return url argument '%s'", queryReturnURI)
response.Error(c, response.ErrAuthInvalidAuthorizeQuery, err)
return
}
returnURL.Path = path.Clean(path.Join("/", returnURL.Path))
stateData["return_uri"] = returnURL.RequestURI()
}
provider := c.Query("provider")
oauthClient, ok := s.oauth[provider]
if !ok {
logger.FromContext(c).Errorf("external OAuth2 provider '%s' is not initialized", provider)
err := fmt.Errorf("provider not initialized")
response.Error(c, response.ErrNotPermitted, err)
return
}
stateData["provider"] = provider
stateUniq, err := randBase64String(16)
if err != nil {
logger.FromContext(c).WithError(err).Errorf("failed to generate state random data")
response.Error(c, response.ErrInternal, err)
return
}
stateData["uniq"] = stateUniq
nonce, err := randBase64String(16)
if err != nil {
logger.FromContext(c).WithError(err).Errorf("failed to generate nonce random data")
response.Error(c, response.ErrInternal, err)
return
}
stateJSON, err := json.Marshal(stateData)
if err != nil {
logger.FromContext(c).WithError(err).Errorf("failed to marshal state json data")
response.Error(c, response.ErrInternal, err)
return
}
mac := hmac.New(sha256.New, s.key)
mac.Write(stateJSON)
signature := mac.Sum(nil)
signedStateJSON := append(signature, stateJSON...)
state := base64.RawURLEncoding.EncodeToString(signedStateJSON)
// Google OAuth uses POST callback which requires SameSite=None for cross-site requests
// GitHub and other providers use GET callback which works with SameSite=Lax
sameSiteMode := http.SameSiteLaxMode
if provider == "google" {
sameSiteMode = http.SameSiteNoneMode
}
maxAge := int(authStateRequestTTL / time.Second)
s.setCallbackCookie(c.Writer, c.Request, authStateCookieName, state, maxAge, sameSiteMode)
s.setCallbackCookie(c.Writer, c.Request, authNonceCookieName, nonce, maxAge, sameSiteMode)
authOpts := []oauth2.AuthCodeOption{
oauth2.SetAuthURLParam("nonce", nonce),
oauth2.SetAuthURLParam("response_mode", "form_post"),
oauth2.SetAuthURLParam("response_type", "code id_token"),
}
http.Redirect(c.Writer, c.Request,
oauthClient.AuthCodeURL(state, authOpts...),
http.StatusTemporaryRedirect)
}
// AuthLoginGetCallback is function to catch login callback from OAuth application with code only
// @Summary Login user from external OAuth application
// @Tags Public
// @Accept json
// @Produce json
// @Param code query string false "Auth code from OAuth provider to exchange token"
// @Success 303 "redirect to registered return_uri path in the state"
// @Failure 400 {object} response.errorResp "invalid login data"
// @Failure 401 {object} response.errorResp "invalid login or password"
// @Failure 403 {object} response.errorResp "login not permitted"
// @Failure 500 {object} response.errorResp "internal error on login"
// @Router /auth/login-callback [get]
func (s *AuthService) AuthLoginGetCallback(c *gin.Context) {
code := c.Query("code")
if code == "" {
response.Error(c, response.ErrAuthInvalidLoginCallbackRequest, fmt.Errorf("code is required"))
return
}
state, err := c.Request.Cookie(authStateCookieName)
if err != nil {
logger.FromContext(c).WithError(err).Errorf("error getting state from cookie")
response.Error(c, response.ErrAuthInvalidAuthorizationState, err)
return
}
queryState := c.Query("state")
if queryState == "" {
logger.FromContext(c).Errorf("error missing state parameter in OAuth callback")
response.Error(c, response.ErrAuthInvalidAuthorizationState, fmt.Errorf("state parameter is required"))
return
}
if queryState != state.Value {
logger.FromContext(c).Errorf("error matching received state to stored one")
response.Error(c, response.ErrAuthInvalidAuthorizationState, nil)
return
}
stateData, err := s.parseState(c, state.Value)
if err != nil {
return
}
s.authLoginCallback(c, stateData, code)
}
// AuthLoginPostCallback is function to catch login callback from OAuth application
// @Summary Login user from external OAuth application
// @Tags Public
// @Accept json
// @Produce json
// @Param json body models.AuthCallback true "Auth form JSON data"
// @Success 303 "redirect to registered return_uri path in the state"
// @Failure 400 {object} response.errorResp "invalid login data"
// @Failure 401 {object} response.errorResp "invalid login or password"
// @Failure 403 {object} response.errorResp "login not permitted"
// @Failure 500 {object} response.errorResp "internal error on login"
// @Router /auth/login-callback [post]
func (s *AuthService) AuthLoginPostCallback(c *gin.Context) {
var (
data models.AuthCallback
err error
)
if err = c.ShouldBind(&data); err != nil || data.Valid() != nil {
if err == nil {
err = data.Valid()
}
logger.FromContext(c).WithError(err).Errorf("error validating request data")
response.Error(c, response.ErrAuthInvalidLoginCallbackRequest, err)
return
}
state, err := c.Request.Cookie(authStateCookieName)
if err != nil {
logger.FromContext(c).WithError(err).Errorf("error getting state from cookie")
response.Error(c, response.ErrAuthInvalidAuthorizationState, err)
return
}
if data.State != state.Value {
logger.FromContext(c).Errorf("error matching received state to stored one")
response.Error(c, response.ErrAuthInvalidAuthorizationState, nil)
return
}
stateData, err := s.parseState(c, state.Value)
if err != nil {
return
}
s.authLoginCallback(c, stateData, data.Code)
}
// AuthLogoutCallback is function to catch logout callback from OAuth application
// @Summary Logout current user from external OAuth application
// @Tags Public
// @Accept json
// @Success 303 {object} response.successResp "logout successful"
// @Router /auth/logout-callback [post]
func (s *AuthService) AuthLogoutCallback(c *gin.Context) {
s.resetSession(c)
http.Redirect(c.Writer, c.Request, "/", http.StatusSeeOther)
}
// AuthLogout is function to logout current user
// @Summary Logout current user via HTTP redirect
// @Tags Public
// @Produce json
// @Param return_uri query string false "URI to redirect user there after logout" default(/)
// @Success 307 "redirect to input return_uri path"
// @Router /auth/logout [get]
func (s *AuthService) AuthLogout(c *gin.Context) {
returnURI := "/"
if returnURL, err := url.Parse(c.Query("return_uri")); err == nil {
if uri := returnURL.RequestURI(); uri != "" {
returnURI = path.Clean(path.Join("/", uri))
}
}
session := sessions.Default(c)
logger.FromContext(c).
WithFields(logrus.Fields{
"uid": session.Get("uid"),
"uhash": session.Get("uhash"),
"rid": session.Get("rid"),
"tid": session.Get("tid"),
"gtm": session.Get("gtm"),
"exp": session.Get("exp"),
"prm": session.Get("prm"),
}).
Info("user made successful logout")
s.resetSession(c)
http.Redirect(c.Writer, c.Request, returnURI, http.StatusTemporaryRedirect)
}
func (s *AuthService) authLoginCallback(c *gin.Context, stateData map[string]string, code string) {
var (
privs []string
role models.Role
user models.User
)
provider := stateData["provider"]
oauthClient, ok := s.oauth[provider]
if !ok {
logger.FromContext(c).Errorf("external OAuth2 provider '%s' is not initialized", provider)
response.Error(c, response.ErrNotPermitted, fmt.Errorf("provider not initialized"))
return
}
ctx := c.Request.Context()
oauth2Token, err := oauthClient.Exchange(ctx, code)
if err != nil {
logger.FromContext(c).WithError(err).Errorf("failed to exchange token")
response.Error(c, response.ErrAuthExchangeTokenFail, err)
return
}
if !oauth2Token.Valid() {
logger.FromContext(c).Errorf("failed to validate OAuth2 token")
response.Error(c, response.ErrAuthVerificationTokenFail, nil)
return
}
nonce, err := c.Request.Cookie(authNonceCookieName)
if err != nil {
logger.FromContext(c).WithError(err).Errorf("error getting nonce from cookie")
response.Error(c, response.ErrAuthInvalidAuthorizationNonce, err)
return
}
email, err := oauthClient.ResolveEmail(ctx, nonce.Value, oauth2Token)
if err != nil {
logger.FromContext(c).WithError(err).Errorf("failed to resolve email")
response.Error(c, response.ErrAuthInvalidUserData, err)
return
}
if !strings.Contains(email, "@") {
logger.FromContext(c).Errorf("invalid email format '%s'", email)
response.Error(c, response.ErrAuthInvalidUserData, fmt.Errorf("invalid email format"))
return
}
username := strings.Split(email, "@")[0]
if username == "" {
logger.FromContext(c).Errorf("empty username from email '%s'", email)
response.Error(c, response.ErrAuthInvalidUserData, fmt.Errorf("empty username"))
return
}
err = s.db.Take(&role, "id = ?", models.RoleUser).Error
if err != nil {
logger.FromContext(c).WithError(err).Errorf("error getting user role '%d'", models.RoleUser)
response.Error(c, response.ErrAuthInvalidServiceData, err)
return
}
err = s.db.Table("privileges").
Where("role_id = ?", models.RoleUser).
Pluck("name", &privs).Error
if err != nil {
logger.FromContext(c).WithError(err).Errorf("error getting user privileges list '%s'", user.Hash)
response.Error(c, response.ErrAuthInvalidServiceData, err)
return
}
filterQuery := "mail = ? AND type = ?"
if err = s.db.Take(&user, filterQuery, email, models.UserTypeOAuth).Error; err != nil {
if errors.Is(err, gorm.ErrRecordNotFound) {
user = models.User{
Hash: rdb.MakeUserHash(email),
Mail: email,
Name: username,
RoleID: models.RoleUser,
Status: "active",
Type: models.UserTypeOAuth,
}
tx := s.db.Begin()
if tx.Error != nil {
logger.FromContext(c).WithError(tx.Error).Errorf("error starting transaction")
response.Error(c, response.ErrInternal, tx.Error)
return
}
if err = tx.Create(&user).Error; err != nil {
tx.Rollback()
logger.FromContext(c).WithError(err).Errorf("error creating user")
response.Error(c, response.ErrInternal, err)
return
}
preferences := models.NewUserPreferences(user.ID)
if err = tx.Create(preferences).Error; err != nil {
tx.Rollback()
logger.FromContext(c).WithError(err).Errorf("error creating user preferences")
response.Error(c, response.ErrInternal, err)
return
}
if err = tx.Commit().Error; err != nil {
logger.FromContext(c).WithError(err).Errorf("error committing transaction")
response.Error(c, response.ErrInternal, err)
return
}
} else {
logger.FromContext(c).WithError(err).Errorf("error searching user by email '%s'", email)
response.Error(c, response.ErrInternal, err)
return
}
} else if err = user.Valid(); err != nil {
logger.FromContext(c).WithError(err).Errorf("error validating user data '%s'", user.Hash)
response.Error(c, response.ErrAuthInvalidUserData, err)
return
}
if user.Status != "active" {
logger.FromContext(c).Errorf("error checking active state for user '%s'", user.Status)
response.Error(c, response.ErrAuthInactiveUser, fmt.Errorf("user is inactive"))
return
}
expires := s.cfg.SessionTimeout
gtm := time.Now().Unix()
exp := time.Now().Add(time.Duration(expires) * time.Second).Unix()
session := sessions.Default(c)
session.Set("uid", user.ID)
session.Set("uhash", user.Hash)
session.Set("rid", user.RoleID)
session.Set("tid", user.Type.String())
session.Set("prm", privs)
session.Set("gtm", gtm)
session.Set("exp", exp)
session.Set("uuid", user.Mail)
session.Set("uname", user.Name)
session.Options(sessions.Options{
HttpOnly: true,
Secure: c.Request.TLS != nil,
SameSite: http.SameSiteLaxMode,
Path: s.cfg.BaseURL,
MaxAge: expires,
})
if err := session.Save(); err != nil {
logger.FromContext(c).WithError(err).Errorf("error saving session")
response.Error(c, response.ErrInternal, err)
return
}
// delete temporary cookies
// Google OAuth uses POST callback which requires SameSite=None for cross-site requests
// GitHub and other providers use GET callback which works with SameSite=Lax
sameSiteMode := http.SameSiteLaxMode
if stateData["provider"] == "google" {
sameSiteMode = http.SameSiteNoneMode
}
s.setCallbackCookie(c.Writer, c.Request, authStateCookieName, "", 0, sameSiteMode)
s.setCallbackCookie(c.Writer, c.Request, authNonceCookieName, "", 0, sameSiteMode)
logger.FromContext(c).
WithFields(logrus.Fields{
"age": expires,
"uid": user.ID,
"uhash": user.Hash,
"rid": user.RoleID,
"tid": user.Type,
"gtm": session.Get("gtm"),
"exp": session.Get("exp"),
"prm": session.Get("prm"),
}).
Infof("user made successful SSO login for '%s' '%s'", user.Mail, user.Name)
if returnURI := stateData["return_uri"]; returnURI == "" {
response.Success(c, http.StatusOK, nil)
} else {
u, err := url.Parse(returnURI)
if err != nil {
response.Success(c, http.StatusOK, nil)
return
}
query := u.Query()
query.Add("status", "success")
u.RawQuery = query.Encode()
http.Redirect(c.Writer, c.Request, u.RequestURI(), http.StatusSeeOther)
}
}
func (s *AuthService) parseState(c *gin.Context, state string) (map[string]string, error) {
var stateData map[string]string
stateJSON, err := base64.RawURLEncoding.DecodeString(state)
if err != nil {
logger.FromContext(c).WithError(err).Errorf("error on getting state as a base64")
response.Error(c, response.ErrAuthInvalidAuthorizationState, err)
return nil, err
}
signatureLen := 32
if len(stateJSON) <= signatureLen {
logger.FromContext(c).Errorf("error on parsing state from json data")
err := fmt.Errorf("unexpected state length")
response.Error(c, response.ErrAuthInvalidAuthorizationState, err)
return nil, err
}
stateSignature := stateJSON[:signatureLen]
stateJSON = stateJSON[signatureLen:]
mac := hmac.New(sha256.New, s.key)
mac.Write(stateJSON)
signature := mac.Sum(nil)
if !hmac.Equal(stateSignature, signature) {
logger.FromContext(c).Errorf("error on matching signature")
err := fmt.Errorf("mismatch state signature")
response.Error(c, response.ErrAuthInvalidAuthorizationState, err)
return nil, err
}
if err := json.Unmarshal(stateJSON, &stateData); err != nil {
logger.FromContext(c).WithError(err).Errorf("error on parsing state from json data")
response.Error(c, response.ErrAuthInvalidAuthorizationState, err)
return nil, err
}
expStr, ok := stateData["exp"]
if !ok || expStr == "" {
err := fmt.Errorf("missing required field: exp")
logger.FromContext(c).WithError(err).Errorf("error on validating state data")
response.Error(c, response.ErrAuthInvalidAuthorizationState, err)
return nil, err
}
if _, ok := stateData["provider"]; !ok {
err := fmt.Errorf("missing required field: provider")
logger.FromContext(c).WithError(err).Errorf("error on validating state data")
response.Error(c, response.ErrAuthInvalidAuthorizationState, err)
return nil, err
}
exp, err := strconv.ParseInt(expStr, 10, 64)
if err != nil {
logger.FromContext(c).WithError(err).Errorf("error on parsing expiration time")
response.Error(c, response.ErrAuthInvalidAuthorizationState, err)
return nil, err
}
if time.Now().Unix() > exp {
logger.FromContext(c).Errorf("error on checking expiration time")
err := fmt.Errorf("state signature expired")
response.Error(c, response.ErrAuthTokenExpired, err)
return nil, err
}
return stateData, nil
}
func (s *AuthService) setCallbackCookie(
w http.ResponseWriter, r *http.Request,
name, value string, maxAge int,
sameSite http.SameSite,
) {
// Check both direct TLS and X-Forwarded-Proto header (for reverse proxy setups)
useTLS := r.TLS != nil || r.Header.Get("X-Forwarded-Proto") == "https"
c := &http.Cookie{
Name: name,
Value: value,
HttpOnly: true,
Secure: useTLS,
SameSite: sameSite,
Path: path.Join(s.cfg.BaseURL, s.cfg.LoginCallbackURL),
MaxAge: maxAge,
}
http.SetCookie(w, c)
}
func (s *AuthService) resetSession(c *gin.Context) {
now := time.Now().Add(-1 * time.Second)
session := sessions.Default(c)
session.Set("gtm", now.Unix())
session.Set("exp", now.Unix())
session.Options(sessions.Options{
HttpOnly: true,
Secure: c.Request.TLS != nil,
SameSite: http.SameSiteLaxMode,
Path: s.cfg.BaseURL,
MaxAge: -1,
})
if err := session.Save(); err != nil {
logger.FromContext(c).WithError(err).Errorf("error resetting session")
}
}
// randBase64String is function to generate random base64 with set length (bytes)
func randBase64String(nByte int) (string, error) {
b := make([]byte, nByte)
if _, err := io.ReadFull(rand.Reader, b); err != nil {
return "", err
}
return base64.RawURLEncoding.EncodeToString(b), nil
}
// randBytes is function to generate random bytes with set length (bytes)
func randBytes(nByte int) ([]byte, error) {
b := make([]byte, nByte)
if _, err := io.ReadFull(rand.Reader, b); err != nil {
return nil, err
}
return b, nil
}
type info struct {
Type string `json:"type"`
Develop bool `json:"develop"`
User models.User `json:"user"`
Role models.Role `json:"role"`
Providers []string `json:"providers"`
Privs []string `json:"privileges"`
OAuth bool `json:"oauth"`
IssuedAt time.Time `json:"issued_at"`
ExpiresAt time.Time `json:"expires_at"`
}
// Info is function to return settings and current information about system and config
// @Summary Retrieve current user and system settings
// @Tags Public
// @Produce json
// @Security BearerAuth
// @Param refresh_cookie query boolean false "boolean arg to refresh current cookie, use explicit false"
// @Success 200 {object} response.successResp{data=info} "info received successful"
// @Failure 403 {object} response.errorResp "getting info not permitted"
// @Failure 404 {object} response.errorResp "user not found"
// @Failure 500 {object} response.errorResp "internal error on getting information about system and config"
// @Router /info [get]
func (s *AuthService) Info(c *gin.Context) {
var resp info
logger.FromContext(c).WithFields(logrus.Fields(c.Keys)).Trace("AuthService.Info")
now := time.Now().Unix()
uhash := c.GetString("uhash")
uid := c.GetUint64("uid")
tid := c.GetString("tid")
exp := c.GetInt64("exp")
gtm := c.GetInt64("gtm")
cpt := c.GetString("cpt")
privs := c.GetStringSlice("prm")
resp.Privs = privs
resp.IssuedAt = time.Unix(gtm, 0).UTC()
resp.ExpiresAt = time.Unix(exp, 0).UTC()
resp.Develop = version.IsDevelopMode()
resp.OAuth = tid == models.UserTypeOAuth.String()
for name := range s.oauth {
resp.Providers = append(resp.Providers, name)
}
logger.FromContext(c).WithFields(logrus.Fields(
map[string]any{
"exp": exp,
"gtm": gtm,
"uhash": uhash,
"now": now,
"cpt": cpt,
"uid": uid,
"tid": tid,
},
)).Trace("AuthService.Info")
if uhash == "" || exp == 0 || gtm == 0 || now > exp {
resp.Type = "guest"
resp.Privs = []string{}
response.Success(c, http.StatusOK, resp)
return
}
err := s.db.Take(&resp.User, "id = ?", uid).Related(&resp.Role).Error
if err != nil {
response.Error(c, response.ErrInfoUserNotFound, err)
return
} else if err = resp.User.Valid(); err != nil {
logger.FromContext(c).WithError(err).Errorf("error validating user data '%s'", resp.User.Hash)
response.Error(c, response.ErrInfoInvalidUserData, err)
return
}
if err = s.db.Table("privileges").Where("role_id = ?", resp.User.RoleID).Pluck("name", &privs).Error; err != nil {
logger.FromContext(c).WithError(err).Errorf("error getting user privileges list '%s'", resp.User.Hash)
response.Error(c, response.ErrInfoInvalidUserData, err)
return
}
if cpt == "automation" {
resp.Type = models.UserTypeAPI.String()
// filter out privileges that are not supported for API tokens
privs = slices.DeleteFunc(privs, func(priv string) bool {
return strings.HasPrefix(priv, "users.") ||
strings.HasPrefix(priv, "roles.") ||
strings.HasPrefix(priv, "settings.user.") ||
strings.HasPrefix(priv, "settings.tokens.")
})
resp.Privs = privs
response.Success(c, http.StatusOK, resp)
return
}
resp.Type = "user"
// check 5 minutes timeout to refresh current token
var fiveMins int64 = 5 * 60
if now >= gtm+fiveMins && c.Query("refresh_cookie") != "false" {
if err = s.refreshCookie(c, &resp, privs); err != nil {
logger.FromContext(c).WithError(err).Errorf("failed to refresh token")
// raise error when there is elapsing last five minutes
if now >= gtm+int64(s.cfg.SessionTimeout)-fiveMins {
response.Error(c, response.ErrAuthRequired, err)
return
}
}
}
// raise error when user has no permissions in the session auth cookie
if resp.Type != "guest" && resp.Privs == nil {
logger.FromContext(c).
WithFields(logrus.Fields{
"uid": resp.User.ID,
"rid": resp.User.RoleID,
"tid": resp.User.Type,
}).
Errorf("failed to get user privileges for '%s' '%s'", resp.User.Mail, resp.User.Name)
response.Error(c, response.ErrInternal, err)
return
}
response.Success(c, http.StatusOK, resp)
}
+287
View File
@@ -0,0 +1,287 @@
package services
import (
"errors"
"net/http"
"slices"
"strconv"
"pentagi/pkg/server/logger"
"pentagi/pkg/server/models"
"pentagi/pkg/server/rdb"
"pentagi/pkg/server/response"
"github.com/gin-gonic/gin"
"github.com/jinzhu/gorm"
)
type containers struct {
Containers []models.Container `json:"containers"`
Total uint64 `json:"total"`
}
type containersGrouped struct {
Grouped []string `json:"grouped"`
Total uint64 `json:"total"`
}
var containersSQLMappers = map[string]any{
"id": "{{table}}.id",
"type": "{{table}}.type",
"name": "{{table}}.name",
"image": "{{table}}.image",
"status": "{{table}}.status",
"local_id": "{{table}}.local_id",
"local_dir": "{{table}}.local_dir",
"flow_id": "{{table}}.flow_id",
"created_at": "{{table}}.created_at",
"updated_at": "{{table}}.updated_at",
"data": "({{table}}.type || ' ' || {{table}}.name || ' ' || {{table}}.status || ' ' || {{table}}.local_id || ' ' || {{table}}.local_dir)",
}
type ContainerService struct {
db *gorm.DB
}
func NewContainerService(db *gorm.DB) *ContainerService {
return &ContainerService{
db: db,
}
}
// GetContainers is a function to return containers list
// @Summary Retrieve containers list
// @Tags Containers
// @Produce json
// @Security BearerAuth
// @Param request query rdb.TableQuery true "query table params"
// @Success 200 {object} response.successResp{data=containers} "containers list received successful"
// @Failure 400 {object} response.errorResp "invalid query request data"
// @Failure 403 {object} response.errorResp "getting containers not permitted"
// @Failure 500 {object} response.errorResp "internal error on getting containers"
// @Router /containers/ [get]
func (s *ContainerService) GetContainers(c *gin.Context) {
var (
err error
query rdb.TableQuery
resp containers
)
if err = c.ShouldBindQuery(&query); err != nil {
logger.FromContext(c).WithError(err).Errorf("error binding query")
response.Error(c, response.ErrContainersInvalidRequest, err)
return
}
uid := c.GetUint64("uid")
privs := c.GetStringSlice("prm")
var scope func(db *gorm.DB) *gorm.DB
if slices.Contains(privs, "containers.admin") {
scope = func(db *gorm.DB) *gorm.DB {
return db.
Joins("INNER JOIN flows f ON f.id = containers.flow_id")
}
} else if slices.Contains(privs, "containers.view") {
scope = func(db *gorm.DB) *gorm.DB {
return db.
Joins("INNER JOIN flows f ON f.id = containers.flow_id").
Where("f.user_id = ?", uid)
}
} else {
logger.FromContext(c).Errorf("error filtering user role permissions: permission not found")
response.Error(c, response.ErrNotPermitted, nil)
return
}
query.Init("containers", containersSQLMappers)
if query.Group != "" {
if _, ok := containersSQLMappers[query.Group]; !ok {
logger.FromContext(c).Errorf("error finding containers grouped: group field not found")
response.Error(c, response.ErrContainersInvalidRequest, errors.New("group field not found"))
return
}
var respGrouped containersGrouped
if respGrouped.Total, err = query.QueryGrouped(s.db, &respGrouped.Grouped, scope); err != nil {
logger.FromContext(c).WithError(err).Errorf("error finding containers grouped")
response.Error(c, response.ErrInternal, err)
return
}
response.Success(c, http.StatusOK, respGrouped)
return
}
if resp.Total, err = query.Query(s.db, &resp.Containers, scope); err != nil {
logger.FromContext(c).WithError(err).Errorf("error finding containers")
response.Error(c, response.ErrInternal, err)
return
}
for i := 0; i < len(resp.Containers); i++ {
if err = resp.Containers[i].Valid(); err != nil {
logger.FromContext(c).WithError(err).Errorf("error validating container data '%d'", resp.Containers[i].ID)
response.Error(c, response.ErrContainersInvalidData, err)
return
}
}
response.Success(c, http.StatusOK, resp)
}
// GetFlowContainers is a function to return containers list by flow id
// @Summary Retrieve containers list by flow id
// @Tags Containers
// @Produce json
// @Security BearerAuth
// @Param flowID path int true "flow id" minimum(0)
// @Param request query rdb.TableQuery true "query table params"
// @Success 200 {object} response.successResp{data=containers} "containers list received successful"
// @Failure 400 {object} response.errorResp "invalid query request data"
// @Failure 403 {object} response.errorResp "getting containers not permitted"
// @Failure 500 {object} response.errorResp "internal error on getting containers"
// @Router /flows/{flowID}/containers/ [get]
func (s *ContainerService) GetFlowContainers(c *gin.Context) {
var (
err error
flowID uint64
query rdb.TableQuery
resp containers
)
if flowID, err = strconv.ParseUint(c.Param("flowID"), 10, 64); err != nil {
logger.FromContext(c).WithError(err).Errorf("error parsing flow id")
response.Error(c, response.ErrContainersInvalidRequest, err)
return
}
if err = c.ShouldBindQuery(&query); err != nil {
logger.FromContext(c).WithError(err).Errorf("error binding query")
response.Error(c, response.ErrContainersInvalidRequest, err)
return
}
uid := c.GetUint64("uid")
privs := c.GetStringSlice("prm")
var scope func(db *gorm.DB) *gorm.DB
if slices.Contains(privs, "containers.admin") {
scope = func(db *gorm.DB) *gorm.DB {
return db.
Joins("INNER JOIN flows f ON f.id = containers.flow_id").
Where("f.id = ?", flowID)
}
} else if slices.Contains(privs, "containers.view") {
scope = func(db *gorm.DB) *gorm.DB {
return db.
Joins("INNER JOIN flows f ON f.id = containers.flow_id").
Where("f.id = ? AND f.user_id = ?", flowID, uid)
}
} else {
logger.FromContext(c).Errorf("error filtering user role permissions: permission not found")
response.Error(c, response.ErrNotPermitted, nil)
return
}
query.Init("containers", containersSQLMappers)
if query.Group != "" {
if _, ok := containersSQLMappers[query.Group]; !ok {
logger.FromContext(c).Errorf("error finding containers grouped: group field not found")
response.Error(c, response.ErrContainersInvalidRequest, errors.New("group field not found"))
return
}
var respGrouped containersGrouped
if respGrouped.Total, err = query.QueryGrouped(s.db, &respGrouped.Grouped, scope); err != nil {
logger.FromContext(c).WithError(err).Errorf("error finding containers grouped")
response.Error(c, response.ErrInternal, err)
return
}
response.Success(c, http.StatusOK, respGrouped)
return
}
if resp.Total, err = query.Query(s.db, &resp.Containers, scope); err != nil {
logger.FromContext(c).WithError(err).Errorf("error finding containers")
response.Error(c, response.ErrInternal, err)
return
}
for i := 0; i < len(resp.Containers); i++ {
if err = resp.Containers[i].Valid(); err != nil {
logger.FromContext(c).WithError(err).Errorf("error validating container data '%d'", resp.Containers[i].ID)
response.Error(c, response.ErrContainersInvalidData, err)
return
}
}
response.Success(c, http.StatusOK, resp)
}
// GetFlowContainer is a function to return container info by id and flow id
// @Summary Retrieve container info by id and flow id
// @Tags Containers
// @Produce json
// @Security BearerAuth
// @Param flowID path int true "flow id" minimum(0)
// @Param containerID path int true "container id" minimum(0)
// @Success 200 {object} response.successResp{data=models.Container} "container info received successful"
// @Failure 403 {object} response.errorResp "getting container not permitted"
// @Failure 404 {object} response.errorResp "container not found"
// @Failure 500 {object} response.errorResp "internal error on getting container"
// @Router /flows/{flowID}/containers/{containerID} [get]
func (s *ContainerService) GetFlowContainer(c *gin.Context) {
var (
err error
containerID uint64
flowID uint64
resp models.Container
)
if flowID, err = strconv.ParseUint(c.Param("flowID"), 10, 64); err != nil {
logger.FromContext(c).WithError(err).Errorf("error parsing flow id")
response.Error(c, response.ErrContainersInvalidRequest, err)
return
}
if containerID, err = strconv.ParseUint(c.Param("containerID"), 10, 64); err != nil {
logger.FromContext(c).WithError(err).Errorf("error parsing container id")
response.Error(c, response.ErrContainersInvalidRequest, err)
return
}
uid := c.GetUint64("uid")
privs := c.GetStringSlice("prm")
var scope func(db *gorm.DB) *gorm.DB
if slices.Contains(privs, "containers.admin") {
scope = func(db *gorm.DB) *gorm.DB {
return db.Where("f.id = ?", flowID)
}
} else if slices.Contains(privs, "containers.view") {
scope = func(db *gorm.DB) *gorm.DB {
return db.Where("f.id = ? AND f.user_id = ?", flowID, uid)
}
} else {
logger.FromContext(c).Errorf("error filtering user role permissions: permission not found")
response.Error(c, response.ErrNotPermitted, nil)
return
}
err = s.db.Model(&resp).
Joins("INNER JOIN flows f ON f.id = flow_id").
Scopes(scope).
Where("containers.id = ?", containerID).
Take(&resp).Error
if err != nil {
logger.FromContext(c).WithError(err).Errorf("error on getting container by id")
if gorm.IsRecordNotFoundError(err) {
response.Error(c, response.ErrContainersNotFound, err)
} else {
response.Error(c, response.ErrInternal, err)
}
return
}
response.Success(c, http.StatusOK, resp)
}
File diff suppressed because it is too large Load Diff
File diff suppressed because it is too large Load Diff
+675
View File
@@ -0,0 +1,675 @@
package services
import (
"encoding/json"
"errors"
"net/http"
"slices"
"strconv"
"pentagi/pkg/controller"
"pentagi/pkg/database"
"pentagi/pkg/graph/subscriptions"
"pentagi/pkg/providers"
"pentagi/pkg/providers/provider"
"pentagi/pkg/server/logger"
"pentagi/pkg/server/models"
"pentagi/pkg/server/rdb"
"pentagi/pkg/server/response"
"github.com/gin-gonic/gin"
"github.com/jinzhu/gorm"
)
type flows struct {
Flows []models.Flow `json:"flows"`
Total uint64 `json:"total"`
}
type flowsGrouped struct {
Grouped []string `json:"grouped"`
Total uint64 `json:"total"`
}
var flowsSQLMappers = map[string]any{
"id": "{{table}}.id",
"status": "{{table}}.status",
"title": "{{table}}.title",
"model": "{{table}}.model",
"model_provider_name": "{{table}}.model_provider_name",
"model_provider_type": "{{table}}.model_provider_type",
"language": "{{table}}.language",
"created_at": "{{table}}.created_at",
"updated_at": "{{table}}.updated_at",
"data": "({{table}}.status || ' ' || {{table}}.title || ' ' || {{table}}.model || ' ' || {{table}}.model_provider || ' ' || {{table}}.language)",
}
type FlowService struct {
db *gorm.DB
pc providers.ProviderController
fc controller.FlowController
ss subscriptions.SubscriptionsController
}
func NewFlowService(
db *gorm.DB,
pc providers.ProviderController,
fc controller.FlowController,
ss subscriptions.SubscriptionsController,
) *FlowService {
return &FlowService{
db: db,
pc: pc,
fc: fc,
ss: ss,
}
}
// GetFlows is a function to return flows list
// @Summary Retrieve flows list
// @Tags Flows
// @Produce json
// @Security BearerAuth
// @Param request query rdb.TableQuery true "query table params"
// @Success 200 {object} response.successResp{data=flows} "flows list received successful"
// @Failure 400 {object} response.errorResp "invalid query request data"
// @Failure 403 {object} response.errorResp "getting flows not permitted"
// @Failure 500 {object} response.errorResp "internal error on getting flows"
// @Router /flows/ [get]
func (s *FlowService) GetFlows(c *gin.Context) {
var (
err error
query rdb.TableQuery
resp flows
)
if err = c.ShouldBindQuery(&query); err != nil {
logger.FromContext(c).WithError(err).Errorf("error binding query")
response.Error(c, response.ErrFlowsInvalidRequest, err)
return
}
uid := c.GetUint64("uid")
privs := c.GetStringSlice("prm")
var scope func(db *gorm.DB) *gorm.DB
if slices.Contains(privs, "flows.admin") {
scope = func(db *gorm.DB) *gorm.DB {
return db
}
} else if slices.Contains(privs, "flows.view") {
scope = func(db *gorm.DB) *gorm.DB {
return db.Where("user_id = ?", uid)
}
} else {
logger.FromContext(c).Errorf("error filtering user role permissions: permission not found")
response.Error(c, response.ErrNotPermitted, nil)
return
}
query.Init("flows", flowsSQLMappers)
if query.Group != "" {
if _, ok := flowsSQLMappers[query.Group]; !ok {
logger.FromContext(c).Errorf("error finding flows grouped: group field not found")
response.Error(c, response.ErrFlowsInvalidRequest, errors.New("group field not found"))
return
}
var respGrouped flowsGrouped
if respGrouped.Total, err = query.QueryGrouped(s.db, &respGrouped.Grouped, scope); err != nil {
logger.FromContext(c).WithError(err).Errorf("error finding flows grouped")
response.Error(c, response.ErrInternal, err)
return
}
response.Success(c, http.StatusOK, respGrouped)
return
}
if resp.Total, err = query.Query(s.db, &resp.Flows, scope); err != nil {
logger.FromContext(c).WithError(err).Errorf("error finding flows")
response.Error(c, response.ErrInternal, err)
return
}
for i := 0; i < len(resp.Flows); i++ {
if err = resp.Flows[i].Valid(); err != nil {
logger.FromContext(c).WithError(err).Errorf("error validating flow data '%d'", resp.Flows[i].ID)
response.Error(c, response.ErrFlowsInvalidData, err)
return
}
}
response.Success(c, http.StatusOK, resp)
}
// GetFlow is a function to return flow by id
// @Summary Retrieve flow by id
// @Tags Flows
// @Produce json
// @Security BearerAuth
// @Param flowID path int true "flow id" minimum(0)
// @Success 200 {object} response.successResp{data=models.Flow} "flow received successful"
// @Failure 403 {object} response.errorResp "getting flow not permitted"
// @Failure 404 {object} response.errorResp "flow not found"
// @Failure 500 {object} response.errorResp "internal error on getting flow"
// @Router /flows/{flowID} [get]
func (s *FlowService) GetFlow(c *gin.Context) {
var (
err error
flowID uint64
resp models.Flow
)
if flowID, err = strconv.ParseUint(c.Param("flowID"), 10, 64); err != nil {
logger.FromContext(c).WithError(err).Errorf("error parsing flow id")
response.Error(c, response.ErrFlowsInvalidRequest, err)
return
}
uid := c.GetUint64("uid")
privs := c.GetStringSlice("prm")
var scope func(db *gorm.DB) *gorm.DB
if slices.Contains(privs, "flows.admin") {
scope = func(db *gorm.DB) *gorm.DB {
return db.Where("id = ?", flowID)
}
} else if slices.Contains(privs, "flows.view") {
scope = func(db *gorm.DB) *gorm.DB {
return db.Where("id = ? AND user_id = ?", flowID, uid)
}
} else {
logger.FromContext(c).Errorf("error filtering user role permissions: permission not found")
response.Error(c, response.ErrNotPermitted, nil)
return
}
if err = s.db.Model(&resp).Scopes(scope).Take(&resp).Error; err != nil {
logger.FromContext(c).WithError(err).Errorf("error on getting flow by id")
if gorm.IsRecordNotFoundError(err) {
response.Error(c, response.ErrFlowsNotFound, err)
} else {
response.Error(c, response.ErrInternal, err)
}
return
}
response.Success(c, http.StatusOK, resp)
}
// GetFlowGraph is a function to return flow graph by id
// @Summary Retrieve flow graph by id
// @Tags Flows
// @Produce json
// @Security BearerAuth
// @Param flowID path int true "flow id" minimum(0)
// @Success 200 {object} response.successResp{data=models.FlowTasksSubtasks} "flow graph received successful"
// @Failure 403 {object} response.errorResp "getting flow graph not permitted"
// @Failure 404 {object} response.errorResp "flow graph not found"
// @Failure 500 {object} response.errorResp "internal error on getting flow graph"
// @Router /flows/{flowID}/graph [get]
func (s *FlowService) GetFlowGraph(c *gin.Context) {
var (
err error
flowID uint64
resp models.FlowTasksSubtasks
tids []uint64
)
if flowID, err = strconv.ParseUint(c.Param("flowID"), 10, 64); err != nil {
logger.FromContext(c).WithError(err).Errorf("error parsing flow id")
response.Error(c, response.ErrFlowsInvalidRequest, err)
return
}
uid := c.GetUint64("uid")
privs := c.GetStringSlice("prm")
var scope func(db *gorm.DB) *gorm.DB
if slices.Contains(privs, "flows.admin") {
scope = func(db *gorm.DB) *gorm.DB {
return db.Where("id = ?", flowID)
}
} else if slices.Contains(privs, "flows.view") {
scope = func(db *gorm.DB) *gorm.DB {
return db.Where("id = ? AND user_id = ?", flowID, uid)
}
} else {
logger.FromContext(c).Errorf("error filtering user role permissions: permission not found")
response.Error(c, response.ErrNotPermitted, nil)
return
}
err = s.db.Model(&resp).
Scopes(scope).
Take(&resp).Error
if err != nil {
logger.FromContext(c).WithError(err).Errorf("error on getting flow by id")
if gorm.IsRecordNotFoundError(err) {
response.Error(c, response.ErrFlowsNotFound, err)
} else {
response.Error(c, response.ErrInternal, err)
}
return
}
isTasksAdmin := slices.Contains(privs, "tasks.admin")
isTasksView := slices.Contains(privs, "tasks.view")
if !(resp.UserID == uid && isTasksView) && !(resp.UserID != uid && isTasksAdmin) {
response.Success(c, http.StatusOK, resp)
return
}
if resp.UserID != uid && !slices.Contains(privs, "tasks.admin") {
response.Success(c, http.StatusOK, resp)
return
}
err = s.db.Model(&resp).Association("tasks").Find(&resp.Tasks).Error
if err != nil {
logger.FromContext(c).WithError(err).Errorf("error on getting flow tasks")
response.Error(c, response.ErrInternal, err)
return
}
isSubtasksAdmin := slices.Contains(privs, "subtasks.admin")
isSubtasksView := slices.Contains(privs, "subtasks.view")
if !(resp.UserID == uid && isSubtasksView) && !(resp.UserID != uid && isSubtasksAdmin) {
response.Success(c, http.StatusOK, resp)
return
}
for _, task := range resp.Tasks {
tids = append(tids, task.ID)
}
var subtasks []models.Subtask
err = s.db.Model(&subtasks).Where("task_id IN (?)", tids).Find(&subtasks).Error
if err != nil {
logger.FromContext(c).WithError(err).Errorf("error on getting flow subtasks")
response.Error(c, response.ErrInternal, err)
return
}
tasksSubtasks := map[uint64][]models.Subtask{}
for _, subtask := range subtasks {
tasksSubtasks[subtask.TaskID] = append(tasksSubtasks[subtask.TaskID], subtask)
}
for i := range resp.Tasks {
resp.Tasks[i].Subtasks = tasksSubtasks[resp.Tasks[i].ID]
}
if err = resp.Valid(); err != nil {
logger.FromContext(c).WithError(err).Errorf("error validating flow data '%d'", flowID)
response.Error(c, response.ErrFlowsInvalidData, err)
return
}
response.Success(c, http.StatusOK, resp)
}
// CreateFlow is a function to create new flow with custom functions
// @Summary Create new flow with custom functions
// @Tags Flows
// @Accept json
// @Produce json
// @Security BearerAuth
// @Param json body models.CreateFlow true "flow model to create"
// @Success 201 {object} response.successResp{data=models.Flow} "flow created successful"
// @Failure 400 {object} response.errorResp "invalid flow request data"
// @Failure 403 {object} response.errorResp "creating flow not permitted"
// @Failure 500 {object} response.errorResp "internal error on creating flow"
// @Router /flows/ [post]
func (s *FlowService) CreateFlow(c *gin.Context) {
var (
err error
flow models.Flow
createFlow models.CreateFlow
)
if err := c.ShouldBindJSON(&createFlow); err != nil {
logger.FromContext(c).WithError(err).Errorf("error binding JSON")
response.Error(c, response.ErrFlowsInvalidRequest, err)
return
}
privs := c.GetStringSlice("prm")
if !slices.Contains(privs, "flows.create") {
logger.FromContext(c).Errorf("error filtering user role permissions: permission not found")
response.Error(c, response.ErrNotPermitted, nil)
return
}
if err := createFlow.Valid(); err != nil {
logger.FromContext(c).WithError(err).Errorf("error validating flow data")
response.Error(c, response.ErrFlowsInvalidData, err)
return
}
uid := c.GetUint64("uid")
prvname := provider.ProviderName(createFlow.Provider)
prv, err := s.pc.GetProvider(c, prvname, int64(uid))
if err != nil {
logger.FromContext(c).WithError(err).Errorf("error getting provider: not found")
response.Error(c, response.ErrInternal, err)
return
}
prvtype := prv.Type()
dbResources, err := validateServiceResources(s.db, uid, privs, createFlow.ResourceIDs)
if err != nil {
logger.FromContext(c).WithError(err).Errorf("error validating resource ids")
response.Error(c, response.ErrFlowsInvalidRequest, err)
return
}
fw, err := s.fc.CreateFlow(c, int64(uid), createFlow.Input, prvname, prvtype, createFlow.Functions, dbResources)
if err != nil {
logger.FromContext(c).WithError(err).Errorf("error creating flow")
response.Error(c, response.ErrInternal, err)
return
}
err = s.db.Model(&flow).Where("id = ?", fw.GetFlowID()).Take(&flow).Error
if err != nil {
logger.FromContext(c).WithError(err).Errorf("error getting flow by id")
response.Error(c, response.ErrInternal, err)
return
}
response.Success(c, http.StatusCreated, flow)
}
// PatchFlow is a function to patch flow
// @Summary Patch flow
// @Tags Flows
// @Accept json
// @Produce json
// @Security BearerAuth
// @Param flowID path int true "flow id" minimum(0)
// @Param json body models.PatchFlow true "flow model to patch"
// @Success 200 {object} response.successResp{data=models.Flow} "flow patched successful"
// @Failure 400 {object} response.errorResp "invalid flow request data"
// @Failure 403 {object} response.errorResp "patching flow not permitted"
// @Failure 500 {object} response.errorResp "internal error on patching flow"
// @Router /flows/{flowID} [put]
func (s *FlowService) PatchFlow(c *gin.Context) {
var (
err error
flow models.Flow
flowID uint64
patchFlow models.PatchFlow
)
if err := c.ShouldBindJSON(&patchFlow); err != nil {
logger.FromContext(c).WithError(err).Errorf("error binding JSON")
response.Error(c, response.ErrFlowsInvalidRequest, err)
return
}
if err := patchFlow.Valid(); err != nil {
logger.FromContext(c).WithError(err).Errorf("error validating flow data")
response.Error(c, response.ErrFlowsInvalidData, err)
return
}
flowID, err = strconv.ParseUint(c.Param("flowID"), 10, 64)
if err != nil {
logger.FromContext(c).WithError(err).Errorf("error parsing flow id")
response.Error(c, response.ErrFlowsInvalidRequest, err)
return
}
uid := c.GetUint64("uid")
privs := c.GetStringSlice("prm")
var scope func(db *gorm.DB) *gorm.DB
if slices.Contains(privs, "flows.admin") {
scope = func(db *gorm.DB) *gorm.DB {
return db.Where("id = ?", flowID)
}
} else if slices.Contains(privs, "flows.edit") {
scope = func(db *gorm.DB) *gorm.DB {
return db.Where("id = ? AND user_id = ?", flowID, uid)
}
} else {
logger.FromContext(c).Errorf("error filtering user role permissions: permission not found")
response.Error(c, response.ErrNotPermitted, nil)
return
}
if err = s.db.Model(&flow).Scopes(scope).Take(&flow).Error; err != nil {
logger.FromContext(c).WithError(err).Errorf("error getting flow by id")
if gorm.IsRecordNotFoundError(err) {
response.Error(c, response.ErrFlowsNotFound, err)
} else {
response.Error(c, response.ErrInternal, err)
}
return
}
fw, err := s.fc.GetFlow(c, int64(flow.ID))
if err != nil {
logger.FromContext(c).WithError(err).Errorf("error getting flow by id in flow controller")
response.Error(c, response.ErrInternal, err)
return
}
switch patchFlow.Action {
case "stop":
if err := fw.Stop(c); err != nil {
logger.FromContext(c).WithError(err).Errorf("error stopping flow")
response.Error(c, response.ErrInternal, err)
return
}
case "finish":
if err := fw.Finish(c); err != nil {
logger.FromContext(c).WithError(err).Errorf("error finishing flow")
response.Error(c, response.ErrInternal, err)
return
}
case "input":
if patchFlow.Input == nil || *patchFlow.Input == "" {
logger.FromContext(c).Errorf("error sending input to flow: input is empty")
response.Error(c, response.ErrFlowsInvalidRequest, nil)
return
}
var prv provider.Provider
if patchFlow.Provider != nil && *patchFlow.Provider != "" {
prv, err = s.pc.GetProvider(c, provider.ProviderName(*patchFlow.Provider), int64(uid))
if err != nil {
logger.FromContext(c).WithError(err).Errorf("error getting provider by name")
response.Error(c, response.ErrInternal, err)
return
}
}
dbResources, err := validateServiceResources(s.db, uid, privs, patchFlow.ResourceIDs)
if err != nil {
logger.FromContext(c).WithError(err).Errorf("error validating resource ids")
response.Error(c, response.ErrFlowsInvalidRequest, err)
return
}
if err := fw.PutInput(c, *patchFlow.Input, prv, dbResources); err != nil {
logger.FromContext(c).WithError(err).Errorf("error sending input to flow")
response.Error(c, response.ErrInternal, err)
return
}
case "rename":
if patchFlow.Name == nil || *patchFlow.Name == "" {
logger.FromContext(c).Errorf("error renaming flow: name is empty")
response.Error(c, response.ErrFlowsInvalidRequest, nil)
return
}
if err := fw.Rename(c, *patchFlow.Name); err != nil {
logger.FromContext(c).WithError(err).Errorf("error renaming flow")
response.Error(c, response.ErrInternal, err)
return
}
default:
logger.FromContext(c).Errorf("error filtering flow action")
response.Error(c, response.ErrFlowsInvalidRequest, nil)
return
}
if err = s.db.Model(&flow).Scopes(scope).Take(&flow).Error; err != nil {
logger.FromContext(c).WithError(err).Errorf("error getting flow by id")
if gorm.IsRecordNotFoundError(err) {
response.Error(c, response.ErrFlowsNotFound, err)
} else {
response.Error(c, response.ErrInternal, err)
}
return
}
response.Success(c, http.StatusOK, flow)
}
// DeleteFlow is a function to delete flow by id
// @Summary Delete flow by id
// @Tags Flows
// @Security BearerAuth
// @Param flowID path int true "flow id" minimum(0)
// @Success 200 {object} response.successResp{data=models.Flow} "flow deleted successful"
// @Failure 403 {object} response.errorResp "deleting flow not permitted"
// @Failure 404 {object} response.errorResp "flow not found"
// @Failure 500 {object} response.errorResp "internal error on deleting flow"
// @Router /flows/{flowID} [delete]
func (s *FlowService) DeleteFlow(c *gin.Context) {
var (
err error
flow models.Flow
flowID uint64
)
flowID, err = strconv.ParseUint(c.Param("flowID"), 10, 64)
if err != nil {
logger.FromContext(c).WithError(err).Errorf("error parsing flow id")
response.Error(c, response.ErrFlowsInvalidRequest, err)
return
}
uid := c.GetUint64("uid")
privs := c.GetStringSlice("prm")
var scope func(db *gorm.DB) *gorm.DB
if slices.Contains(privs, "flows.admin") {
scope = func(db *gorm.DB) *gorm.DB {
return db.Where("id = ?", flowID)
}
} else if slices.Contains(privs, "flows.delete") {
scope = func(db *gorm.DB) *gorm.DB {
return db.Where("id = ? AND user_id = ?", flowID, uid)
}
} else {
logger.FromContext(c).Errorf("error filtering user role permissions: permission not found")
response.Error(c, response.ErrNotPermitted, nil)
return
}
if err = s.db.Model(&flow).Scopes(scope).Take(&flow).Error; err != nil {
logger.FromContext(c).WithError(err).Errorf("error getting flow by id")
if gorm.IsRecordNotFoundError(err) {
response.Error(c, response.ErrFlowsNotFound, err)
} else {
response.Error(c, response.ErrInternal, err)
}
return
}
if err := s.fc.FinishFlow(c, int64(flow.ID)); err != nil {
logger.FromContext(c).WithError(err).Errorf("error stopping flow")
response.Error(c, response.ErrInternal, err)
return
}
var containers []models.Container
err = s.db.Model(&containers).Where("flow_id = ?", flow.ID).Find(&containers).Error
if err != nil {
logger.FromContext(c).WithError(err).Errorf("error getting flow containers")
response.Error(c, response.ErrInternal, err)
return
}
if err = s.db.Scopes(scope).Delete(&flow).Error; err != nil {
logger.FromContext(c).WithError(err).Errorf("error deleting flow by id")
if gorm.IsRecordNotFoundError(err) {
response.Error(c, response.ErrFlowsNotFound, err)
} else {
response.Error(c, response.ErrInternal, err)
}
return
}
// Best-effort cleanup: remove accumulated long-term memory documents for this
// flow from the vector store. They will never be re-used after the flow is gone.
if err := s.db.Exec(
`DELETE FROM langchain_pg_embedding
WHERE collection_id = (SELECT uuid FROM langchain_pg_collection WHERE name = 'langchain')
AND COALESCE(cmetadata ->> 'doc_type', '') = 'memory'
AND (cmetadata ->> 'flow_id') = ?`,
strconv.FormatUint(flow.ID, 10),
).Error; err != nil {
logger.FromContext(c).WithError(err).Warnf("failed to clean up memory documents for deleted flow %d", flow.ID)
}
flowDB, err := convertFlowToDatabase(flow)
if err != nil {
logger.FromContext(c).WithError(err).Errorf("error converting flow to database")
response.Error(c, response.ErrInternal, err)
return
}
containersDB := make([]database.Container, 0, len(containers))
for _, container := range containers {
containersDB = append(containersDB, convertContainerToDatabase(container))
}
if s.ss != nil {
publisher := s.ss.NewFlowPublisher(int64(flow.UserID), int64(flow.ID))
publisher.FlowUpdated(c, flowDB, containersDB)
publisher.FlowDeleted(c, flowDB, containersDB)
}
response.Success(c, http.StatusOK, flow)
}
func convertFlowToDatabase(flow models.Flow) (database.Flow, error) {
functions, err := json.Marshal(flow.Functions)
if err != nil {
return database.Flow{}, err
}
return database.Flow{
ID: int64(flow.ID),
Status: database.FlowStatus(flow.Status),
Title: flow.Title,
Model: flow.Model,
ModelProviderName: flow.ModelProviderName,
Language: flow.Language,
Functions: functions,
UserID: int64(flow.UserID),
CreatedAt: database.TimeToNullTime(flow.CreatedAt),
UpdatedAt: database.TimeToNullTime(flow.UpdatedAt),
DeletedAt: database.PtrTimeToNullTime(flow.DeletedAt),
TraceID: database.PtrStringToNullString(flow.TraceID),
ModelProviderType: database.ProviderType(flow.ModelProviderType),
ToolCallIDTemplate: flow.ToolCallIDTemplate,
}, nil
}
func convertContainerToDatabase(container models.Container) database.Container {
return database.Container{
ID: int64(container.ID),
Type: database.ContainerType(container.Type),
Name: container.Name,
Image: container.Image,
Status: database.ContainerStatus(container.Status),
LocalID: database.StringToNullString(container.LocalID),
LocalDir: database.StringToNullString(container.LocalDir),
FlowID: int64(container.FlowID),
CreatedAt: database.TimeToNullTime(container.CreatedAt),
UpdatedAt: database.TimeToNullTime(container.UpdatedAt),
}
}
+227
View File
@@ -0,0 +1,227 @@
package services
import (
"context"
"fmt"
"net/http"
"slices"
"strings"
"time"
"pentagi/pkg/config"
"pentagi/pkg/controller"
"pentagi/pkg/database"
"pentagi/pkg/database/knowledge"
"pentagi/pkg/graph"
"pentagi/pkg/graph/subscriptions"
"pentagi/pkg/providers"
"pentagi/pkg/server/auth"
"pentagi/pkg/server/logger"
"pentagi/pkg/templates"
"github.com/99designs/gqlgen/graphql"
"github.com/99designs/gqlgen/graphql/handler"
"github.com/99designs/gqlgen/graphql/handler/extension"
"github.com/99designs/gqlgen/graphql/handler/lru"
"github.com/99designs/gqlgen/graphql/handler/transport"
"github.com/99designs/gqlgen/graphql/playground"
"github.com/gin-gonic/gin"
"github.com/gorilla/websocket"
"github.com/sirupsen/logrus"
"github.com/vektah/gqlparser/v2/ast"
"github.com/vxcontrol/cloud/anonymizer"
)
var (
_ graphql.RawParams // @lint:ignore U1000
_ graphql.Response // @lint:ignore U1000
)
type GraphqlService struct {
srv *handler.Server
play http.HandlerFunc
}
type originValidator struct {
allowAll bool
allowed []string
wildcards [][]string
wrappers []string
}
func NewGraphqlService(
db *database.Queries,
cfg *config.Config,
baseURL string,
origins []string,
tokenCache *auth.TokenCache,
providers providers.ProviderController,
controller controller.FlowController,
subscriptions subscriptions.SubscriptionsController,
knowledgeStore knowledge.KnowledgeStore,
replacer anonymizer.Replacer,
) *GraphqlService {
srv := handler.New(graph.NewExecutableSchema(graph.Config{Resolvers: &graph.Resolver{
DB: db,
Config: cfg,
Logger: logrus.StandardLogger().WithField("component", "pentagi-gql-bl"),
TokenCache: tokenCache,
DefaultPrompter: templates.NewDefaultPrompter(),
ProvidersCtrl: providers,
Controller: controller,
Subscriptions: subscriptions,
Knowledge: knowledgeStore,
Replacer: replacer,
}}))
component := "pentagi-gql"
srv.AroundResponses(logger.WithGqlLogger(component))
logger := logrus.WithField("component", component)
srv.AddTransport(transport.Options{})
srv.AddTransport(transport.GET{})
srv.AddTransport(transport.POST{})
srv.AddTransport(transport.MultipartForm{
MaxMemory: 32 << 20, // 32MB
})
srv.SetQueryCache(lru.New[*ast.QueryDocument](1000))
srv.Use(extension.Introspection{})
srv.Use(extension.AutomaticPersistedQuery{
Cache: lru.New[string](100),
})
srv.Use(extension.FixedComplexityLimit(20000))
ov := newOriginValidator(origins)
// Add transport to support GraphQL subscriptions
srv.AddTransport(&transport.Websocket{
KeepAlivePingInterval: 10 * time.Second,
InitFunc: func(ctx context.Context, initPayload transport.InitPayload) (context.Context, *transport.InitPayload, error) {
uid, err := graph.GetUserID(ctx)
if err != nil {
return nil, nil, fmt.Errorf("unauthorized: invalid user: %v", err)
}
logger.WithField("uid", uid).Info("graphql websocket upgrade")
return ctx, &initPayload, nil
},
Upgrader: websocket.Upgrader{
CheckOrigin: func(r *http.Request) (allowed bool) {
return ov.validateOrigin(r.Header.Get("Origin"), r.Host)
},
EnableCompression: true,
},
})
return &GraphqlService{
srv: srv,
play: playground.Handler("GraphQL", baseURL+"/graphql"),
}
}
// ServeGraphql is a function to perform graphql requests
// @Summary Perform graphql requests
// @Tags GraphQL
// @Accept json
// @Produce json
// @Security BearerAuth
// @Param json body graphql.RawParams true "graphql request"
// @Success 200 {object} graphql.Response "graphql response"
// @Failure 400 {object} graphql.Response "invalid graphql request data"
// @Failure 403 {object} graphql.Response "unauthorized"
// @Failure 500 {object} graphql.Response "internal error on graphql request"
// @Router /graphql [post]
func (s *GraphqlService) ServeGraphql(c *gin.Context) {
uid := c.GetUint64("uid")
tid := c.GetString("tid")
privs := c.GetStringSlice("prm")
savedCtx := c.Request.Context()
defer func() {
c.Request = c.Request.WithContext(savedCtx)
}()
ctx := savedCtx
ctx = graph.SetUserID(ctx, uid)
ctx = graph.SetUserType(ctx, tid)
ctx = graph.SetUserPermissions(ctx, privs)
c.Request = c.Request.WithContext(ctx)
s.srv.ServeHTTP(c.Writer, c.Request)
}
func (s *GraphqlService) ServeGraphqlPlayground(c *gin.Context) {
s.play.ServeHTTP(c.Writer, c.Request)
}
func newOriginValidator(origins []string) *originValidator {
var wRules [][]string
for _, o := range origins {
if !strings.Contains(o, "*") {
continue
}
if c := strings.Count(o, "*"); c > 1 {
continue
}
i := strings.Index(o, "*")
if i == 0 {
wRules = append(wRules, []string{"*", o[1:]})
continue
}
if i == (len(o) - 1) {
wRules = append(wRules, []string{o[:i], "*"})
continue
}
wRules = append(wRules, []string{o[:i], o[i+1:]})
}
return &originValidator{
allowAll: slices.Contains(origins, "*"),
allowed: origins,
wildcards: wRules,
wrappers: []string{"http://", "https://", "ws://", "wss://"},
}
}
func (ov *originValidator) validateOrigin(origin, host string) bool {
if ov.allowAll {
// CORS for origin '*' is allowed
return true
}
if len(origin) == 0 {
// request is not a CORS request
return true
}
for _, wrapper := range ov.wrappers {
if origin == wrapper+host {
// request is not a CORS request but have origin header
return true
}
}
if slices.Contains(ov.allowed, origin) {
return true
}
for _, w := range ov.wildcards {
if w[0] == "*" && strings.HasSuffix(origin, w[1]) {
return true
}
if w[1] == "*" && strings.HasPrefix(origin, w[0]) {
return true
}
if strings.HasPrefix(origin, w[0]) && strings.HasSuffix(origin, w[1]) {
return true
}
}
return false
}
+519
View File
@@ -0,0 +1,519 @@
package services
import (
"errors"
"net/http"
"strings"
knowledgepkg "pentagi/pkg/database/knowledge"
gqlmodel "pentagi/pkg/graph/model"
"pentagi/pkg/server/auth"
"pentagi/pkg/server/logger"
"pentagi/pkg/server/models"
"pentagi/pkg/server/rdb"
"pentagi/pkg/server/response"
"github.com/gin-gonic/gin"
"github.com/jinzhu/gorm"
)
// knowledgeSQLMappers maps rdb filter field names to the corresponding SQL
// expressions against langchain_pg_embedding / its cmetadata JSON column.
// Fields ending in _id are treated as numeric by rdb (operator = instead of like),
// so numeric JSON fields use custom function mappers that cast cmetadata values.
var knowledgeSQLMappers = map[string]any{
// Primary key exposed as "id" so TableQuery default ordering still works via SetOrders override.
"id": "langchain_pg_embedding.uuid::text",
// Text fields from cmetadata — rdb wraps them in LOWER(…::text) automatically.
"doc_type": "(langchain_pg_embedding.cmetadata ->> 'doc_type')",
"question": "(langchain_pg_embedding.cmetadata ->> 'question')",
"description": "(langchain_pg_embedding.cmetadata ->> 'description')",
"guide_type": "(langchain_pg_embedding.cmetadata ->> 'guide_type')",
"answer_type": "(langchain_pg_embedding.cmetadata ->> 'answer_type')",
"code_lang": "(langchain_pg_embedding.cmetadata ->> 'code_lang')",
// Searchable concat of common text fields for generic "data" filter.
"data": "((langchain_pg_embedding.cmetadata ->> 'doc_type') || ' ' || " +
"COALESCE(langchain_pg_embedding.cmetadata ->> 'question', ''))",
// Numeric ID fields: stored as JSON numbers; cast to bigint before comparison.
"user_id": func(q *rdb.TableQuery, db *gorm.DB, value any) *gorm.DB {
return applyKnowledgeIntFilter(db, "user_id", value)
},
"flow_id": func(q *rdb.TableQuery, db *gorm.DB, value any) *gorm.DB {
return applyKnowledgeIntFilter(db, "flow_id", value)
},
"task_id": func(q *rdb.TableQuery, db *gorm.DB, value any) *gorm.DB {
return applyKnowledgeIntFilter(db, "task_id", value)
},
"subtask_id": func(q *rdb.TableQuery, db *gorm.DB, value any) *gorm.DB {
return applyKnowledgeIntFilter(db, "subtask_id", value)
},
// Boolean field: 'true' / anything else (JSON booleans serialise as 'true'/'false').
"manual": func(q *rdb.TableQuery, db *gorm.DB, value any) *gorm.DB {
if b, ok := value.(bool); ok {
if b {
return db.Where("(langchain_pg_embedding.cmetadata ->> 'manual') = 'true'")
}
return db.Where("(langchain_pg_embedding.cmetadata ->> 'manual') IS DISTINCT FROM 'true'")
}
return db
},
}
// applyKnowledgeIntFilter adds a bigint equality / IN filter against a cmetadata integer field.
func applyKnowledgeIntFilter(db *gorm.DB, field string, value any) *gorm.DB {
expr := "(langchain_pg_embedding.cmetadata ->> '" + field + "')::bigint"
switch v := value.(type) {
case float64:
return db.Where(expr+" = ?", int64(v))
case []any:
var ids []int64
for _, vi := range v {
if f, ok := vi.(float64); ok {
ids = append(ids, int64(f))
}
}
if len(ids) > 0 {
return db.Where(expr+" IN (?)", ids)
}
}
return db
}
// knowledgeGrouped is the response shape for grouped list queries.
type knowledgeGrouped struct {
Grouped []string `json:"grouped"`
Total uint64 `json:"total"`
}
// KnowledgeService exposes the knowledge (pgvector) store via a REST API
// that mirrors the GraphQL knowledge operations.
//
// Authorization model:
// - knowledge.admin → read/write any document (no user_id filtering)
// - knowledge.view → list/get own documents (filtered by user_id)
// - knowledge.create → create a new document
// - knowledge.edit → update own document (admin: any)
// - knowledge.delete → delete own document (admin: any)
// - knowledge.search → semantic search (own docs for regular users)
type KnowledgeService struct {
db *gorm.DB
store knowledgepkg.KnowledgeStore
}
// NewKnowledgeService creates a KnowledgeService.
// db is the gorm connection used for paginated list queries via rdb.TableQuery.
// store may be nil when the embedding provider is not configured; in that
// case embedding-dependent endpoints return 503.
func NewKnowledgeService(db *gorm.DB, store knowledgepkg.KnowledgeStore) *KnowledgeService {
return &KnowledgeService{db: db, store: store}
}
func isKnowledgeAdmin(c *gin.Context) bool {
return auth.LookupPerm(c.GetStringSlice("prm"), "knowledge.admin")
}
// ---- ListDocuments ----------------------------------------------------------
// ListDocuments returns a paginated, sortable, filterable list of knowledge documents.
// It uses the standard rdb.TableQuery protocol (page, pageSize, sort[], filters[])
// plus an additional `with_content` boolean query parameter.
//
// Filterable fields:
// - id UUID (exact match)
// - doc_type answer | guide | code (text, like/=)
// - question question text (text, like)
// - description description text (text, like)
// - guide_type guide sub-type (text, like/=)
// - answer_type answer sub-type (text, like/=)
// - code_lang code language (text, like/=)
// - manual boolean flag (bool, =)
// - user_id owner user ID (int64, =/in)
// - flow_id originating flow ID (int64, =/in)
// - task_id originating task ID (int64, =/in)
// - subtask_id originating subtask ID (int64, =/in)
// - data full-text across doc_type + question (text, like)
//
// @Summary List knowledge documents
// @Tags Knowledge
// @Produce json
// @Security BearerAuth
// @Param request query rdb.TableQuery true "pagination / sort / filter params"
// @Param with_content query bool false "include document text in the response"
// @Success 200 {object} response.successResp{data=models.KnowledgeDocList}
// @Failure 400 {object} response.errorResp "invalid query parameters"
// @Failure 403 {object} response.errorResp "not permitted"
// @Failure 500 {object} response.errorResp "internal error"
// @Router /knowledge/ [get]
func (s *KnowledgeService) ListDocuments(c *gin.Context) {
uid := int64(c.GetUint64("uid"))
admin := isKnowledgeAdmin(c)
var query rdb.TableQuery
if err := c.ShouldBindQuery(&query); err != nil {
logger.FromContext(c).WithError(err).Error("error binding list query")
response.Error(c, response.ErrKnowledgeInvalidRequest, err)
return
}
withContent := c.Query("with_content") == "true" || c.Query("with_content") == "1"
if err := query.Init("langchain_pg_embedding", knowledgeSQLMappers); err != nil {
// Init fails when the requested group field is not a sortable string column
// (e.g. a custom function mapper like user_id).
logger.FromContext(c).WithError(err).Error("error initialising knowledge query")
response.Error(c, response.ErrKnowledgeInvalidRequest, err)
return
}
// Base scope: JOIN the collection table, restrict to 'langchain', exclude memory.
baseScope := func(db *gorm.DB) *gorm.DB {
return db.
Joins("JOIN langchain_pg_collection ON langchain_pg_embedding.collection_id = langchain_pg_collection.uuid").
Where("langchain_pg_collection.name = ?", "langchain").
Where("COALESCE(langchain_pg_embedding.cmetadata ->> 'doc_type', '') NOT IN (?)", []string{"memory"})
}
// Authorization scope: admins see all documents; regular users see only their own.
authScope := func(db *gorm.DB) *gorm.DB {
if admin {
return db
}
return db.Where("(langchain_pg_embedding.cmetadata ->> 'user_id')::bigint = ?", uid)
}
// ---- Grouped query -------------------------------------------------------
// When a group field is requested the response is a list of distinct values
// rather than full document rows (mirrors the pattern used in other services).
if query.Group != "" {
if _, ok := knowledgeSQLMappers[query.Group]; !ok {
logger.FromContext(c).Errorf("group field %q not found in knowledge mappers", query.Group)
response.Error(c, response.ErrKnowledgeInvalidRequest, errors.New("group field not found"))
return
}
var resp knowledgeGrouped
var err error
if resp.Total, err = query.QueryGrouped(s.db, &resp.Grouped, baseScope, authScope); err != nil {
logger.FromContext(c).WithError(err).Error("error querying knowledge documents grouped")
response.Error(c, response.ErrInternal, err)
return
}
response.Success(c, http.StatusOK, resp)
return
}
// ---- Paginated / sorted query --------------------------------------------
// Override the default "ORDER BY id DESC" — the PK column is uuid, not id.
query.SetOrders([]func(*gorm.DB) *gorm.DB{
func(db *gorm.DB) *gorm.DB {
return db.Order("langchain_pg_embedding.uuid DESC")
},
})
// SELECT only the columns the API exposes; never return the embedding vector.
query.SetFind(func(out any) func(*gorm.DB) *gorm.DB {
docCol := "'' AS document"
if withContent {
docCol = "COALESCE(langchain_pg_embedding.document, '') AS document"
}
sel := "langchain_pg_embedding.uuid::text AS id, " +
docCol + ", " +
"COALESCE(langchain_pg_embedding.cmetadata::text, '{}') AS cmetadata"
return func(db *gorm.DB) *gorm.DB {
return db.Select(sel).Find(out)
}
})
var rows []models.KnowledgeEmbeddingRow
total, err := query.Query(s.db, &rows, baseScope, authScope)
if err != nil {
logger.FromContext(c).WithError(err).Error("error querying knowledge documents")
response.Error(c, response.ErrInternal, err)
return
}
response.Success(c, http.StatusOK, models.KnowledgeDocListFromRows(rows, total, withContent))
}
// ---- GetDocument ------------------------------------------------------------
// GetDocument returns a single knowledge document by its UUID.
//
// @Summary Get knowledge document
// @Tags Knowledge
// @Produce json
// @Security BearerAuth
// @Param id path string true "Document UUID"
// @Success 200 {object} response.successResp{data=models.KnowledgeDocEntry}
// @Failure 403 {object} response.errorResp "not permitted or wrong owner"
// @Failure 404 {object} response.errorResp "document not found"
// @Failure 500 {object} response.errorResp "internal error"
// @Router /knowledge/{id} [get]
func (s *KnowledgeService) GetDocument(c *gin.Context) {
uid := int64(c.GetUint64("uid"))
admin := isKnowledgeAdmin(c)
id := c.Param("id")
if id == "" {
response.Error(c, response.ErrKnowledgeInvalidRequest, errors.New("document id is required"))
return
}
ctx := c.Request.Context()
var (
doc *gqlmodel.KnowledgeDocument
err error
)
if admin {
doc, err = s.store.GetDocument(ctx, id)
} else {
doc, err = s.store.GetUserDocument(ctx, uid, id)
}
if err != nil {
logger.FromContext(c).WithError(err).Errorf("error getting knowledge document %s", id)
response.Error(c, response.ErrKnowledgeNotFound, err)
return
}
response.Success(c, http.StatusOK, models.KnowledgeDocFromGQL(doc))
}
// ---- SearchDocuments --------------------------------------------------------
// SearchDocuments performs a semantic similarity search over the knowledge store.
//
// @Summary Semantic search in knowledge base
// @Tags Knowledge
// @Accept json
// @Produce json
// @Security BearerAuth
// @Param json body models.KnowledgeSearchRequest true "Search request"
// @Success 200 {object} response.successResp{data=models.KnowledgeSearchResult}
// @Failure 400 {object} response.errorResp "invalid request"
// @Failure 403 {object} response.errorResp "not permitted"
// @Failure 503 {object} response.errorResp "embedding provider not configured"
// @Failure 500 {object} response.errorResp "internal error"
// @Router /knowledge/search [post]
func (s *KnowledgeService) SearchDocuments(c *gin.Context) {
uid := int64(c.GetUint64("uid"))
admin := isKnowledgeAdmin(c)
var req models.KnowledgeSearchRequest
if err := c.ShouldBindJSON(&req); err != nil {
logger.FromContext(c).WithError(err).Error("error binding search request")
response.Error(c, response.ErrKnowledgeInvalidRequest, err)
return
}
if err := req.Valid(); err != nil {
logger.FromContext(c).WithError(err).Error("invalid search request")
response.Error(c, response.ErrKnowledgeInvalidRequest, err)
return
}
ctx := c.Request.Context()
filter := req.ToGQLFilter()
var (
results []*gqlmodel.KnowledgeDocumentWithScore
err error
)
if admin {
results, err = s.store.SearchDocuments(ctx, req.Query, filter, req.Limit)
} else {
results, err = s.store.SearchUserDocuments(ctx, uid, req.Query, filter, req.Limit)
}
if err != nil {
if isKnowledgeStoreUnavailable(err) {
response.Error(c, response.ErrKnowledgeStoreUnavail, err)
return
}
logger.FromContext(c).WithError(err).Error("error searching knowledge documents")
response.Error(c, response.ErrInternal, err)
return
}
response.Success(c, http.StatusOK, models.KnowledgeSearchResultFromGQL(results))
}
// ---- CreateDocument ---------------------------------------------------------
// CreateDocument creates a new knowledge document and computes its embedding.
//
// @Summary Create knowledge document
// @Tags Knowledge
// @Accept json
// @Produce json
// @Security BearerAuth
// @Param json body models.CreateKnowledgeDocRequest true "Create request"
// @Success 201 {object} response.successResp{data=models.KnowledgeDocEntry}
// @Failure 400 {object} response.errorResp "invalid request"
// @Failure 403 {object} response.errorResp "not permitted"
// @Failure 503 {object} response.errorResp "embedding provider not configured"
// @Failure 500 {object} response.errorResp "internal error"
// @Router /knowledge/ [post]
func (s *KnowledgeService) CreateDocument(c *gin.Context) {
uid := int64(c.GetUint64("uid"))
var req models.CreateKnowledgeDocRequest
if err := c.ShouldBindJSON(&req); err != nil {
logger.FromContext(c).WithError(err).Error("error binding create request")
response.Error(c, response.ErrKnowledgeInvalidRequest, err)
return
}
if err := req.Valid(); err != nil {
logger.FromContext(c).WithError(err).Error("invalid create request")
response.Error(c, response.ErrKnowledgeInvalidRequest, err)
return
}
doc, err := s.store.CreateDocument(c.Request.Context(), uid, req.ToGQL())
if err != nil {
if isKnowledgeStoreUnavailable(err) {
response.Error(c, response.ErrKnowledgeStoreUnavail, err)
return
}
logger.FromContext(c).WithError(err).Error("error creating knowledge document")
response.Error(c, response.ErrInternal, err)
return
}
response.Success(c, http.StatusCreated, models.KnowledgeDocFromGQL(doc))
}
// ---- UpdateDocument ---------------------------------------------------------
// UpdateDocument updates a knowledge document, re-computing its embedding.
// Admin can update any document; regular users can only update their own.
//
// @Summary Update knowledge document
// @Tags Knowledge
// @Accept json
// @Produce json
// @Security BearerAuth
// @Param id path string true "Document UUID"
// @Param json body models.UpdateKnowledgeDocRequest true "Update request"
// @Success 200 {object} response.successResp{data=models.KnowledgeDocEntry}
// @Failure 400 {object} response.errorResp "invalid request"
// @Failure 403 {object} response.errorResp "not permitted or wrong owner"
// @Failure 404 {object} response.errorResp "document not found"
// @Failure 503 {object} response.errorResp "embedding provider not configured"
// @Failure 500 {object} response.errorResp "internal error"
// @Router /knowledge/{id} [put]
func (s *KnowledgeService) UpdateDocument(c *gin.Context) {
uid := int64(c.GetUint64("uid"))
admin := isKnowledgeAdmin(c)
id := c.Param("id")
if id == "" {
response.Error(c, response.ErrKnowledgeInvalidRequest, errors.New("document id is required"))
return
}
var req models.UpdateKnowledgeDocRequest
if err := c.ShouldBindJSON(&req); err != nil {
logger.FromContext(c).WithError(err).Error("error binding update request")
response.Error(c, response.ErrKnowledgeInvalidRequest, err)
return
}
if err := req.Valid(); err != nil {
logger.FromContext(c).WithError(err).Error("invalid update request")
response.Error(c, response.ErrKnowledgeInvalidRequest, err)
return
}
ctx := c.Request.Context()
input := req.ToGQL()
var (
doc *gqlmodel.KnowledgeDocument
err error
)
if admin {
doc, err = s.store.UpdateDocument(ctx, uid, id, input)
} else {
doc, err = s.store.UpdateUserDocument(ctx, uid, id, input)
}
if err != nil {
if isKnowledgeStoreUnavailable(err) {
response.Error(c, response.ErrKnowledgeStoreUnavail, err)
return
}
if isKnowledgeNotFound(err) {
response.Error(c, response.ErrKnowledgeNotFound, err)
return
}
logger.FromContext(c).WithError(err).Errorf("error updating knowledge document %s", id)
response.Error(c, response.ErrInternal, err)
return
}
response.Success(c, http.StatusOK, models.KnowledgeDocFromGQL(doc))
}
// ---- DeleteDocument ---------------------------------------------------------
// DeleteDocument deletes a knowledge document by UUID.
// Admin can delete any document; regular users can only delete their own.
//
// @Summary Delete knowledge document
// @Tags Knowledge
// @Produce json
// @Security BearerAuth
// @Param id path string true "Document UUID"
// @Success 200 {object} response.successResp "document deleted"
// @Failure 403 {object} response.errorResp "not permitted or wrong owner"
// @Failure 404 {object} response.errorResp "document not found"
// @Failure 500 {object} response.errorResp "internal error"
// @Router /knowledge/{id} [delete]
func (s *KnowledgeService) DeleteDocument(c *gin.Context) {
uid := int64(c.GetUint64("uid"))
admin := isKnowledgeAdmin(c)
id := c.Param("id")
if id == "" {
response.Error(c, response.ErrKnowledgeInvalidRequest, errors.New("document id is required"))
return
}
ctx := c.Request.Context()
var err error
if admin {
err = s.store.DeleteDocument(ctx, uid, id)
} else {
err = s.store.DeleteUserDocument(ctx, uid, id)
}
if err != nil {
if isKnowledgeNotFound(err) {
response.Error(c, response.ErrKnowledgeNotFound, err)
return
}
logger.FromContext(c).WithError(err).Errorf("error deleting knowledge document %s", id)
response.Error(c, response.ErrInternal, err)
return
}
response.Success(c, http.StatusOK, gin.H{"message": "knowledge document deleted successfully"})
}
// ---- error classification ---------------------------------------------------
// isKnowledgeStoreUnavailable returns true when the error originates from a
// missing embedding provider (store or embedder not configured).
func isKnowledgeStoreUnavailable(err error) bool {
return err != nil && strings.Contains(err.Error(), "knowledge: embedding provider")
}
// isKnowledgeNotFound returns true when the underlying DB returned no rows,
// which means the document does not exist or does not belong to the caller.
func isKnowledgeNotFound(err error) bool {
return err != nil && (strings.Contains(err.Error(), "no rows") ||
strings.Contains(err.Error(), "not found"))
}
+221
View File
@@ -0,0 +1,221 @@
package services
import (
"errors"
"net/http"
"slices"
"strconv"
"pentagi/pkg/server/logger"
"pentagi/pkg/server/models"
"pentagi/pkg/server/rdb"
"pentagi/pkg/server/response"
"github.com/gin-gonic/gin"
"github.com/jinzhu/gorm"
)
type msglogs struct {
MsgLogs []models.Msglog `json:"msglogs"`
Total uint64 `json:"total"`
}
type msglogsGrouped struct {
Grouped []string `json:"grouped"`
Total uint64 `json:"total"`
}
var msglogsSQLMappers = map[string]any{
"id": "{{table}}.id",
"type": "{{table}}.type",
"message": "{{table}}.message",
"thinking": "{{table}}.thinking",
"result": "{{table}}.result",
"result_format": "{{table}}.result_format",
"flow_id": "{{table}}.flow_id",
"task_id": "{{table}}.task_id",
"subtask_id": "{{table}}.subtask_id",
"created_at": "{{table}}.created_at",
"data": "({{table}}.type || ' ' || {{table}}.message || ' ' || {{table}}.thinking || ' ' || {{table}}.result)",
}
type MsglogService struct {
db *gorm.DB
}
func NewMsglogService(db *gorm.DB) *MsglogService {
return &MsglogService{
db: db,
}
}
// GetMsglogs is a function to return msglogs list
// @Summary Retrieve msglogs list
// @Tags Msglogs
// @Produce json
// @Security BearerAuth
// @Param request query rdb.TableQuery true "query table params"
// @Success 200 {object} response.successResp{data=msglogs} "msglogs list received successful"
// @Failure 400 {object} response.errorResp "invalid query request data"
// @Failure 403 {object} response.errorResp "getting msglogs not permitted"
// @Failure 500 {object} response.errorResp "internal error on getting msglogs"
// @Router /msglogs/ [get]
func (s *MsglogService) GetMsglogs(c *gin.Context) {
var (
err error
query rdb.TableQuery
resp msglogs
)
if err = c.ShouldBindQuery(&query); err != nil {
logger.FromContext(c).WithError(err).Errorf("error binding query")
response.Error(c, response.ErrMsglogsInvalidRequest, err)
return
}
uid := c.GetUint64("uid")
privs := c.GetStringSlice("prm")
var scope func(db *gorm.DB) *gorm.DB
if slices.Contains(privs, "msglogs.admin") {
scope = func(db *gorm.DB) *gorm.DB {
return db.
Joins("INNER JOIN flows f ON f.id = msglogs.flow_id")
}
} else if slices.Contains(privs, "msglogs.view") {
scope = func(db *gorm.DB) *gorm.DB {
return db.
Joins("INNER JOIN flows f ON f.id = msglogs.flow_id").
Where("f.user_id = ?", uid)
}
} else {
logger.FromContext(c).Errorf("error filtering user role permissions: permission not found")
response.Error(c, response.ErrNotPermitted, nil)
return
}
query.Init("msglogs", msglogsSQLMappers)
if query.Group != "" {
if _, ok := msglogsSQLMappers[query.Group]; !ok {
logger.FromContext(c).Errorf("error finding msglogs grouped: group field not found")
response.Error(c, response.ErrMsglogsInvalidRequest, errors.New("group field not found"))
return
}
var respGrouped msglogsGrouped
if respGrouped.Total, err = query.QueryGrouped(s.db, &respGrouped.Grouped, scope); err != nil {
logger.FromContext(c).WithError(err).Errorf("error finding msglogs grouped")
response.Error(c, response.ErrInternal, err)
return
}
response.Success(c, http.StatusOK, respGrouped)
return
}
if resp.Total, err = query.Query(s.db, &resp.MsgLogs, scope); err != nil {
logger.FromContext(c).WithError(err).Errorf("error finding msglogs")
response.Error(c, response.ErrInternal, err)
return
}
for i := 0; i < len(resp.MsgLogs); i++ {
if err = resp.MsgLogs[i].Valid(); err != nil {
logger.FromContext(c).WithError(err).Errorf("error validating msglog data '%d'", resp.MsgLogs[i].ID)
response.Error(c, response.ErrMsglogsInvalidData, err)
return
}
}
response.Success(c, http.StatusOK, resp)
}
// GetFlowMsglogs is a function to return msglogs list by flow id
// @Summary Retrieve msglogs list by flow id
// @Tags Msglogs
// @Produce json
// @Security BearerAuth
// @Param flowID path int true "flow id" minimum(0)
// @Param request query rdb.TableQuery true "query table params"
// @Success 200 {object} response.successResp{data=msglogs} "msglogs list received successful"
// @Failure 400 {object} response.errorResp "invalid query request data"
// @Failure 403 {object} response.errorResp "getting msglogs not permitted"
// @Failure 500 {object} response.errorResp "internal error on getting msglogs"
// @Router /flows/{flowID}/msglogs/ [get]
func (s *MsglogService) GetFlowMsglogs(c *gin.Context) {
var (
err error
flowID uint64
query rdb.TableQuery
resp msglogs
)
if flowID, err = strconv.ParseUint(c.Param("flowID"), 10, 64); err != nil {
logger.FromContext(c).WithError(err).Errorf("error parsing flow id")
response.Error(c, response.ErrMsglogsInvalidRequest, err)
return
}
if err = c.ShouldBindQuery(&query); err != nil {
logger.FromContext(c).WithError(err).Errorf("error binding query")
response.Error(c, response.ErrMsglogsInvalidRequest, err)
return
}
uid := c.GetUint64("uid")
privs := c.GetStringSlice("prm")
var scope func(db *gorm.DB) *gorm.DB
if slices.Contains(privs, "msglogs.admin") {
scope = func(db *gorm.DB) *gorm.DB {
return db.
Joins("INNER JOIN flows f ON f.id = msglogs.flow_id").
Where("f.id = ?", flowID)
}
} else if slices.Contains(privs, "msglogs.view") {
scope = func(db *gorm.DB) *gorm.DB {
return db.
Joins("INNER JOIN flows f ON f.id = msglogs.flow_id").
Where("f.id = ? AND f.user_id = ?", flowID, uid)
}
} else {
logger.FromContext(c).Errorf("error filtering user role permissions: permission not found")
response.Error(c, response.ErrNotPermitted, nil)
return
}
query.Init("msglogs", msglogsSQLMappers)
if query.Group != "" {
if _, ok := msglogsSQLMappers[query.Group]; !ok {
logger.FromContext(c).Errorf("error finding msglogs grouped: group field not found")
response.Error(c, response.ErrMsglogsInvalidRequest, errors.New("group field not found"))
return
}
var respGrouped msglogsGrouped
if respGrouped.Total, err = query.QueryGrouped(s.db, &respGrouped.Grouped, scope); err != nil {
logger.FromContext(c).WithError(err).Errorf("error finding msglogs grouped")
response.Error(c, response.ErrInternal, err)
return
}
response.Success(c, http.StatusOK, respGrouped)
return
}
if resp.Total, err = query.Query(s.db, &resp.MsgLogs, scope); err != nil {
logger.FromContext(c).WithError(err).Errorf("error finding msglogs")
response.Error(c, response.ErrInternal, err)
return
}
for i := 0; i < len(resp.MsgLogs); i++ {
if err = resp.MsgLogs[i].Valid(); err != nil {
logger.FromContext(c).WithError(err).Errorf("error validating msglog data '%d'", resp.MsgLogs[i].ID)
response.Error(c, response.ErrMsglogsInvalidData, err)
return
}
}
response.Success(c, http.StatusOK, resp)
}
+391
View File
@@ -0,0 +1,391 @@
package services
import (
"errors"
"net/http"
"slices"
"pentagi/pkg/server/logger"
"pentagi/pkg/server/models"
"pentagi/pkg/server/rdb"
"pentagi/pkg/server/response"
"pentagi/pkg/templates"
"github.com/gin-gonic/gin"
"github.com/jinzhu/gorm"
)
type prompts struct {
Prompts []models.Prompt `json:"prompts"`
Total uint64 `json:"total"`
}
type promptsGrouped struct {
Grouped []string `json:"grouped"`
Total uint64 `json:"total"`
}
var promptsSQLMappers = map[string]any{
"type": "{{table}}.type",
"prompt": "{{table}}.prompt",
"created_at": "{{table}}.created_at",
"updated_at": "{{table}}.updated_at",
"data": "({{table}}.type || ' ' || {{table}}.prompt)",
}
type PromptService struct {
db *gorm.DB
prompter templates.Prompter
}
func NewPromptService(db *gorm.DB) *PromptService {
return &PromptService{
db: db,
prompter: templates.NewDefaultPrompter(),
}
}
// GetPrompts is a function to return prompts list
// @Summary Retrieve prompts list
// @Tags Prompts
// @Produce json
// @Security BearerAuth
// @Param request query rdb.TableQuery true "query table params"
// @Success 200 {object} response.successResp{data=prompts} "prompts list received successful"
// @Failure 400 {object} response.errorResp "invalid query request data"
// @Failure 403 {object} response.errorResp "getting prompts not permitted"
// @Failure 500 {object} response.errorResp "internal error on getting prompts"
// @Router /prompts/ [get]
func (s *PromptService) GetPrompts(c *gin.Context) {
var (
err error
query rdb.TableQuery
resp prompts
)
if err = c.ShouldBindQuery(&query); err != nil {
logger.FromContext(c).WithError(err).Errorf("error binding query")
response.Error(c, response.ErrPromptsInvalidRequest, err)
return
}
privs := c.GetStringSlice("prm")
if !slices.Contains(privs, "settings.prompts.view") {
logger.FromContext(c).Errorf("error filtering user role permissions: permission not found")
response.Error(c, response.ErrNotPermitted, nil)
return
}
uid := c.GetUint64("uid")
scope := func(db *gorm.DB) *gorm.DB {
return db.Where("user_id = ?", uid)
}
query.Init("prompts", promptsSQLMappers)
if query.Group != "" {
if _, ok := promptsSQLMappers[query.Group]; !ok {
logger.FromContext(c).Errorf("error finding prompts grouped: group field not found")
response.Error(c, response.ErrPromptsInvalidRequest, errors.New("group field not found"))
return
}
var respGrouped promptsGrouped
if respGrouped.Total, err = query.QueryGrouped(s.db, &respGrouped.Grouped, scope); err != nil {
logger.FromContext(c).WithError(err).Errorf("error finding prompts grouped")
response.Error(c, response.ErrInternal, err)
return
}
response.Success(c, http.StatusOK, respGrouped)
return
}
if resp.Total, err = query.Query(s.db, &resp.Prompts, scope); err != nil {
logger.FromContext(c).WithError(err).Errorf("error finding prompts")
response.Error(c, response.ErrInternal, err)
return
}
for i := 0; i < len(resp.Prompts); i++ {
if err = resp.Prompts[i].Valid(); err != nil {
logger.FromContext(c).WithError(err).Errorf("error validating prompt data '%s'", resp.Prompts[i].Type)
response.Error(c, response.ErrPromptsInvalidData, err)
return
}
}
response.Success(c, http.StatusOK, resp)
}
// GetPrompt is a function to return prompt by type
// @Summary Retrieve prompt by type
// @Tags Prompts
// @Produce json
// @Security BearerAuth
// @Param promptType path string true "prompt type"
// @Success 200 {object} response.successResp{data=models.Prompt} "prompt received successful"
// @Failure 400 {object} response.errorResp "invalid prompt request data"
// @Failure 403 {object} response.errorResp "getting prompt not permitted"
// @Failure 404 {object} response.errorResp "prompt not found"
// @Failure 500 {object} response.errorResp "internal error on getting prompt"
// @Router /prompts/{promptType} [get]
func (s *PromptService) GetPrompt(c *gin.Context) {
var (
err error
promptType models.PromptType = models.PromptType(c.Param("promptType"))
resp models.Prompt
)
if err = models.PromptType(promptType).Valid(); err != nil {
logger.FromContext(c).WithError(err).Errorf("error validating prompt type '%s'", promptType)
response.Error(c, response.ErrPromptsInvalidRequest, err)
return
}
privs := c.GetStringSlice("prm")
if !slices.Contains(privs, "settings.prompts.view") {
logger.FromContext(c).Errorf("error filtering user role permissions: permission not found")
response.Error(c, response.ErrNotPermitted, nil)
return
}
uid := c.GetUint64("uid")
scope := func(db *gorm.DB) *gorm.DB {
return db.Where("type = ? AND user_id = ?", promptType, uid)
}
if err = s.db.Scopes(scope).Take(&resp).Error; err != nil {
logger.FromContext(c).WithError(err).Errorf("error finding prompt by type")
if errors.Is(err, gorm.ErrRecordNotFound) {
response.Error(c, response.ErrPromptsNotFound, err)
} else {
response.Error(c, response.ErrInternal, err)
}
return
}
if err = resp.Valid(); err != nil {
logger.FromContext(c).WithError(err).Errorf("error validating prompt data '%s'", resp.Type)
response.Error(c, response.ErrPromptsInvalidData, err)
return
}
response.Success(c, http.StatusOK, resp)
}
// PatchPrompt is a function to update prompt by type
// @Summary Update prompt
// @Tags Prompts
// @Accept json
// @Produce json
// @Security BearerAuth
// @Param promptType path string true "prompt type"
// @Param json body models.PatchPrompt true "prompt model to update"
// @Success 200 {object} response.successResp{data=models.Prompt} "prompt updated successful"
// @Success 201 {object} response.successResp{data=models.Prompt} "prompt created successful"
// @Failure 400 {object} response.errorResp "invalid prompt request data"
// @Failure 403 {object} response.errorResp "updating prompt not permitted"
// @Failure 404 {object} response.errorResp "prompt not found"
// @Failure 500 {object} response.errorResp "internal error on updating prompt"
// @Router /prompts/{promptType} [put]
func (s *PromptService) PatchPrompt(c *gin.Context) {
var (
err error
prompt models.PatchPrompt
promptType models.PromptType = models.PromptType(c.Param("promptType"))
resp models.Prompt
)
if err = c.ShouldBindJSON(&prompt); err != nil {
logger.FromContext(c).WithError(err).Errorf("error binding JSON")
response.Error(c, response.ErrPromptsInvalidRequest, err)
return
} else if err = prompt.Valid(); err != nil {
logger.FromContext(c).WithError(err).Errorf("error validating prompt JSON")
response.Error(c, response.ErrPromptsInvalidRequest, err)
return
} else if err = promptType.Valid(); err != nil {
logger.FromContext(c).WithError(err).Errorf("error validating prompt type '%s'", promptType)
response.Error(c, response.ErrPromptsInvalidRequest, err)
return
}
privs := c.GetStringSlice("prm")
if !slices.Contains(privs, "settings.prompts.edit") {
logger.FromContext(c).Errorf("error filtering user role permissions: permission not found")
response.Error(c, response.ErrNotPermitted, nil)
return
}
uid := c.GetUint64("uid")
scope := func(db *gorm.DB) *gorm.DB {
return db.Where("type = ? AND user_id = ?", promptType, uid)
}
err = s.db.Scopes(scope).Take(&resp).Error
if err != nil && errors.Is(err, gorm.ErrRecordNotFound) {
resp = models.Prompt{
Type: promptType,
UserID: uid,
Prompt: prompt.Prompt,
}
if err = s.db.Create(&resp).Error; err != nil {
logger.FromContext(c).WithError(err).Errorf("error creating prompt by type '%s'", promptType)
response.Error(c, response.ErrInternal, err)
return
}
response.Success(c, http.StatusCreated, resp)
return
} else if err != nil {
logger.FromContext(c).WithError(err).Errorf("error finding updated prompt by type '%s'", promptType)
response.Error(c, response.ErrInternal, err)
return
}
resp.Prompt = prompt.Prompt
err = s.db.Scopes(scope).Save(&resp).Error
if err != nil {
logger.FromContext(c).WithError(err).Errorf("error updating prompt by type '%s'", promptType)
response.Error(c, response.ErrInternal, err)
return
}
response.Success(c, http.StatusOK, resp)
}
// ResetPrompt is a function to reset prompt by type to default value
// @Summary Reset prompt by type to default value
// @Tags Prompts
// @Accept json
// @Produce json
// @Security BearerAuth
// @Param promptType path string true "prompt type"
// @Success 200 {object} response.successResp{data=models.Prompt} "prompt reset successful"
// @Success 201 {object} response.successResp{data=models.Prompt} "prompt created with default value successful"
// @Failure 400 {object} response.errorResp "invalid prompt request data"
// @Failure 403 {object} response.errorResp "updating prompt not permitted"
// @Failure 404 {object} response.errorResp "prompt not found"
// @Failure 500 {object} response.errorResp "internal error on resetting prompt"
// @Router /prompts/{promptType}/default [post]
func (s *PromptService) ResetPrompt(c *gin.Context) {
var (
err error
promptType models.PromptType = models.PromptType(c.Param("promptType"))
resp models.Prompt
)
if err = promptType.Valid(); err != nil {
logger.FromContext(c).WithError(err).Errorf("error validating prompt type '%s'", promptType)
response.Error(c, response.ErrPromptsInvalidRequest, err)
return
}
privs := c.GetStringSlice("prm")
if !slices.Contains(privs, "settings.prompts.edit") {
logger.FromContext(c).Errorf("error filtering user role permissions: permission not found")
response.Error(c, response.ErrNotPermitted, nil)
return
}
uid := c.GetUint64("uid")
scope := func(db *gorm.DB) *gorm.DB {
return db.Where("type = ? AND user_id = ?", promptType, uid)
}
template, err := s.prompter.GetTemplate(templates.PromptType(promptType))
if err != nil {
logger.FromContext(c).WithError(err).Errorf("error getting template '%s'", promptType)
response.Error(c, response.ErrPromptsInvalidRequest, err)
return
}
err = s.db.Scopes(scope).Take(&resp).Error
if err != nil && errors.Is(err, gorm.ErrRecordNotFound) {
resp = models.Prompt{
Type: promptType,
UserID: uid,
Prompt: template,
}
err = s.db.Create(&resp).Error
if err != nil {
logger.FromContext(c).WithError(err).Errorf("error creating default prompt by type '%s'", promptType)
response.Error(c, response.ErrInternal, err)
}
response.Success(c, http.StatusCreated, resp)
return
} else if err != nil {
logger.FromContext(c).WithError(err).Errorf("error finding updated prompt by type '%s'", promptType)
response.Error(c, response.ErrInternal, err)
return
}
resp.Prompt = template
err = s.db.Scopes(scope).Save(&resp).Error
if err != nil {
logger.FromContext(c).WithError(err).Errorf("error resetting prompt by type '%s'", promptType)
response.Error(c, response.ErrInternal, err)
return
}
response.Success(c, http.StatusOK, resp)
}
// DeletePrompt is a function to delete prompt by type
// @Summary Delete prompt by type
// @Tags Prompts
// @Produce json
// @Security BearerAuth
// @Param promptType path string true "prompt type"
// @Success 200 {object} response.successResp "prompt deleted successful"
// @Failure 400 {object} response.errorResp "invalid prompt request data"
// @Failure 403 {object} response.errorResp "deleting prompt not permitted"
// @Failure 404 {object} response.errorResp "prompt not found"
// @Failure 500 {object} response.errorResp "internal error on deleting prompt"
// @Router /prompts/{promptType} [delete]
func (s *PromptService) DeletePrompt(c *gin.Context) {
var (
err error
promptType models.PromptType = models.PromptType(c.Param("promptType"))
resp models.Prompt
)
if err = promptType.Valid(); err != nil {
logger.FromContext(c).WithError(err).Errorf("error validating prompt type '%s'", promptType)
response.Error(c, response.ErrPromptsInvalidRequest, err)
return
}
privs := c.GetStringSlice("prm")
if !slices.Contains(privs, "settings.prompts.edit") {
logger.FromContext(c).Errorf("error filtering user role permissions: permission not found")
response.Error(c, response.ErrNotPermitted, nil)
return
}
uid := c.GetUint64("uid")
scope := func(db *gorm.DB) *gorm.DB {
return db.Where("type = ? AND user_id = ?", promptType, uid)
}
if err = s.db.Scopes(scope).Take(&resp).Error; err != nil {
logger.FromContext(c).WithError(err).Errorf("error finding prompt by type '%s'", promptType)
if errors.Is(err, gorm.ErrRecordNotFound) {
response.Error(c, response.ErrPromptsNotFound, err)
} else {
response.Error(c, response.ErrInternal, err)
}
return
}
if err = s.db.Scopes(scope).Delete(&resp).Error; err != nil {
logger.FromContext(c).WithError(err).Errorf("error deleting prompt by type '%s'", promptType)
response.Error(c, response.ErrInternal, err)
return
}
response.Success(c, http.StatusOK, nil)
}
+123
View File
@@ -0,0 +1,123 @@
package services
import (
"net/http"
"slices"
"pentagi/pkg/providers"
"pentagi/pkg/providers/pconfig"
"pentagi/pkg/providers/provider"
"pentagi/pkg/server/logger"
"pentagi/pkg/server/models"
"pentagi/pkg/server/response"
"github.com/gin-gonic/gin"
)
type ProviderService struct {
providers providers.ProviderController
}
func NewProviderService(providers providers.ProviderController) *ProviderService {
return &ProviderService{
providers: providers,
}
}
// GetProviders is a function to return providers list
// @Summary Retrieve providers list
// @Tags Providers
// @Produce json
// @Security BearerAuth
// @Success 200 {object} response.successResp{data=[]models.ProviderInfo} "providers list received successful"
// @Failure 403 {object} response.errorResp "getting providers not permitted"
// @Router /providers/ [get]
func (s *ProviderService) GetProviders(c *gin.Context) {
privs := c.GetStringSlice("prm")
if !slices.Contains(privs, "providers.view") {
logger.FromContext(c).Errorf("error filtering user role permissions: permission not found")
response.Error(c, response.ErrNotPermitted, nil)
return
}
providers, err := s.providers.GetProviders(c, int64(c.GetUint64("uid")))
if err != nil {
logger.FromContext(c).Errorf("error getting providers: %v", err)
response.Error(c, response.ErrInternal, nil)
return
}
providerInfos := make([]models.ProviderInfo, len(providers))
for i, name := range providers.ListNames() {
prv := providers[name]
providerInfos[i] = models.ProviderInfo{
Name: name.String(),
Type: models.ProviderType(prv.Type()),
DefaultModel: prv.Model(pconfig.OptionsTypePrimaryAgent),
Models: buildModelInfos(prv),
}
}
response.Success(c, http.StatusOK, providerInfos)
}
func buildModelInfos(prv provider.Provider) []models.ModelInfo {
modelsConfig := prv.GetModels()
// Build lookup: model name -> price from ModelsConfig (models.yml)
modelConfigPrice := make(map[string]*pconfig.PriceInfo, len(modelsConfig))
for _, mc := range modelsConfig {
if mc.Price != nil {
modelConfigPrice[mc.Name] = mc.Price
}
}
// Collect unique models actually in use across all agent types.
// For price priority: AgentConfig.Price > ModelsConfig price.
type entry struct {
price *pconfig.PriceInfo
agentTypes []string
}
seen := make(map[string]*entry)
order := make([]string, 0)
for _, optType := range pconfig.AllAgentTypes {
modelName := prv.Model(optType)
if modelName == "" {
continue
}
e, exists := seen[modelName]
if !exists {
price := prv.GetPriceInfo(optType)
if price == nil {
price = modelConfigPrice[modelName]
}
e = &entry{price: price}
seen[modelName] = e
order = append(order, modelName)
}
e.agentTypes = append(e.agentTypes, string(optType))
}
modelInfos := make([]models.ModelInfo, 0, len(order))
for _, name := range order {
e := seen[name]
mi := models.ModelInfo{
Name: name,
AgentTypes: e.agentTypes,
}
if e.price != nil {
mi.PriceInfo = &models.ModelPriceInfo{
Input: e.price.Input,
Output: e.price.Output,
CacheRead: e.price.CacheRead,
CacheWrite: e.price.CacheWrite,
}
}
modelInfos = append(modelInfos, mi)
}
return modelInfos
}
File diff suppressed because it is too large Load Diff
File diff suppressed because it is too large Load Diff
+164
View File
@@ -0,0 +1,164 @@
package services
import (
"errors"
"net/http"
"slices"
"strconv"
"pentagi/pkg/server/logger"
"pentagi/pkg/server/models"
"pentagi/pkg/server/rdb"
"pentagi/pkg/server/response"
"github.com/gin-gonic/gin"
"github.com/jinzhu/gorm"
)
type roles struct {
Roles []models.RolePrivileges `json:"roles"`
Total uint64 `json:"total"`
}
var rolesSQLMappers = map[string]any{
"id": "{{table}}.id",
"name": "{{table}}.name",
"data": "{{table}}.name",
}
type RoleService struct {
db *gorm.DB
}
func NewRoleService(db *gorm.DB) *RoleService {
return &RoleService{
db: db,
}
}
// GetRoles is a function to return roles list
// @Summary Retrieve roles list
// @Tags Roles
// @Produce json
// @Param request query rdb.TableQuery true "query table params"
// @Success 200 {object} response.successResp{data=roles} "roles list received successful"
// @Failure 400 {object} response.errorResp "invalid query request data"
// @Failure 403 {object} response.errorResp "getting roles not permitted"
// @Failure 500 {object} response.errorResp "internal error on getting roles"
// @Router /roles/ [get]
func (s *RoleService) GetRoles(c *gin.Context) {
var (
err error
query rdb.TableQuery
resp roles
rids []uint64
)
if err = c.ShouldBindQuery(&query); err != nil {
logger.FromContext(c).WithError(err).Errorf("error binding query")
response.Error(c, response.ErrRolesInvalidRequest, err)
return
}
rid := c.GetUint64("rid")
privs := c.GetStringSlice("prm")
scope := func(db *gorm.DB) *gorm.DB {
if !slices.Contains(privs, "roles.view") {
return db.Where("role_id = ?", rid)
}
return db
}
query.Init("roles", rolesSQLMappers)
if resp.Total, err = query.Query(s.db, &resp.Roles, scope); err != nil {
logger.FromContext(c).WithError(err).Errorf("error finding roles")
response.Error(c, response.ErrInternal, err)
return
}
for _, role := range resp.Roles {
rids = append(rids, role.ID)
}
var privsObjs []models.Privilege
if err = s.db.Find(&privsObjs, "role_id IN (?)", rids).Error; err != nil {
logger.FromContext(c).WithError(err).Errorf("error finding linked roles")
response.Error(c, response.ErrInternal, err)
return
}
privsRoles := make(map[uint64][]models.Privilege)
for i := range privsObjs {
privsRoles[privsObjs[i].RoleID] = append(privsRoles[privsObjs[i].RoleID], privsObjs[i])
}
for i := range resp.Roles {
resp.Roles[i].Privileges = privsRoles[resp.Roles[i].ID]
}
for i := 0; i < len(resp.Roles); i++ {
if err = resp.Roles[i].Valid(); err != nil {
logger.FromContext(c).WithError(err).Errorf("error validating role data '%d'", resp.Roles[i].ID)
response.Error(c, response.ErrRolesInvalidData, err)
return
}
}
response.Success(c, http.StatusOK, resp)
}
// GetRole is a function to return role by id
// @Summary Retrieve role by id
// @Tags Roles
// @Produce json
// @Param roleID path uint64 true "role id"
// @Success 200 {object} response.successResp{data=models.RolePrivileges} "role received successful"
// @Failure 400 {object} response.errorResp "invalid query request data"
// @Failure 403 {object} response.errorResp "getting role not permitted"
// @Failure 500 {object} response.errorResp "internal error on getting role"
// @Router /roles/{roleID} [get]
func (s *RoleService) GetRole(c *gin.Context) {
var (
err error
resp models.RolePrivileges
roleID uint64
)
rid := c.GetUint64("rid")
privs := c.GetStringSlice("prm")
scope := func(db *gorm.DB) *gorm.DB {
if !slices.Contains(privs, "roles.view") {
return db.Where("role_id = ?", rid)
}
return db
}
if roleID, err = strconv.ParseUint(c.Param("roleID"), 10, 64); err != nil {
logger.FromContext(c).WithError(err).Errorf("error parsing role id")
response.Error(c, response.ErrRolesInvalidRequest, err)
return
}
if err := s.db.Scopes(scope).Take(&resp, "id = ?", roleID).Error; err != nil {
logger.FromContext(c).WithError(err).Errorf("error finding role by id")
if errors.Is(err, gorm.ErrRecordNotFound) {
response.Error(c, response.ErrRolesNotFound, err)
} else {
response.Error(c, response.ErrInternal, err)
}
return
}
if err := s.db.Model(&resp).Association("privileges").Find(&resp.Privileges).Error; err != nil {
logger.FromContext(c).WithError(err).Errorf("error finding role privileges by role model")
response.Error(c, response.ErrInternal, err)
return
}
if err := resp.Valid(); err != nil {
logger.FromContext(c).WithError(err).Errorf("error validating role data '%d'", resp.ID)
response.Error(c, response.ErrRolesInvalidData, err)
return
}
response.Success(c, http.StatusOK, resp)
}
+354
View File
@@ -0,0 +1,354 @@
package services
import (
"errors"
"fmt"
"net/http"
"path/filepath"
"slices"
"strconv"
"pentagi/pkg/server/logger"
"pentagi/pkg/server/models"
"pentagi/pkg/server/rdb"
"pentagi/pkg/server/response"
"github.com/gin-gonic/gin"
"github.com/jinzhu/gorm"
)
type screenshots struct {
Screenshots []models.Screenshot `json:"screenshots"`
Total uint64 `json:"total"`
}
type screenshotsGrouped struct {
Grouped []string `json:"grouped"`
Total uint64 `json:"total"`
}
var screenshotsSQLMappers = map[string]any{
"id": "{{table}}.id",
"name": "{{table}}.name",
"url": "{{table}}.url",
"flow_id": "{{table}}.flow_id",
"task_id": "{{table}}.task_id",
"subtask_id": "{{table}}.subtask_id",
"created_at": "{{table}}.created_at",
"data": "({{table}}.name || ' ' || {{table}}.url)",
}
type ScreenshotService struct {
db *gorm.DB
dataDir string
}
func NewScreenshotService(db *gorm.DB, dataDir string) *ScreenshotService {
return &ScreenshotService{
db: db,
dataDir: dataDir,
}
}
// GetScreenshots is a function to return screenshots list
// @Summary Retrieve screenshots list
// @Tags Screenshots
// @Produce json
// @Security BearerAuth
// @Param request query rdb.TableQuery true "query table params"
// @Success 200 {object} response.successResp{data=screenshots} "screenshots list received successful"
// @Failure 400 {object} response.errorResp "invalid query request data"
// @Failure 403 {object} response.errorResp "getting screenshots not permitted"
// @Failure 500 {object} response.errorResp "internal error on getting screenshots"
// @Router /screenshots/ [get]
func (s *ScreenshotService) GetScreenshots(c *gin.Context) {
var (
err error
query rdb.TableQuery
resp screenshots
)
if err = c.ShouldBindQuery(&query); err != nil {
logger.FromContext(c).WithError(err).Errorf("error binding query")
response.Error(c, response.ErrScreenshotsInvalidRequest, err)
return
}
uid := c.GetUint64("uid")
privs := c.GetStringSlice("prm")
var scope func(db *gorm.DB) *gorm.DB
if slices.Contains(privs, "screenshots.admin") {
scope = func(db *gorm.DB) *gorm.DB {
return db.
Joins("INNER JOIN flows f ON f.id = screenshots.flow_id")
}
} else if slices.Contains(privs, "screenshots.view") {
scope = func(db *gorm.DB) *gorm.DB {
return db.
Joins("INNER JOIN flows f ON f.id = screenshots.flow_id").
Where("f.user_id = ?", uid)
}
} else {
logger.FromContext(c).Errorf("error filtering user role permissions: permission not found")
response.Error(c, response.ErrNotPermitted, nil)
return
}
query.Init("screenshots", screenshotsSQLMappers)
if query.Group != "" {
if _, ok := screenshotsSQLMappers[query.Group]; !ok {
logger.FromContext(c).Errorf("error finding screenshots grouped: group field not found")
response.Error(c, response.ErrScreenshotsInvalidRequest, errors.New("group field not found"))
return
}
var respGrouped screenshotsGrouped
if respGrouped.Total, err = query.QueryGrouped(s.db, &respGrouped.Grouped, scope); err != nil {
logger.FromContext(c).WithError(err).Errorf("error finding screenshots grouped")
response.Error(c, response.ErrInternal, err)
return
}
response.Success(c, http.StatusOK, respGrouped)
return
}
if resp.Total, err = query.Query(s.db, &resp.Screenshots, scope); err != nil {
logger.FromContext(c).WithError(err).Errorf("error finding screenshots")
response.Error(c, response.ErrInternal, err)
return
}
for i := 0; i < len(resp.Screenshots); i++ {
if err = resp.Screenshots[i].Valid(); err != nil {
logger.FromContext(c).WithError(err).Errorf("error validating screenshot data '%d'", resp.Screenshots[i].ID)
response.Error(c, response.ErrScreenshotsInvalidData, err)
return
}
}
response.Success(c, http.StatusOK, resp)
}
// GetFlowScreenshots is a function to return screenshots list by flow id
// @Summary Retrieve screenshots list by flow id
// @Tags Screenshots
// @Produce json
// @Security BearerAuth
// @Param flowID path int true "flow id" minimum(0)
// @Param request query rdb.TableQuery true "query table params"
// @Success 200 {object} response.successResp{data=screenshots} "screenshots list received successful"
// @Failure 400 {object} response.errorResp "invalid query request data"
// @Failure 403 {object} response.errorResp "getting screenshots not permitted"
// @Failure 500 {object} response.errorResp "internal error on getting screenshots"
// @Router /flows/{flowID}/screenshots/ [get]
func (s *ScreenshotService) GetFlowScreenshots(c *gin.Context) {
var (
err error
flowID uint64
query rdb.TableQuery
resp screenshots
)
if flowID, err = strconv.ParseUint(c.Param("flowID"), 10, 64); err != nil {
logger.FromContext(c).WithError(err).Errorf("error parsing flow id")
response.Error(c, response.ErrScreenshotsInvalidRequest, err)
return
}
if err = c.ShouldBindQuery(&query); err != nil {
logger.FromContext(c).WithError(err).Errorf("error binding query")
response.Error(c, response.ErrScreenshotsInvalidRequest, err)
return
}
uid := c.GetUint64("uid")
privs := c.GetStringSlice("prm")
var scope func(db *gorm.DB) *gorm.DB
if slices.Contains(privs, "screenshots.admin") {
scope = func(db *gorm.DB) *gorm.DB {
return db.
Joins("INNER JOIN flows f ON f.id = screenshots.flow_id").
Where("f.id = ?", flowID)
}
} else if slices.Contains(privs, "screenshots.view") {
scope = func(db *gorm.DB) *gorm.DB {
return db.
Joins("INNER JOIN flows f ON f.id = screenshots.flow_id").
Where("f.id = ? AND f.user_id = ?", flowID, uid)
}
} else {
logger.FromContext(c).Errorf("error filtering user role permissions: permission not found")
response.Error(c, response.ErrNotPermitted, nil)
return
}
query.Init("screenshots", screenshotsSQLMappers)
if query.Group != "" {
if _, ok := screenshotsSQLMappers[query.Group]; !ok {
logger.FromContext(c).Errorf("error finding screenshots grouped: group field not found")
response.Error(c, response.ErrScreenshotsInvalidRequest, errors.New("group field not found"))
return
}
var respGrouped screenshotsGrouped
if respGrouped.Total, err = query.QueryGrouped(s.db, &respGrouped.Grouped, scope); err != nil {
logger.FromContext(c).WithError(err).Errorf("error finding screenshots grouped")
response.Error(c, response.ErrInternal, err)
return
}
response.Success(c, http.StatusOK, respGrouped)
return
}
if resp.Total, err = query.Query(s.db, &resp.Screenshots, scope); err != nil {
logger.FromContext(c).WithError(err).Errorf("error finding screenshots")
response.Error(c, response.ErrInternal, err)
return
}
for i := 0; i < len(resp.Screenshots); i++ {
if err = resp.Screenshots[i].Valid(); err != nil {
logger.FromContext(c).WithError(err).Errorf("error validating screenshot data '%d'", resp.Screenshots[i].ID)
response.Error(c, response.ErrScreenshotsInvalidData, err)
return
}
}
response.Success(c, http.StatusOK, resp)
}
// GetFlowScreenshot is a function to return screenshot info by id and flow id
// @Summary Retrieve screenshot info by id and flow id
// @Tags Screenshots
// @Produce json
// @Security BearerAuth
// @Param flowID path int true "flow id" minimum(0)
// @Param screenshotID path int true "screenshot id" minimum(0)
// @Success 200 {object} response.successResp{data=models.Screenshot} "screenshot info received successful"
// @Failure 403 {object} response.errorResp "getting screenshot not permitted"
// @Failure 404 {object} response.errorResp "screenshot not found"
// @Failure 500 {object} response.errorResp "internal error on getting screenshot"
// @Router /flows/{flowID}/screenshots/{screenshotID} [get]
func (s *ScreenshotService) GetFlowScreenshot(c *gin.Context) {
var (
err error
flowID uint64
screenshotID uint64
resp models.Screenshot
)
if flowID, err = strconv.ParseUint(c.Param("flowID"), 10, 64); err != nil {
logger.FromContext(c).WithError(err).Errorf("error parsing flow id")
response.Error(c, response.ErrScreenshotsInvalidRequest, err)
return
}
if screenshotID, err = strconv.ParseUint(c.Param("screenshotID"), 10, 64); err != nil {
logger.FromContext(c).WithError(err).Errorf("error parsing screenshot id")
response.Error(c, response.ErrScreenshotsInvalidRequest, err)
return
}
uid := c.GetUint64("uid")
privs := c.GetStringSlice("prm")
var scope func(db *gorm.DB) *gorm.DB
if slices.Contains(privs, "screenshots.admin") {
scope = func(db *gorm.DB) *gorm.DB {
return db.Where("f.id = ?", flowID)
}
} else if slices.Contains(privs, "screenshots.view") {
scope = func(db *gorm.DB) *gorm.DB {
return db.Where("f.id = ? AND f.user_id = ?", flowID, uid)
}
} else {
logger.FromContext(c).Errorf("error filtering user role permissions: permission not found")
response.Error(c, response.ErrNotPermitted, nil)
return
}
err = s.db.Model(&resp).
Joins("INNER JOIN flows f ON f.id = flow_id").
Scopes(scope).
Where("screenshots.id = ?", screenshotID).
Take(&resp).Error
if err != nil {
logger.FromContext(c).WithError(err).Errorf("error on getting screenshot by id")
if gorm.IsRecordNotFoundError(err) {
response.Error(c, response.ErrScreenshotsNotFound, err)
} else {
response.Error(c, response.ErrInternal, err)
}
return
}
response.Success(c, http.StatusOK, resp)
}
// GetFlowScreenshotFile is a function to return screenshot file by id and flow id
// @Summary Retrieve screenshot file by id and flow id
// @Tags Screenshots
// @Produce png,json
// @Security BearerAuth
// @Param flowID path int true "flow id" minimum(0)
// @Param screenshotID path int true "screenshot id" minimum(0)
// @Success 200 {file} file "screenshot file"
// @Failure 403 {object} response.errorResp "getting screenshot not permitted"
// @Failure 500 {object} response.errorResp "internal error on getting screenshot"
// @Router /flows/{flowID}/screenshots/{screenshotID}/file [get]
func (s *ScreenshotService) GetFlowScreenshotFile(c *gin.Context) {
var (
err error
flowID uint64
screenshotID uint64
resp models.Screenshot
)
if flowID, err = strconv.ParseUint(c.Param("flowID"), 10, 64); err != nil {
logger.FromContext(c).WithError(err).Errorf("error parsing flow id")
response.Error(c, response.ErrScreenshotsInvalidRequest, err)
return
}
if screenshotID, err = strconv.ParseUint(c.Param("screenshotID"), 10, 64); err != nil {
logger.FromContext(c).WithError(err).Errorf("error parsing screenshot id")
response.Error(c, response.ErrScreenshotsInvalidRequest, err)
return
}
uid := c.GetUint64("uid")
privs := c.GetStringSlice("prm")
var scope func(db *gorm.DB) *gorm.DB
if slices.Contains(privs, "screenshots.admin") {
scope = func(db *gorm.DB) *gorm.DB {
return db.Where("f.id = ?", flowID)
}
} else if slices.Contains(privs, "screenshots.download") {
scope = func(db *gorm.DB) *gorm.DB {
return db.Where("f.id = ? AND f.user_id = ?", flowID, uid)
}
} else {
logger.FromContext(c).Errorf("error filtering user role permissions: permission not found")
response.Error(c, response.ErrNotPermitted, nil)
return
}
err = s.db.Model(&resp).
Joins("INNER JOIN flows f ON f.id = flow_id").
Scopes(scope).
Where("screenshots.id = ?", screenshotID).
Take(&resp).Error
if err != nil {
logger.FromContext(c).WithError(err).Errorf("error on getting screenshot by id")
if gorm.IsRecordNotFoundError(err) {
response.Error(c, response.ErrScreenshotsNotFound, err)
} else {
response.Error(c, response.ErrInternal, err)
}
return
}
flowDirName := fmt.Sprintf("flow-%d", resp.FlowID)
c.File(filepath.Join(s.dataDir, "screenshots", flowDirName, resp.Name))
}
+221
View File
@@ -0,0 +1,221 @@
package services
import (
"errors"
"net/http"
"slices"
"strconv"
"pentagi/pkg/server/logger"
"pentagi/pkg/server/models"
"pentagi/pkg/server/rdb"
"pentagi/pkg/server/response"
"github.com/gin-gonic/gin"
"github.com/jinzhu/gorm"
)
type searchlogs struct {
SearchLogs []models.Searchlog `json:"searchlogs"`
Total uint64 `json:"total"`
}
type searchlogsGrouped struct {
Grouped []string `json:"grouped"`
Total uint64 `json:"total"`
}
var searchlogsSQLMappers = map[string]any{
"id": "{{table}}.id",
"initiator": "{{table}}.initiator",
"executor": "{{table}}.executor",
"engine": "{{table}}.engine",
"query": "{{table}}.query",
"result": "{{table}}.result",
"flow_id": "{{table}}.flow_id",
"task_id": "{{table}}.task_id",
"subtask_id": "{{table}}.subtask_id",
"created_at": "{{table}}.created_at",
"data": "({{table}}.query || ' ' || {{table}}.result)",
}
type SearchlogService struct {
db *gorm.DB
}
func NewSearchlogService(db *gorm.DB) *SearchlogService {
return &SearchlogService{
db: db,
}
}
// GetSearchlogs is a function to return searchlogs list
// @Summary Retrieve searchlogs list
// @Tags Searchlogs
// @Produce json
// @Security BearerAuth
// @Param request query rdb.TableQuery true "query table params"
// @Success 200 {object} response.successResp{data=searchlogs} "searchlogs list received successful"
// @Failure 400 {object} response.errorResp "invalid query request data"
// @Failure 403 {object} response.errorResp "getting searchlogs not permitted"
// @Failure 500 {object} response.errorResp "internal error on getting searchlogs"
// @Router /searchlogs/ [get]
func (s *SearchlogService) GetSearchlogs(c *gin.Context) {
var (
err error
query rdb.TableQuery
resp searchlogs
)
if err = c.ShouldBindQuery(&query); err != nil {
logger.FromContext(c).WithError(err).Errorf("error binding query")
response.Error(c, response.ErrSearchlogsInvalidRequest, err)
return
}
uid := c.GetUint64("uid")
privs := c.GetStringSlice("prm")
var scope func(db *gorm.DB) *gorm.DB
if slices.Contains(privs, "searchlogs.admin") {
scope = func(db *gorm.DB) *gorm.DB {
return db.
Joins("INNER JOIN flows f ON f.id = flow_id")
}
} else if slices.Contains(privs, "searchlogs.view") {
scope = func(db *gorm.DB) *gorm.DB {
return db.
Joins("INNER JOIN flows f ON f.id = flow_id").
Where("f.user_id = ?", uid)
}
} else {
logger.FromContext(c).Errorf("error filtering user role permissions: permission not found")
response.Error(c, response.ErrNotPermitted, nil)
return
}
query.Init("searchlogs", searchlogsSQLMappers)
if query.Group != "" {
if _, ok := searchlogsSQLMappers[query.Group]; !ok {
logger.FromContext(c).Errorf("error finding searchlogs grouped: group field not found")
response.Error(c, response.ErrSearchlogsInvalidRequest, errors.New("group field not found"))
return
}
var respGrouped searchlogsGrouped
if respGrouped.Total, err = query.QueryGrouped(s.db, &respGrouped.Grouped, scope); err != nil {
logger.FromContext(c).WithError(err).Errorf("error finding searchlogs grouped")
response.Error(c, response.ErrInternal, err)
return
}
response.Success(c, http.StatusOK, respGrouped)
return
}
if resp.Total, err = query.Query(s.db, &resp.SearchLogs, scope); err != nil {
logger.FromContext(c).WithError(err).Errorf("error finding searchlogs")
response.Error(c, response.ErrInternal, err)
return
}
for i := 0; i < len(resp.SearchLogs); i++ {
if err = resp.SearchLogs[i].Valid(); err != nil {
logger.FromContext(c).WithError(err).Errorf("error validating searchlog data '%d'", resp.SearchLogs[i].ID)
response.Error(c, response.ErrSearchlogsInvalidData, err)
return
}
}
response.Success(c, http.StatusOK, resp)
}
// GetFlowSearchlogs is a function to return searchlogs list by flow id
// @Summary Retrieve searchlogs list by flow id
// @Tags Searchlogs
// @Produce json
// @Security BearerAuth
// @Param flowID path int true "flow id" minimum(0)
// @Param request query rdb.TableQuery true "query table params"
// @Success 200 {object} response.successResp{data=searchlogs} "searchlogs list received successful"
// @Failure 400 {object} response.errorResp "invalid query request data"
// @Failure 403 {object} response.errorResp "getting searchlogs not permitted"
// @Failure 500 {object} response.errorResp "internal error on getting searchlogs"
// @Router /flows/{flowID}/searchlogs/ [get]
func (s *SearchlogService) GetFlowSearchlogs(c *gin.Context) {
var (
err error
flowID uint64
query rdb.TableQuery
resp searchlogs
)
if flowID, err = strconv.ParseUint(c.Param("flowID"), 10, 64); err != nil {
logger.FromContext(c).WithError(err).Errorf("error parsing flow id")
response.Error(c, response.ErrSearchlogsInvalidRequest, err)
return
}
if err = c.ShouldBindQuery(&query); err != nil {
logger.FromContext(c).WithError(err).Errorf("error binding query")
response.Error(c, response.ErrSearchlogsInvalidRequest, err)
return
}
uid := c.GetUint64("uid")
privs := c.GetStringSlice("prm")
var scope func(db *gorm.DB) *gorm.DB
if slices.Contains(privs, "searchlogs.admin") {
scope = func(db *gorm.DB) *gorm.DB {
return db.
Joins("INNER JOIN flows f ON f.id = flow_id").
Where("f.id = ?", flowID)
}
} else if slices.Contains(privs, "searchlogs.view") {
scope = func(db *gorm.DB) *gorm.DB {
return db.
Joins("INNER JOIN flows f ON f.id = flow_id").
Where("f.id = ? AND f.user_id = ?", flowID, uid)
}
} else {
logger.FromContext(c).Errorf("error filtering user role permissions: permission not found")
response.Error(c, response.ErrNotPermitted, nil)
return
}
query.Init("searchlogs", searchlogsSQLMappers)
if query.Group != "" {
if _, ok := searchlogsSQLMappers[query.Group]; !ok {
logger.FromContext(c).Errorf("error finding searchlogs grouped: group field not found")
response.Error(c, response.ErrSearchlogsInvalidRequest, errors.New("group field not found"))
return
}
var respGrouped searchlogsGrouped
if respGrouped.Total, err = query.QueryGrouped(s.db, &respGrouped.Grouped, scope); err != nil {
logger.FromContext(c).WithError(err).Errorf("error finding searchlogs grouped")
response.Error(c, response.ErrInternal, err)
return
}
response.Success(c, http.StatusOK, respGrouped)
return
}
if resp.Total, err = query.Query(s.db, &resp.SearchLogs, scope); err != nil {
logger.FromContext(c).WithError(err).Errorf("error finding searchlogs")
response.Error(c, response.ErrInternal, err)
return
}
for i := 0; i < len(resp.SearchLogs); i++ {
if err = resp.SearchLogs[i].Valid(); err != nil {
logger.FromContext(c).WithError(err).Errorf("error validating searchlog data '%d'", resp.SearchLogs[i].ID)
response.Error(c, response.ErrSearchlogsInvalidData, err)
return
}
}
response.Success(c, http.StatusOK, resp)
}
+50
View File
@@ -0,0 +1,50 @@
package services
import (
"net/http"
"slices"
"pentagi/pkg/config"
"pentagi/pkg/server/logger"
"pentagi/pkg/server/models"
"pentagi/pkg/server/response"
"pentagi/pkg/version"
"github.com/gin-gonic/gin"
)
type SettingsService struct {
cfg *config.Config
}
func NewSettingsService(cfg *config.Config) *SettingsService {
return &SettingsService{cfg: cfg}
}
// GetSettings is a function to return settings
// @Summary Retrieve settings
// @Tags Settings
// @Produce json
// @Security BearerAuth
// @Success 200 {object} response.successResp{data=models.Settings} "settings received successful"
// @Failure 403 {object} response.errorResp "getting settings not permitted"
// @Router /settings/ [get]
func (s *SettingsService) GetSettings(c *gin.Context) {
privs := c.GetStringSlice("prm")
if !slices.Contains(privs, "settings.view") {
logger.FromContext(c).Errorf("error filtering user role permissions: permission not found")
response.Error(c, response.ErrNotPermitted, nil)
return
}
settings := models.Settings{
Debug: s.cfg.Debug,
AskUser: s.cfg.AskUser,
Version: version.GetBinaryVersion(),
DockerInside: s.cfg.DockerInside,
IsDevelopMode: version.IsDevelopMode(),
AssistantUseAgents: s.cfg.AssistantUseAgents,
}
response.Success(c, http.StatusOK, settings)
}
+321
View File
@@ -0,0 +1,321 @@
package services
import (
"errors"
"net/http"
"slices"
"strconv"
"pentagi/pkg/server/logger"
"pentagi/pkg/server/models"
"pentagi/pkg/server/rdb"
"pentagi/pkg/server/response"
"github.com/gin-gonic/gin"
"github.com/jinzhu/gorm"
)
type subtasks struct {
Subtasks []models.Subtask `json:"subtasks"`
Total uint64 `json:"total"`
}
type subtasksGrouped struct {
Grouped []string `json:"grouped"`
Total uint64 `json:"total"`
}
var subtasksSQLMappers = map[string]any{
"id": "{{table}}.id",
"status": "{{table}}.status",
"title": "{{table}}.title",
"description": "{{table}}.description",
"context": "{{table}}.context",
"result": "{{table}}.result",
"task_id": "{{table}}.task_id",
"created_at": "{{table}}.created_at",
"updated_at": "{{table}}.updated_at",
"data": "({{table}}.status || ' ' || {{table}}.title || ' ' || {{table}}.description || ' ' || {{table}}.context || ' ' || {{table}}.result)",
}
type SubtaskService struct {
db *gorm.DB
}
func NewSubtaskService(db *gorm.DB) *SubtaskService {
return &SubtaskService{
db: db,
}
}
// GetFlowSubtasks is a function to return flow subtasks list
// @Summary Retrieve flow subtasks list
// @Tags Subtasks
// @Produce json
// @Security BearerAuth
// @Param flowID path int true "flow id" minimum(0)
// @Param request query rdb.TableQuery true "query table params"
// @Success 200 {object} response.successResp{data=subtasks} "flow subtasks list received successful"
// @Failure 400 {object} response.errorResp "invalid query request data"
// @Failure 403 {object} response.errorResp "getting flow subtasks not permitted"
// @Failure 500 {object} response.errorResp "internal error on getting flow subtasks"
// @Router /flows/{flowID}/subtasks/ [get]
func (s *SubtaskService) GetFlowSubtasks(c *gin.Context) {
var (
err error
flowID uint64
query rdb.TableQuery
resp subtasks
)
if err = c.ShouldBindQuery(&query); err != nil {
logger.FromContext(c).WithError(err).Errorf("error binding query")
response.Error(c, response.ErrSubtasksInvalidRequest, err)
return
}
if flowID, err = strconv.ParseUint(c.Param("flowID"), 10, 64); err != nil {
logger.FromContext(c).WithError(err).Errorf("error parsing flow id")
response.Error(c, response.ErrSubtasksInvalidRequest, err)
return
}
uid := c.GetUint64("uid")
privs := c.GetStringSlice("prm")
var scope func(db *gorm.DB) *gorm.DB
if slices.Contains(privs, "subtasks.admin") {
scope = func(db *gorm.DB) *gorm.DB {
return db.
Joins("INNER JOIN tasks t ON t.id = subtasks.task_id").
Joins("INNER JOIN flows f ON f.id = t.flow_id").
Where("f.id = ?", flowID)
}
} else if slices.Contains(privs, "subtasks.view") {
scope = func(db *gorm.DB) *gorm.DB {
return db.
Joins("INNER JOIN tasks t ON t.id = subtasks.task_id").
Joins("INNER JOIN flows f ON f.id = t.flow_id").
Where("f.id = ? AND f.user_id = ?", flowID, uid)
}
} else {
logger.FromContext(c).Errorf("error filtering user role permissions: permission not found")
response.Error(c, response.ErrNotPermitted, nil)
return
}
query.Init("subtasks", subtasksSQLMappers)
if query.Group != "" {
if _, ok := subtasksSQLMappers[query.Group]; !ok {
logger.FromContext(c).Errorf("error finding subtasks grouped: group field not found")
response.Error(c, response.ErrSubtasksInvalidRequest, errors.New("group field not found"))
return
}
var respGrouped subtasksGrouped
if respGrouped.Total, err = query.QueryGrouped(s.db, &respGrouped.Grouped, scope); err != nil {
logger.FromContext(c).WithError(err).Errorf("error finding subtasks grouped")
response.Error(c, response.ErrInternal, err)
return
}
response.Success(c, http.StatusOK, respGrouped)
return
}
if resp.Total, err = query.Query(s.db, &resp.Subtasks, scope); err != nil {
logger.FromContext(c).WithError(err).Errorf("error finding subtasks")
response.Error(c, response.ErrInternal, err)
return
}
for i := 0; i < len(resp.Subtasks); i++ {
if err = resp.Subtasks[i].Valid(); err != nil {
logger.FromContext(c).WithError(err).Errorf("error validating subtask data '%d'", resp.Subtasks[i].ID)
response.Error(c, response.ErrSubtasksInvalidData, err)
return
}
}
response.Success(c, http.StatusOK, resp)
}
// GetFlowTaskSubtasks is a function to return flow task subtasks list
// @Summary Retrieve flow task subtasks list
// @Tags Subtasks
// @Produce json
// @Security BearerAuth
// @Param flowID path int true "flow id" minimum(0)
// @Param taskID path int true "task id" minimum(0)
// @Param request query rdb.TableQuery true "query table params"
// @Success 200 {object} response.successResp{data=subtasks} "flow task subtasks list received successful"
// @Failure 400 {object} response.errorResp "invalid query request data"
// @Failure 403 {object} response.errorResp "getting flow task subtasks not permitted"
// @Failure 500 {object} response.errorResp "internal error on getting flow subtasks"
// @Router /flows/{flowID}/tasks/{taskID}/subtasks/ [get]
func (s *SubtaskService) GetFlowTaskSubtasks(c *gin.Context) {
var (
err error
flowID uint64
taskID uint64
query rdb.TableQuery
resp subtasks
)
if err = c.ShouldBindQuery(&query); err != nil {
logger.FromContext(c).WithError(err).Errorf("error binding query")
response.Error(c, response.ErrSubtasksInvalidRequest, err)
return
}
if flowID, err = strconv.ParseUint(c.Param("flowID"), 10, 64); err != nil {
logger.FromContext(c).WithError(err).Errorf("error parsing flow id")
response.Error(c, response.ErrSubtasksInvalidRequest, err)
return
}
if taskID, err = strconv.ParseUint(c.Param("taskID"), 10, 64); err != nil {
logger.FromContext(c).WithError(err).Errorf("error parsing task id")
response.Error(c, response.ErrSubtasksInvalidRequest, err)
return
}
uid := c.GetUint64("uid")
privs := c.GetStringSlice("prm")
var scope func(db *gorm.DB) *gorm.DB
if slices.Contains(privs, "subtasks.admin") {
scope = func(db *gorm.DB) *gorm.DB {
return db.
Joins("INNER JOIN tasks t ON t.id = subtasks.task_id").
Joins("INNER JOIN flows f ON f.id = t.flow_id").
Where("f.id = ? AND t.id = ?", flowID, taskID)
}
} else if slices.Contains(privs, "subtasks.view") {
scope = func(db *gorm.DB) *gorm.DB {
return db.
Joins("INNER JOIN tasks t ON t.id = subtasks.task_id").
Joins("INNER JOIN flows f ON f.id = t.flow_id").
Where("f.id = ? AND f.user_id = ? AND t.id = ?", flowID, uid, taskID)
}
} else {
logger.FromContext(c).Errorf("error filtering user role permissions: permission not found")
response.Error(c, response.ErrNotPermitted, nil)
return
}
query.Init("subtasks", subtasksSQLMappers)
if query.Group != "" {
if _, ok := subtasksSQLMappers[query.Group]; !ok {
logger.FromContext(c).Errorf("error finding subtasks grouped: group field not found")
response.Error(c, response.ErrSubtasksInvalidRequest, errors.New("group field not found"))
return
}
var respGrouped subtasksGrouped
if respGrouped.Total, err = query.QueryGrouped(s.db, &respGrouped.Grouped, scope); err != nil {
logger.FromContext(c).WithError(err).Errorf("error finding subtasks grouped")
response.Error(c, response.ErrInternal, err)
return
}
response.Success(c, http.StatusOK, respGrouped)
return
}
if resp.Total, err = query.Query(s.db, &resp.Subtasks, scope); err != nil {
logger.FromContext(c).WithError(err).Errorf("error finding subtasks")
response.Error(c, response.ErrInternal, err)
return
}
for i := 0; i < len(resp.Subtasks); i++ {
if err = resp.Subtasks[i].Valid(); err != nil {
logger.FromContext(c).WithError(err).Errorf("error validating subtask data '%d'", resp.Subtasks[i].ID)
response.Error(c, response.ErrSubtasksInvalidData, err)
return
}
}
response.Success(c, http.StatusOK, resp)
}
// GetFlowTaskSubtask is a function to return flow task subtask by id
// @Summary Retrieve flow task subtask by id
// @Tags Subtasks
// @Produce json
// @Security BearerAuth
// @Param flowID path int true "flow id" minimum(0)
// @Param taskID path int true "task id" minimum(0)
// @Param subtaskID path int true "subtask id" minimum(0)
// @Success 200 {object} response.successResp{data=models.Subtask} "flow task subtask received successful"
// @Failure 403 {object} response.errorResp "getting flow task subtask not permitted"
// @Failure 404 {object} response.errorResp "flow task subtask not found"
// @Failure 500 {object} response.errorResp "internal error on getting flow task subtask"
// @Router /flows/{flowID}/tasks/{taskID}/subtasks/{subtaskID} [get]
func (s *SubtaskService) GetFlowTaskSubtask(c *gin.Context) {
var (
err error
flowID uint64
taskID uint64
subtaskID uint64
resp models.Subtask
)
if flowID, err = strconv.ParseUint(c.Param("flowID"), 10, 64); err != nil {
logger.FromContext(c).WithError(err).Errorf("error parsing flow id")
response.Error(c, response.ErrSubtasksInvalidRequest, err)
return
}
if taskID, err = strconv.ParseUint(c.Param("taskID"), 10, 64); err != nil {
logger.FromContext(c).WithError(err).Errorf("error parsing task id")
response.Error(c, response.ErrSubtasksInvalidRequest, err)
return
}
if subtaskID, err = strconv.ParseUint(c.Param("subtaskID"), 10, 64); err != nil {
logger.FromContext(c).WithError(err).Errorf("error parsing subtask id")
response.Error(c, response.ErrSubtasksInvalidRequest, err)
return
}
uid := c.GetUint64("uid")
privs := c.GetStringSlice("prm")
var scope func(db *gorm.DB) *gorm.DB
if slices.Contains(privs, "subtasks.admin") {
scope = func(db *gorm.DB) *gorm.DB {
return db.
Joins("INNER JOIN tasks t ON t.id = subtasks.task_id").
Joins("INNER JOIN flows f ON f.id = t.flow_id").
Where("f.id = ? AND t.id = ?", flowID, taskID)
}
} else if slices.Contains(privs, "subtasks.view") {
scope = func(db *gorm.DB) *gorm.DB {
return db.
Joins("INNER JOIN tasks t ON t.id = subtasks.task_id").
Joins("INNER JOIN flows f ON f.id = t.flow_id").
Where("f.id = ? AND f.user_id = ? AND t.id = ?", flowID, uid, taskID)
}
} else {
logger.FromContext(c).Errorf("error filtering user role permissions: permission not found")
response.Error(c, response.ErrNotPermitted, nil)
return
}
err = s.db.Model(&resp).
Scopes(scope).
Where("subtasks.id = ?", subtaskID).
Take(&resp).Error
if err != nil {
logger.FromContext(c).WithError(err).Errorf("error on getting flow task subtask by id")
if gorm.IsRecordNotFoundError(err) {
response.Error(c, response.ErrSubtasksNotFound, err)
} else {
response.Error(c, response.ErrInternal, err)
}
return
}
response.Success(c, http.StatusOK, resp)
}
+307
View File
@@ -0,0 +1,307 @@
package services
import (
"errors"
"net/http"
"slices"
"strconv"
"pentagi/pkg/server/logger"
"pentagi/pkg/server/models"
"pentagi/pkg/server/rdb"
"pentagi/pkg/server/response"
"github.com/gin-gonic/gin"
"github.com/jinzhu/gorm"
)
type tasks struct {
Tasks []models.Task `json:"tasks"`
Total uint64 `json:"total"`
}
type tasksGrouped struct {
Grouped []string `json:"grouped"`
Total uint64 `json:"total"`
}
var tasksSQLMappers = map[string]any{
"id": "{{table}}.id",
"status": "{{table}}.status",
"title": "{{table}}.title",
"input": "{{table}}.input",
"result": "{{table}}.result",
"flow_id": "{{table}}.flow_id",
"created_at": "{{table}}.created_at",
"updated_at": "{{table}}.updated_at",
"data": "({{table}}.status || ' ' || {{table}}.title || ' ' || {{table}}.input || ' ' || {{table}}.result)",
}
type TaskService struct {
db *gorm.DB
}
func NewTaskService(db *gorm.DB) *TaskService {
return &TaskService{
db: db,
}
}
// GetFlowTasks is a function to return flow tasks list
// @Summary Retrieve flow tasks list
// @Tags Tasks
// @Produce json
// @Security BearerAuth
// @Param flowID path int true "flow id" minimum(0)
// @Param request query rdb.TableQuery true "query table params"
// @Success 200 {object} response.successResp{data=tasks} "flow tasks list received successful"
// @Failure 400 {object} response.errorResp "invalid query request data"
// @Failure 403 {object} response.errorResp "getting flow tasks not permitted"
// @Failure 500 {object} response.errorResp "internal error on getting flow tasks"
// @Router /flows/{flowID}/tasks/ [get]
func (s *TaskService) GetFlowTasks(c *gin.Context) {
var (
err error
flowID uint64
query rdb.TableQuery
resp tasks
)
if err = c.ShouldBindQuery(&query); err != nil {
logger.FromContext(c).WithError(err).Errorf("error binding query")
response.Error(c, response.ErrTasksInvalidRequest, err)
return
}
if flowID, err = strconv.ParseUint(c.Param("flowID"), 10, 64); err != nil {
logger.FromContext(c).WithError(err).Errorf("error parsing flow id")
response.Error(c, response.ErrTasksInvalidRequest, err)
return
}
uid := c.GetUint64("uid")
privs := c.GetStringSlice("prm")
var scope func(db *gorm.DB) *gorm.DB
if slices.Contains(privs, "tasks.admin") {
scope = func(db *gorm.DB) *gorm.DB {
return db.
Joins("INNER JOIN flows f ON f.id = tasks.flow_id").
Where("f.id = ?", flowID)
}
} else if slices.Contains(privs, "tasks.view") {
scope = func(db *gorm.DB) *gorm.DB {
return db.
Joins("INNER JOIN flows f ON f.id = tasks.flow_id").
Where("f.id = ? AND f.user_id = ?", flowID, uid)
}
} else {
logger.FromContext(c).Errorf("error filtering user role permissions: permission not found")
response.Error(c, response.ErrNotPermitted, nil)
return
}
query.Init("tasks", tasksSQLMappers)
if query.Group != "" {
if _, ok := tasksSQLMappers[query.Group]; !ok {
logger.FromContext(c).Errorf("error finding tasks grouped: group field not found")
response.Error(c, response.ErrTasksInvalidRequest, errors.New("group field not found"))
return
}
var respGrouped tasksGrouped
if respGrouped.Total, err = query.QueryGrouped(s.db, &respGrouped.Grouped, scope); err != nil {
logger.FromContext(c).WithError(err).Errorf("error finding tasks grouped")
response.Error(c, response.ErrInternal, err)
return
}
response.Success(c, http.StatusOK, respGrouped)
return
}
if resp.Total, err = query.Query(s.db, &resp.Tasks, scope); err != nil {
logger.FromContext(c).WithError(err).Errorf("error finding tasks")
response.Error(c, response.ErrInternal, err)
return
}
for i := 0; i < len(resp.Tasks); i++ {
if err = resp.Tasks[i].Valid(); err != nil {
logger.FromContext(c).WithError(err).Errorf("error validating task data '%d'", resp.Tasks[i].ID)
response.Error(c, response.ErrTasksInvalidData, err)
return
}
}
response.Success(c, http.StatusOK, resp)
}
// GetFlowTask is a function to return flow task by id
// @Summary Retrieve flow task by id
// @Tags Tasks
// @Produce json
// @Security BearerAuth
// @Param flowID path int true "flow id" minimum(0)
// @Param taskID path int true "task id" minimum(0)
// @Success 200 {object} response.successResp{data=models.Task} "flow task received successful"
// @Failure 403 {object} response.errorResp "getting flow task not permitted"
// @Failure 404 {object} response.errorResp "flow task not found"
// @Failure 500 {object} response.errorResp "internal error on getting flow task"
// @Router /flows/{flowID}/tasks/{taskID} [get]
func (s *TaskService) GetFlowTask(c *gin.Context) {
var (
err error
flowID uint64
taskID uint64
resp models.Task
)
if flowID, err = strconv.ParseUint(c.Param("flowID"), 10, 64); err != nil {
logger.FromContext(c).WithError(err).Errorf("error parsing flow id")
response.Error(c, response.ErrTasksInvalidRequest, err)
return
}
if taskID, err = strconv.ParseUint(c.Param("taskID"), 10, 64); err != nil {
logger.FromContext(c).WithError(err).Errorf("error parsing task id")
response.Error(c, response.ErrTasksInvalidRequest, err)
return
}
uid := c.GetUint64("uid")
privs := c.GetStringSlice("prm")
var scope func(db *gorm.DB) *gorm.DB
if slices.Contains(privs, "tasks.admin") {
scope = func(db *gorm.DB) *gorm.DB {
return db.
Joins("INNER JOIN flows f ON f.id = tasks.flow_id").
Where("f.id = ?", flowID)
}
} else if slices.Contains(privs, "tasks.view") {
scope = func(db *gorm.DB) *gorm.DB {
return db.
Joins("INNER JOIN flows f ON f.id = tasks.flow_id").
Where("f.id = ? AND f.user_id = ?", flowID, uid)
}
} else {
logger.FromContext(c).Errorf("error filtering user role permissions: permission not found")
response.Error(c, response.ErrNotPermitted, nil)
return
}
err = s.db.Model(&resp).
Scopes(scope).
Where("tasks.id = ?", taskID).
Take(&resp).Error
if err != nil {
logger.FromContext(c).WithError(err).Errorf("error on getting flow task by id")
if gorm.IsRecordNotFoundError(err) {
response.Error(c, response.ErrTasksNotFound, err)
} else {
response.Error(c, response.ErrInternal, err)
}
return
}
response.Success(c, http.StatusOK, resp)
}
// GetFlowTaskGraph is a function to return flow task graph by id
// @Summary Retrieve flow task graph by id
// @Tags Tasks
// @Produce json
// @Security BearerAuth
// @Param flowID path int true "flow id" minimum(0)
// @Param taskID path int true "task id" minimum(0)
// @Success 200 {object} response.successResp{data=models.FlowTasksSubtasks} "flow task graph received successful"
// @Failure 403 {object} response.errorResp "getting flow task graph not permitted"
// @Failure 404 {object} response.errorResp "flow task graph not found"
// @Failure 500 {object} response.errorResp "internal error on getting flow task graph"
// @Router /flows/{flowID}/tasks/{taskID}/graph [get]
func (s *TaskService) GetFlowTaskGraph(c *gin.Context) {
var (
err error
flow models.Flow
flowID uint64
taskID uint64
resp models.TaskSubtasks
)
if flowID, err = strconv.ParseUint(c.Param("flowID"), 10, 64); err != nil {
logger.FromContext(c).WithError(err).Errorf("error parsing flow id")
response.Error(c, response.ErrTasksInvalidRequest, err)
return
}
if taskID, err = strconv.ParseUint(c.Param("taskID"), 10, 64); err != nil {
logger.FromContext(c).WithError(err).Errorf("error parsing task id")
response.Error(c, response.ErrTasksInvalidRequest, err)
return
}
uid := c.GetUint64("uid")
privs := c.GetStringSlice("prm")
var scope func(db *gorm.DB) *gorm.DB
if slices.Contains(privs, "tasks.admin") {
scope = func(db *gorm.DB) *gorm.DB {
return db.Where("f.id = ?", flowID)
}
} else if slices.Contains(privs, "tasks.view") {
scope = func(db *gorm.DB) *gorm.DB {
return db.Where("f.id = ? AND f.user_id = ?", flowID, uid)
}
} else {
logger.FromContext(c).Errorf("error filtering user role permissions: permission not found")
response.Error(c, response.ErrNotPermitted, nil)
return
}
err = s.db.Model(&resp).
Joins("INNER JOIN flows f ON f.id = tasks.flow_id").
Scopes(scope).
Where("tasks.id = ?", taskID).
Take(&resp).Error
if err != nil {
logger.FromContext(c).WithError(err).Errorf("error on getting flow task by id")
if gorm.IsRecordNotFoundError(err) {
response.Error(c, response.ErrTasksNotFound, err)
} else {
response.Error(c, response.ErrInternal, err)
}
return
}
err = s.db.Where("id = ?", resp.FlowID).Take(&flow).Error
if err != nil {
logger.FromContext(c).WithError(err).Errorf("error on getting flow by id")
if gorm.IsRecordNotFoundError(err) {
response.Error(c, response.ErrTasksNotFound, err)
} else {
response.Error(c, response.ErrInternal, err)
}
return
}
isSubtasksAdmin := slices.Contains(privs, "subtasks.admin")
isSubtasksView := slices.Contains(privs, "subtasks.view")
if !(flow.UserID == uid && isSubtasksView) && !(flow.UserID != uid && isSubtasksAdmin) {
response.Success(c, http.StatusOK, resp)
return
}
err = s.db.Model(&resp).Association("subtasks").Find(&resp.Subtasks).Error
if err != nil {
logger.FromContext(c).WithError(err).Errorf("error on getting task subtasks")
response.Error(c, response.ErrInternal, err)
return
}
if err = resp.Valid(); err != nil {
logger.FromContext(c).WithError(err).Errorf("error validating task data '%d'", taskID)
response.Error(c, response.ErrTasksInvalidData, err)
return
}
response.Success(c, http.StatusOK, resp)
}
+219
View File
@@ -0,0 +1,219 @@
package services
import (
"errors"
"net/http"
"slices"
"strconv"
"pentagi/pkg/server/logger"
"pentagi/pkg/server/models"
"pentagi/pkg/server/rdb"
"pentagi/pkg/server/response"
"github.com/gin-gonic/gin"
"github.com/jinzhu/gorm"
)
type termlogs struct {
TermLogs []models.Termlog `json:"termlogs"`
Total uint64 `json:"total"`
}
type termlogsGrouped struct {
Grouped []string `json:"grouped"`
Total uint64 `json:"total"`
}
var termlogsSQLMappers = map[string]any{
"id": "{{table}}.id",
"type": "{{table}}.type",
"text": "{{table}}.text",
"container_id": "{{table}}.container_id",
"flow_id": "{{table}}.flow_id",
"task_id": "{{table}}.task_id",
"subtask_id": "{{table}}.subtask_id",
"created_at": "{{table}}.created_at",
"data": "({{table}}.type || ' ' || {{table}}.text)",
}
type TermlogService struct {
db *gorm.DB
}
func NewTermlogService(db *gorm.DB) *TermlogService {
return &TermlogService{
db: db,
}
}
// GetTermlogs is a function to return termlogs list
// @Summary Retrieve termlogs list
// @Tags Termlogs
// @Produce json
// @Security BearerAuth
// @Param request query rdb.TableQuery true "query table params"
// @Success 200 {object} response.successResp{data=termlogs} "termlogs list received successful"
// @Failure 400 {object} response.errorResp "invalid query request data"
// @Failure 403 {object} response.errorResp "getting termlogs not permitted"
// @Failure 500 {object} response.errorResp "internal error on getting termlogs"
// @Router /termlogs/ [get]
func (s *TermlogService) GetTermlogs(c *gin.Context) {
var (
err error
query rdb.TableQuery
resp termlogs
)
if err = c.ShouldBindQuery(&query); err != nil {
logger.FromContext(c).WithError(err).Errorf("error binding query")
response.Error(c, response.ErrTermlogsInvalidRequest, err)
return
}
uid := c.GetUint64("uid")
privs := c.GetStringSlice("prm")
var scope func(db *gorm.DB) *gorm.DB
if slices.Contains(privs, "termlogs.admin") {
scope = func(db *gorm.DB) *gorm.DB {
return db.
Joins("INNER JOIN flows f ON f.id = flow_id")
}
} else if slices.Contains(privs, "termlogs.view") {
scope = func(db *gorm.DB) *gorm.DB {
return db.
Joins("INNER JOIN flows f ON f.id = flow_id").
Where("f.user_id = ?", uid)
}
} else {
logger.FromContext(c).Errorf("error filtering user role permissions: permission not found")
response.Error(c, response.ErrNotPermitted, nil)
return
}
query.Init("termlogs", termlogsSQLMappers)
if query.Group != "" {
if _, ok := termlogsSQLMappers[query.Group]; !ok {
logger.FromContext(c).Errorf("error finding termlogs grouped: group field not found")
response.Error(c, response.ErrTermlogsInvalidRequest, errors.New("group field not found"))
return
}
var respGrouped termlogsGrouped
if respGrouped.Total, err = query.QueryGrouped(s.db, &respGrouped.Grouped, scope); err != nil {
logger.FromContext(c).WithError(err).Errorf("error finding termlogs grouped")
response.Error(c, response.ErrInternal, err)
return
}
response.Success(c, http.StatusOK, respGrouped)
return
}
if resp.Total, err = query.Query(s.db, &resp.TermLogs, scope); err != nil {
logger.FromContext(c).WithError(err).Errorf("error finding termlogs")
response.Error(c, response.ErrInternal, err)
return
}
for i := 0; i < len(resp.TermLogs); i++ {
if err = resp.TermLogs[i].Valid(); err != nil {
logger.FromContext(c).WithError(err).Errorf("error validating termlog data '%d'", resp.TermLogs[i].ID)
response.Error(c, response.ErrTermlogsInvalidData, err)
return
}
}
response.Success(c, http.StatusOK, resp)
}
// GetFlowTermlogs is a function to return termlogs list by flow id
// @Summary Retrieve termlogs list by flow id
// @Tags Termlogs
// @Produce json
// @Security BearerAuth
// @Param flowID path int true "flow id" minimum(0)
// @Param request query rdb.TableQuery true "query table params"
// @Success 200 {object} response.successResp{data=termlogs} "termlogs list received successful"
// @Failure 400 {object} response.errorResp "invalid query request data"
// @Failure 403 {object} response.errorResp "getting termlogs not permitted"
// @Failure 500 {object} response.errorResp "internal error on getting termlogs"
// @Router /flows/{flowID}/termlogs/ [get]
func (s *TermlogService) GetFlowTermlogs(c *gin.Context) {
var (
err error
flowID uint64
query rdb.TableQuery
resp termlogs
)
if flowID, err = strconv.ParseUint(c.Param("flowID"), 10, 64); err != nil {
logger.FromContext(c).WithError(err).Errorf("error parsing flow id")
response.Error(c, response.ErrTermlogsInvalidRequest, err)
return
}
if err = c.ShouldBindQuery(&query); err != nil {
logger.FromContext(c).WithError(err).Errorf("error binding query")
response.Error(c, response.ErrTermlogsInvalidRequest, err)
return
}
uid := c.GetUint64("uid")
privs := c.GetStringSlice("prm")
var scope func(db *gorm.DB) *gorm.DB
if slices.Contains(privs, "termlogs.admin") {
scope = func(db *gorm.DB) *gorm.DB {
return db.
Joins("INNER JOIN flows f ON f.id = flow_id").
Where("f.id = ?", flowID)
}
} else if slices.Contains(privs, "termlogs.view") {
scope = func(db *gorm.DB) *gorm.DB {
return db.
Joins("INNER JOIN flows f ON f.id = flow_id").
Where("f.id = ? AND f.user_id = ?", flowID, uid)
}
} else {
logger.FromContext(c).Errorf("error filtering user role permissions: permission not found")
response.Error(c, response.ErrNotPermitted, nil)
return
}
query.Init("termlogs", termlogsSQLMappers)
if query.Group != "" {
if _, ok := termlogsSQLMappers[query.Group]; !ok {
logger.FromContext(c).Errorf("error finding termlogs grouped: group field not found")
response.Error(c, response.ErrTermlogsInvalidRequest, errors.New("group field not found"))
return
}
var respGrouped termlogsGrouped
if respGrouped.Total, err = query.QueryGrouped(s.db, &respGrouped.Grouped, scope); err != nil {
logger.FromContext(c).WithError(err).Errorf("error finding termlogs grouped")
response.Error(c, response.ErrInternal, err)
return
}
response.Success(c, http.StatusOK, respGrouped)
return
}
if resp.Total, err = query.Query(s.db, &resp.TermLogs, scope); err != nil {
logger.FromContext(c).WithError(err).Errorf("error finding termlogs")
response.Error(c, response.ErrInternal, err)
return
}
for i := 0; i < len(resp.TermLogs); i++ {
if err = resp.TermLogs[i].Valid(); err != nil {
logger.FromContext(c).WithError(err).Errorf("error validating termlog data '%d'", resp.TermLogs[i].ID)
response.Error(c, response.ErrTermlogsInvalidData, err)
return
}
}
response.Success(c, http.StatusOK, resp)
}
+290
View File
@@ -0,0 +1,290 @@
package services
import (
"errors"
"net/http"
"slices"
"strconv"
"pentagi/pkg/server/logger"
"pentagi/pkg/server/models"
"pentagi/pkg/server/rdb"
"pentagi/pkg/server/response"
"github.com/gin-gonic/gin"
"github.com/jinzhu/gorm"
)
type toolcalls struct {
Toolcalls []models.Toolcall `json:"toolcalls"`
Total uint64 `json:"total"`
}
type toolcallsGrouped struct {
Grouped []string `json:"grouped"`
Total uint64 `json:"total"`
}
var toolcallsSQLMappers = map[string]any{
"id": "{{table}}.id",
"call_id": "{{table}}.call_id",
"status": "{{table}}.status",
"name": "{{table}}.name",
"args": "{{table}}.args",
"result": "{{table}}.result",
"duration_seconds": "{{table}}.duration_seconds",
"flow_id": "{{table}}.flow_id",
"task_id": "{{table}}.task_id",
"subtask_id": "{{table}}.subtask_id",
"created_at": "{{table}}.created_at",
"updated_at": "{{table}}.updated_at",
"data": "({{table}}.name || ' ' || {{table}}.call_id || ' ' || {{table}}.args::text || ' ' || {{table}}.result)",
}
type ToolcallService struct {
db *gorm.DB
}
func NewToolcallService(db *gorm.DB) *ToolcallService {
return &ToolcallService{
db: db,
}
}
// GetToolcalls is a function to return toolcalls list
// @Summary Retrieve toolcalls list
// @Tags Toolcalls
// @Produce json
// @Security BearerAuth
// @Param request query rdb.TableQuery true "query table params"
// @Success 200 {object} response.successResp{data=toolcalls} "toolcalls list received successful"
// @Failure 400 {object} response.errorResp "invalid query request data"
// @Failure 403 {object} response.errorResp "getting toolcalls not permitted"
// @Failure 500 {object} response.errorResp "internal error on getting toolcalls"
// @Router /toolcalls/ [get]
func (s *ToolcallService) GetToolcalls(c *gin.Context) {
var (
err error
query rdb.TableQuery
resp toolcalls
)
if err = c.ShouldBindQuery(&query); err != nil {
logger.FromContext(c).WithError(err).Errorf("error binding query")
response.Error(c, response.ErrToolcallsInvalidRequest, err)
return
}
uid := c.GetUint64("uid")
privs := c.GetStringSlice("prm")
var scope func(db *gorm.DB) *gorm.DB
if slices.Contains(privs, "toolcalls.admin") {
scope = func(db *gorm.DB) *gorm.DB {
return db.
Joins("INNER JOIN flows f ON f.id = toolcalls.flow_id")
}
} else if slices.Contains(privs, "toolcalls.view") {
scope = func(db *gorm.DB) *gorm.DB {
return db.
Joins("INNER JOIN flows f ON f.id = toolcalls.flow_id").
Where("f.user_id = ?", uid)
}
} else {
logger.FromContext(c).Errorf("error filtering user role permissions: permission not found")
response.Error(c, response.ErrNotPermitted, nil)
return
}
query.Init("toolcalls", toolcallsSQLMappers)
if query.Group != "" {
if _, ok := toolcallsSQLMappers[query.Group]; !ok {
logger.FromContext(c).Errorf("error finding toolcalls grouped: group field not found")
response.Error(c, response.ErrToolcallsInvalidRequest, errors.New("group field not found"))
return
}
var respGrouped toolcallsGrouped
if respGrouped.Total, err = query.QueryGrouped(s.db, &respGrouped.Grouped, scope); err != nil {
logger.FromContext(c).WithError(err).Errorf("error finding toolcalls grouped")
response.Error(c, response.ErrInternal, err)
return
}
response.Success(c, http.StatusOK, respGrouped)
return
}
if resp.Total, err = query.Query(s.db, &resp.Toolcalls, scope); err != nil {
logger.FromContext(c).WithError(err).Errorf("error finding toolcalls")
response.Error(c, response.ErrInternal, err)
return
}
for i := range resp.Toolcalls {
if err = resp.Toolcalls[i].Valid(); err != nil {
logger.FromContext(c).WithError(err).Errorf("error validating toolcall data '%d'", resp.Toolcalls[i].ID)
response.Error(c, response.ErrToolcallsInvalidData, err)
return
}
}
response.Success(c, http.StatusOK, resp)
}
// GetFlowToolcalls is a function to return toolcalls list by flow id
// @Summary Retrieve toolcalls list by flow id
// @Tags Toolcalls
// @Produce json
// @Security BearerAuth
// @Param flowID path int true "flow id" minimum(0)
// @Param request query rdb.TableQuery true "query table params"
// @Success 200 {object} response.successResp{data=toolcalls} "toolcalls list received successful"
// @Failure 400 {object} response.errorResp "invalid query request data"
// @Failure 403 {object} response.errorResp "getting toolcalls not permitted"
// @Failure 500 {object} response.errorResp "internal error on getting toolcalls"
// @Router /flows/{flowID}/toolcalls/ [get]
func (s *ToolcallService) GetFlowToolcalls(c *gin.Context) {
var (
err error
flowID uint64
query rdb.TableQuery
resp toolcalls
)
if flowID, err = strconv.ParseUint(c.Param("flowID"), 10, 64); err != nil {
logger.FromContext(c).WithError(err).Errorf("error parsing flow id")
response.Error(c, response.ErrToolcallsInvalidRequest, err)
return
}
if err = c.ShouldBindQuery(&query); err != nil {
logger.FromContext(c).WithError(err).Errorf("error binding query")
response.Error(c, response.ErrToolcallsInvalidRequest, err)
return
}
uid := c.GetUint64("uid")
privs := c.GetStringSlice("prm")
var scope func(db *gorm.DB) *gorm.DB
if slices.Contains(privs, "toolcalls.admin") {
scope = func(db *gorm.DB) *gorm.DB {
return db.
Joins("INNER JOIN flows f ON f.id = toolcalls.flow_id").
Where("f.id = ?", flowID)
}
} else if slices.Contains(privs, "toolcalls.view") {
scope = func(db *gorm.DB) *gorm.DB {
return db.
Joins("INNER JOIN flows f ON f.id = toolcalls.flow_id").
Where("f.id = ? AND f.user_id = ?", flowID, uid)
}
} else {
logger.FromContext(c).Errorf("error filtering user role permissions: permission not found")
response.Error(c, response.ErrNotPermitted, nil)
return
}
query.Init("toolcalls", toolcallsSQLMappers)
if query.Group != "" {
if _, ok := toolcallsSQLMappers[query.Group]; !ok {
logger.FromContext(c).Errorf("error finding toolcalls grouped: group field not found")
response.Error(c, response.ErrToolcallsInvalidRequest, errors.New("group field not found"))
return
}
var respGrouped toolcallsGrouped
if respGrouped.Total, err = query.QueryGrouped(s.db, &respGrouped.Grouped, scope); err != nil {
logger.FromContext(c).WithError(err).Errorf("error finding toolcalls grouped")
response.Error(c, response.ErrInternal, err)
return
}
response.Success(c, http.StatusOK, respGrouped)
return
}
if resp.Total, err = query.Query(s.db, &resp.Toolcalls, scope); err != nil {
logger.FromContext(c).WithError(err).Errorf("error finding toolcalls")
response.Error(c, response.ErrInternal, err)
return
}
for i := range resp.Toolcalls {
if err = resp.Toolcalls[i].Valid(); err != nil {
logger.FromContext(c).WithError(err).Errorf("error validating toolcall data '%d'", resp.Toolcalls[i].ID)
response.Error(c, response.ErrToolcallsInvalidData, err)
return
}
}
response.Success(c, http.StatusOK, resp)
}
// GetFlowToolcall is a function to return toolcall info by id and flow id
// @Summary Retrieve toolcall info by id and flow id
// @Tags Toolcalls
// @Produce json
// @Security BearerAuth
// @Param flowID path int true "flow id" minimum(0)
// @Param toolcallID path int true "toolcall id" minimum(0)
// @Success 200 {object} response.successResp{data=models.Toolcall} "toolcall info received successful"
// @Failure 400 {object} response.errorResp "invalid request data"
// @Failure 403 {object} response.errorResp "getting toolcall not permitted"
// @Failure 404 {object} response.errorResp "toolcall not found"
// @Failure 500 {object} response.errorResp "internal error on getting toolcall"
// @Router /flows/{flowID}/toolcalls/{toolcallID} [get]
func (s *ToolcallService) GetFlowToolcall(c *gin.Context) {
var (
err error
flowID uint64
toolcallID uint64
resp models.Toolcall
)
if flowID, err = strconv.ParseUint(c.Param("flowID"), 10, 64); err != nil {
logger.FromContext(c).WithError(err).Errorf("error parsing flow id")
response.Error(c, response.ErrToolcallsInvalidRequest, err)
return
}
if toolcallID, err = strconv.ParseUint(c.Param("toolcallID"), 10, 64); err != nil {
logger.FromContext(c).WithError(err).Errorf("error parsing toolcall id")
response.Error(c, response.ErrToolcallsInvalidRequest, err)
return
}
uid := c.GetUint64("uid")
privs := c.GetStringSlice("prm")
var scope func(db *gorm.DB) *gorm.DB
if slices.Contains(privs, "toolcalls.admin") {
scope = func(db *gorm.DB) *gorm.DB {
return db.Where("f.id = ?", flowID)
}
} else if slices.Contains(privs, "toolcalls.view") {
scope = func(db *gorm.DB) *gorm.DB {
return db.Where("f.id = ? AND f.user_id = ?", flowID, uid)
}
} else {
logger.FromContext(c).Errorf("error filtering user role permissions: permission not found")
response.Error(c, response.ErrNotPermitted, nil)
return
}
err = s.db.Model(&resp).
Joins("INNER JOIN flows f ON f.id = toolcalls.flow_id").
Scopes(scope).
Where("toolcalls.id = ?", toolcallID).
Take(&resp).Error
if err != nil {
logger.FromContext(c).WithError(err).Errorf("error on getting toolcall by id")
if gorm.IsRecordNotFoundError(err) {
response.Error(c, response.ErrToolcallsNotFound, err)
} else {
response.Error(c, response.ErrInternal, err)
}
return
}
response.Success(c, http.StatusOK, resp)
}
+684
View File
@@ -0,0 +1,684 @@
package services
import (
"errors"
"net/http"
"slices"
"pentagi/pkg/server/auth"
"pentagi/pkg/server/logger"
"pentagi/pkg/server/models"
"pentagi/pkg/server/rdb"
"pentagi/pkg/server/response"
"github.com/gin-gonic/gin"
"github.com/jinzhu/gorm"
"golang.org/x/crypto/bcrypt"
)
type users struct {
Users []models.UserRole `json:"users"`
Total uint64 `json:"total"`
}
type usersGrouped struct {
Grouped []string `json:"grouped"`
Total uint64 `json:"total"`
}
var usersSQLMappers = map[string]any{
"id": "{{table}}.id",
"hash": "{{table}}.hash",
"type": "{{table}}.type",
"mail": "{{table}}.mail",
"name": "{{table}}.name",
"role_id": "{{table}}.role_id",
"status": "{{table}}.status",
"created_at": "{{table}}.created_at",
"data": "({{table}}.hash || ' ' || {{table}}.mail || ' ' || {{table}}.name || ' ' || {{table}}.status)",
}
type UserService struct {
db *gorm.DB
userCache *auth.UserCache
}
func NewUserService(db *gorm.DB, userCache *auth.UserCache) *UserService {
return &UserService{
db: db,
userCache: userCache,
}
}
// GetCurrentUser is a function to return account information
// @Summary Retrieve current user information
// @Tags Users
// @Produce json
// @Success 200 {object} response.successResp{data=models.UserRolePrivileges} "user info received successful"
// @Failure 403 {object} response.errorResp "getting current user not permitted"
// @Failure 404 {object} response.errorResp "current user not found"
// @Failure 500 {object} response.errorResp "internal error on getting current user"
// @Router /user/ [get]
func (s *UserService) GetCurrentUser(c *gin.Context) {
var (
err error
resp models.UserRolePrivileges
)
uid := c.GetUint64("uid")
if err = s.db.Take(&resp.User, "id = ?", uid).Error; err != nil {
logger.FromContext(c).WithError(err).Errorf("error finding current user")
if errors.Is(err, gorm.ErrRecordNotFound) {
response.Error(c, response.ErrUsersNotFound, err)
} else {
response.Error(c, response.ErrInternal, err)
}
return
}
if err = s.db.Take(&resp.Role, "id = ?", resp.User.RoleID).Error; err != nil {
logger.FromContext(c).WithError(err).Errorf("error finding role by role id")
if errors.Is(err, gorm.ErrRecordNotFound) {
response.Error(c, response.ErrGetUserModelsNotFound, err)
} else {
response.Error(c, response.ErrInternal, err)
}
return
}
if err = s.db.Model(&resp.Role).Association("privileges").Find(&resp.Role.Privileges).Error; err != nil {
logger.FromContext(c).WithError(err).Errorf("error finding privileges by role id")
if errors.Is(err, gorm.ErrRecordNotFound) {
response.Error(c, response.ErrGetUserModelsNotFound, err)
} else {
response.Error(c, response.ErrInternal, err)
}
return
}
if err = resp.Valid(); err != nil {
logger.FromContext(c).WithError(err).Errorf("error validating user data '%s'", resp.Hash)
response.Error(c, response.ErrUsersInvalidData, err)
return
}
response.Success(c, http.StatusOK, resp)
}
// ChangePasswordCurrentUser is a function to update account password
// @Summary Update password for current user (account)
// @Tags Users
// @Accept json
// @Produce json
// @Param json body models.Password true "container to validate and update account password"
// @Success 200 {object} response.successResp "account password updated successful"
// @Failure 400 {object} response.errorResp "invalid account password form data"
// @Failure 403 {object} response.errorResp "updating account password not permitted"
// @Failure 404 {object} response.errorResp "current user not found"
// @Failure 500 {object} response.errorResp "internal error on updating account password"
// @Router /user/password [put]
func (s *UserService) ChangePasswordCurrentUser(c *gin.Context) {
var (
encPass []byte
err error
form models.Password
user models.UserPassword
)
if err = c.ShouldBindJSON(&form); err != nil || form.Valid() != nil {
if err == nil {
err = form.Valid()
}
logger.FromContext(c).WithError(err).Errorf("error binding JSON")
response.Error(c, response.ErrChangePasswordCurrentUserInvalidPassword, err)
return
}
uid := c.GetUint64("uid")
scope := func(db *gorm.DB) *gorm.DB {
return db.Where("id = ?", uid)
}
if err = s.db.Scopes(scope).Take(&user).Error; err != nil {
logger.FromContext(c).WithError(err).Errorf("error finding current user")
if errors.Is(err, gorm.ErrRecordNotFound) {
response.Error(c, response.ErrUsersNotFound, err)
} else {
response.Error(c, response.ErrInternal, err)
}
return
} else if err = user.Valid(); err != nil {
logger.FromContext(c).WithError(err).Errorf("error validating user data '%s'", user.Hash)
response.Error(c, response.ErrUsersInvalidData, err)
return
}
if err = bcrypt.CompareHashAndPassword([]byte(user.Password), []byte(form.CurrentPassword)); err != nil {
logger.FromContext(c).WithError(err).Errorf("error checking password for current user")
response.Error(c, response.ErrChangePasswordCurrentUserInvalidCurrentPassword, err)
return
}
if encPass, err = rdb.EncryptPassword(form.Password); err != nil {
logger.FromContext(c).WithError(err).Errorf("error making new password for current user")
response.Error(c, response.ErrChangePasswordCurrentUserInvalidNewPassword, err)
return
}
// Use map to update fields to avoid GORM ignoring zero values (false for bool)
updates := map[string]any{
"password": string(encPass),
"password_change_required": false,
}
if err = s.db.Model(&user).Scopes(scope).Updates(updates).Error; err != nil {
logger.FromContext(c).WithError(err).Errorf("error updating password for current user")
response.Error(c, response.ErrInternal, err)
return
}
response.Success(c, http.StatusOK, struct{}{})
}
// GetUsers returns users list
// @Summary Retrieve users list by filters
// @Tags Users
// @Produce json
// @Param request query rdb.TableQuery true "query table params"
// @Success 200 {object} response.successResp{data=users} "users list received successful"
// @Failure 400 {object} response.errorResp "invalid query request data"
// @Failure 403 {object} response.errorResp "getting users not permitted"
// @Failure 500 {object} response.errorResp "internal error on getting users"
// @Router /users/ [get]
func (s *UserService) GetUsers(c *gin.Context) {
var (
err error
query rdb.TableQuery
resp users
rids []uint64
roles []models.Role
)
if err = c.ShouldBindQuery(&query); err != nil {
logger.FromContext(c).WithError(err).Errorf("error binding query")
response.Error(c, response.ErrUsersInvalidRequest, err)
return
}
uid := c.GetUint64("uid")
privs := c.GetStringSlice("prm")
scope := func(db *gorm.DB) *gorm.DB {
if !slices.Contains(privs, "users.view") {
return db.Where("id = ?", uid)
}
return db
}
query.Init("users", usersSQLMappers)
if query.Group != "" {
if _, ok := usersSQLMappers[query.Group]; !ok {
logger.FromContext(c).Errorf("error finding users grouped: group field not found")
response.Error(c, response.ErrUsersInvalidRequest, errors.New("group field not found"))
return
}
var respGrouped usersGrouped
if respGrouped.Total, err = query.QueryGrouped(s.db, &respGrouped.Grouped, scope); err != nil {
logger.FromContext(c).WithError(err).Errorf("error finding users grouped")
response.Error(c, response.ErrInternal, err)
return
}
response.Success(c, http.StatusOK, respGrouped)
return
}
if resp.Total, err = query.Query(s.db, &resp.Users, scope); err != nil {
logger.FromContext(c).WithError(err).Errorf("error finding users")
response.Error(c, response.ErrInternal, err)
return
}
for _, user := range resp.Users {
rids = append(rids, user.RoleID)
}
if err = s.db.Find(&roles, "id IN (?)", rids).Error; err != nil {
logger.FromContext(c).WithError(err).Errorf("error finding linked roles")
response.Error(c, response.ErrInternal, err)
return
}
for i := range resp.Users {
roleID := resp.Users[i].RoleID
for _, role := range roles {
if roleID == role.ID {
resp.Users[i].Role = role
break
}
}
}
for i := 0; i < len(resp.Users); i++ {
if err = resp.Users[i].Valid(); err != nil {
logger.FromContext(c).WithError(err).Errorf("error validating user data '%s'", resp.Users[i].Hash)
response.Error(c, response.ErrUsersInvalidData, err)
return
}
}
response.Success(c, http.StatusOK, resp)
}
// GetUser is a function to return user by hash
// @Summary Retrieve user by hash
// @Tags Users
// @Produce json
// @Param hash path string true "hash in hex format (md5)" minlength(32) maxlength(32)
// @Success 200 {object} response.successResp{data=models.UserRolePrivileges} "user received successful"
// @Failure 403 {object} response.errorResp "getting user not permitted"
// @Failure 404 {object} response.errorResp "user not found"
// @Failure 500 {object} response.errorResp "internal error on getting user"
// @Router /users/{hash} [get]
func (s *UserService) GetUser(c *gin.Context) {
var (
err error
hash string = c.Param("hash")
resp models.UserRolePrivileges
)
uhash := c.GetString("uhash")
privs := c.GetStringSlice("prm")
if !slices.Contains(privs, "users.view") && uhash != hash {
logger.FromContext(c).Errorf("error filtering user role permissions: permission not found")
response.Error(c, response.ErrNotPermitted, nil)
return
}
if err = s.db.Take(&resp.User, "hash = ?", hash).Error; err != nil {
logger.FromContext(c).WithError(err).Errorf("error finding user by hash")
if errors.Is(err, gorm.ErrRecordNotFound) {
response.Error(c, response.ErrUsersNotFound, err)
} else {
response.Error(c, response.ErrInternal, err)
}
return
}
if err = s.db.Take(&resp.Role, "id = ?", resp.User.RoleID).Error; err != nil {
logger.FromContext(c).WithError(err).Errorf("error finding role by role id")
if errors.Is(err, gorm.ErrRecordNotFound) {
response.Error(c, response.ErrGetUserModelsNotFound, err)
} else {
response.Error(c, response.ErrInternal, err)
}
return
}
if err = s.db.Model(&resp.Role).Association("privileges").Find(&resp.Role.Privileges).Error; err != nil {
logger.FromContext(c).WithError(err).Errorf("error finding privileges by role id")
if errors.Is(err, gorm.ErrRecordNotFound) {
response.Error(c, response.ErrGetUserModelsNotFound, err)
} else {
response.Error(c, response.ErrInternal, err)
}
return
}
if err = resp.Valid(); err != nil {
logger.FromContext(c).WithError(err).Errorf("error validating user data '%s'", resp.Hash)
response.Error(c, response.ErrUsersInvalidData, err)
return
}
response.Success(c, http.StatusOK, resp)
}
// CreateUser is a function to create new user
// @Summary Create new user
// @Tags Users
// @Accept json
// @Produce json
// @Param json body models.UserPassword true "user model to create from"
// @Success 201 {object} response.successResp{data=models.UserRole} "user created successful"
// @Failure 400 {object} response.errorResp "invalid user request data"
// @Failure 403 {object} response.errorResp "creating user not permitted"
// @Failure 500 {object} response.errorResp "internal error on creating user"
// @Router /users/ [post]
func (s *UserService) CreateUser(c *gin.Context) {
var (
encPassword []byte
err error
resp models.UserRole
user models.UserPassword
)
if err = c.ShouldBindJSON(&user); err != nil {
logger.FromContext(c).WithError(err).Errorf("error binding JSON")
response.Error(c, response.ErrUsersInvalidRequest, err)
return
}
rid := c.GetUint64("rid")
privs := c.GetStringSlice("prm")
if !slices.Contains(privs, "users.create") {
logger.FromContext(c).Errorf("error filtering user role permissions: permission not found")
response.Error(c, response.ErrNotPermitted, nil)
return
}
privsCurrentUser, err := s.GetUserPrivileges(c, rid)
if err != nil {
logger.FromContext(c).WithError(err).Errorf("error getting current user privileges")
response.Error(c, response.ErrInternal, err)
return
}
privsNewUser, err := s.GetUserPrivileges(c, user.RoleID)
if err != nil {
logger.FromContext(c).WithError(err).Errorf("error getting new user privileges")
response.Error(c, response.ErrInternal, err)
return
}
if !s.CheckPrivilege(c, privsCurrentUser, privsNewUser) {
logger.FromContext(c).Errorf("error checking new user privileges")
response.Error(c, response.ErrNotPermitted, nil)
return
}
user.ID = 0
user.Hash = rdb.MakeUserHash(user.Name)
if err = user.Valid(); err != nil {
logger.FromContext(c).WithError(err).Errorf("error validating user")
response.Error(c, response.ErrCreateUserInvalidUser, err)
return
}
if encPassword, err = rdb.EncryptPassword(user.Password); err != nil {
logger.FromContext(c).WithError(err).Errorf("error encoding password")
response.Error(c, response.ErrInternal, err)
return
} else {
user.Password = string(encPassword)
}
tx := s.db.Begin()
if tx.Error != nil {
logger.FromContext(c).WithError(tx.Error).Errorf("error starting transaction")
response.Error(c, response.ErrInternal, tx.Error)
return
}
if err = tx.Create(&user).Error; err != nil {
tx.Rollback()
logger.FromContext(c).WithError(err).Errorf("error creating user")
response.Error(c, response.ErrInternal, err)
return
}
preferences := models.NewUserPreferences(user.ID)
if err = tx.Create(preferences).Error; err != nil {
tx.Rollback()
logger.FromContext(c).WithError(err).Errorf("error creating user preferences")
response.Error(c, response.ErrInternal, err)
return
}
if err = tx.Commit().Error; err != nil {
logger.FromContext(c).WithError(err).Errorf("error committing transaction")
response.Error(c, response.ErrInternal, err)
return
}
if err = s.db.Take(&resp.User, "hash = ?", user.Hash).Error; err != nil {
logger.FromContext(c).WithError(err).Errorf("error finding user by hash")
response.Error(c, response.ErrInternal, err)
return
}
if err = s.db.Take(&resp.Role, "id = ?", resp.User.RoleID).Error; err != nil {
logger.FromContext(c).WithError(err).Errorf("error finding role by role id")
response.Error(c, response.ErrInternal, err)
return
}
if err = resp.Valid(); err != nil {
logger.FromContext(c).WithError(err).Errorf("error validating user data '%s'", resp.Hash)
response.Error(c, response.ErrUsersInvalidData, err)
return
}
s.userCache.Invalidate(resp.User.ID)
response.Success(c, http.StatusCreated, resp)
}
// PatchUser is a function to update user by hash
// @Summary Update user
// @Tags Users
// @Accept json
// @Produce json
// @Param hash path string true "user hash in hex format (md5)" minlength(32) maxlength(32)
// @Param json body models.UserPassword true "user model to update"
// @Success 200 {object} response.successResp{data=models.UserRole} "user updated successful"
// @Failure 400 {object} response.errorResp "invalid user request data"
// @Failure 403 {object} response.errorResp "updating user not permitted"
// @Failure 404 {object} response.errorResp "user not found"
// @Failure 500 {object} response.errorResp "internal error on updating user"
// @Router /users/{hash} [put]
func (s *UserService) PatchUser(c *gin.Context) {
var (
err error
hash = c.Param("hash")
resp models.UserRole
user models.UserPassword
)
if err = c.ShouldBindJSON(&user); err != nil {
logger.FromContext(c).WithError(err).Errorf("error binding JSON")
response.Error(c, response.ErrUsersInvalidRequest, err)
return
} else if hash != user.Hash {
logger.FromContext(c).Errorf("mismatch user hash to requested one")
response.Error(c, response.ErrUsersInvalidRequest, nil)
return
} else if err = user.User.Valid(); err != nil {
logger.FromContext(c).WithError(err).Errorf("error validating user JSON")
response.Error(c, response.ErrUsersInvalidRequest, err)
return
} else if err = user.Valid(); user.Password != "" && err != nil {
logger.FromContext(c).WithError(err).Errorf("error validating user password")
response.Error(c, response.ErrUsersInvalidRequest, err)
return
}
uid := c.GetUint64("uid")
uhash := c.GetString("uhash")
privs := c.GetStringSlice("prm")
scope := func(db *gorm.DB) *gorm.DB {
if slices.Contains(privs, "users.edit") {
return db.Where("hash = ?", hash)
} else {
return db.Where("hash = ? AND id = ?", hash, uid)
}
}
if !slices.Contains(privs, "users.edit") && uhash != hash {
logger.FromContext(c).Errorf("error filtering user role permissions: permission not found")
response.Error(c, response.ErrNotPermitted, nil)
return
}
// Check if user exists before updating
var existingUser models.User
if err = s.db.Scopes(scope).Take(&existingUser).Error; err != nil {
logger.FromContext(c).WithError(err).Errorf("error finding user by hash")
if errors.Is(err, gorm.ErrRecordNotFound) {
response.Error(c, response.ErrUsersNotFound, err)
} else {
response.Error(c, response.ErrInternal, err)
}
return
}
if user.Password != "" {
var encPassword []byte
encPassword, err = rdb.EncryptPassword(user.Password)
if err != nil {
logger.FromContext(c).WithError(err).Errorf("error encoding password")
response.Error(c, response.ErrInternal, err)
return
}
// Use map to update fields to avoid GORM ignoring zero values (false for bool)
updates := map[string]any{
"name": user.Name,
"status": user.Status,
"password": string(encPassword),
"password_change_required": false,
}
err = s.db.Model(&existingUser).Updates(updates).Error
} else {
updates := map[string]any{
"name": user.Name,
"status": user.Status,
}
err = s.db.Model(&existingUser).Updates(updates).Error
}
if err != nil {
logger.FromContext(c).WithError(err).Errorf("error updating user by hash '%s'", hash)
response.Error(c, response.ErrInternal, err)
return
}
if err = s.db.Scopes(scope).Take(&resp.User).Error; err != nil {
logger.FromContext(c).WithError(err).Errorf("error finding user by hash")
if errors.Is(err, gorm.ErrRecordNotFound) {
response.Error(c, response.ErrUsersNotFound, err)
} else {
response.Error(c, response.ErrInternal, err)
}
return
}
if err = s.db.Take(&resp.Role, "id = ?", resp.User.RoleID).Error; err != nil {
logger.FromContext(c).WithError(err).Errorf("error finding role by role id")
if errors.Is(err, gorm.ErrRecordNotFound) {
response.Error(c, response.ErrPatchUserModelsNotFound, err)
} else {
response.Error(c, response.ErrInternal, err)
}
return
}
if err = resp.Valid(); err != nil {
logger.FromContext(c).WithError(err).Errorf("error validating user data '%s'", resp.Hash)
response.Error(c, response.ErrInternal, err)
return
}
s.userCache.Invalidate(resp.User.ID)
response.Success(c, http.StatusOK, resp)
}
// DeleteUser is a function to delete user by hash
// @Summary Delete user by hash
// @Tags Users
// @Produce json
// @Param hash path string true "hash in hex format (md5)" minlength(32) maxlength(32)
// @Success 200 {object} response.successResp "user deleted successful"
// @Failure 403 {object} response.errorResp "deleting user not permitted"
// @Failure 404 {object} response.errorResp "user not found"
// @Failure 500 {object} response.errorResp "internal error on deleting user"
// @Router /users/{hash} [delete]
func (s *UserService) DeleteUser(c *gin.Context) {
var (
err error
hash string = c.Param("hash")
user models.UserRole
)
uid := c.GetUint64("uid")
uhash := c.GetString("uhash")
privs := c.GetStringSlice("prm")
scope := func(db *gorm.DB) *gorm.DB {
if slices.Contains(privs, "users.delete") {
return db.Where("hash = ?", hash)
} else {
return db.Where("hash = ? AND id = ?", hash, uid)
}
}
if !slices.Contains(privs, "users.delete") && uhash != hash {
logger.FromContext(c).Errorf("error filtering user role permissions: permission not found")
response.Error(c, response.ErrNotPermitted, nil)
return
}
if err = s.db.Scopes(scope).Take(&user.User).Error; err != nil {
logger.FromContext(c).WithError(err).Errorf("error finding user by hash")
if errors.Is(err, gorm.ErrRecordNotFound) {
response.Error(c, response.ErrUsersNotFound, err)
} else {
response.Error(c, response.ErrInternal, err)
}
return
}
if err = s.db.Take(&user.Role, "id = ?", user.User.RoleID).Error; err != nil {
logger.FromContext(c).WithError(err).Errorf("error finding role by role id")
if errors.Is(err, gorm.ErrRecordNotFound) {
response.Error(c, response.ErrDeleteUserModelsNotFound, err)
} else {
response.Error(c, response.ErrInternal, err)
}
return
}
if err = user.Valid(); err != nil {
logger.FromContext(c).WithError(err).Errorf("error validating user data '%s'", user.Hash)
response.Error(c, response.ErrUsersInvalidData, err)
return
}
if err = s.db.Delete(&user.User).Error; err != nil {
logger.FromContext(c).WithError(err).Errorf("error deleting user by hash '%s'", hash)
response.Error(c, response.ErrInternal, err)
return
}
s.userCache.Invalidate(user.ID)
response.Success(c, http.StatusOK, struct{}{})
}
// GetUserPrivileges is a function to return user privileges
func (s *UserService) GetUserPrivileges(c *gin.Context, rid uint64) ([]string, error) {
var (
err error
privs []string
resp []models.Privilege
)
if err = s.db.Model(&models.Privilege{}).Where("role_id = ?", rid).Find(&resp).Error; err != nil {
return nil, err
}
for _, p := range resp {
if err = p.Valid(); err != nil {
logger.FromContext(c).WithError(err).Errorf("error validating privilege data '%s'", p.Name)
return nil, err
}
privs = append(privs, p.Name)
}
return privs, nil
}
// CheckPrivilege is a function to check if user has privilege
func (s *UserService) CheckPrivilege(c *gin.Context, privsCurrentUser, privsNewUser []string) bool {
for _, priv := range privsNewUser {
if !slices.Contains(privsCurrentUser, priv) {
return false
}
}
return true
}
+354
View File
@@ -0,0 +1,354 @@
package services
import (
"bytes"
"encoding/json"
"net/http"
"net/http/httptest"
"testing"
"pentagi/pkg/server/auth"
"pentagi/pkg/server/models"
"github.com/gin-gonic/gin"
"github.com/jinzhu/gorm"
_ "github.com/jinzhu/gorm/dialects/sqlite"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
func TestCreateUser_CreatesUserPreferences(t *testing.T) {
db := setupTestDB(t)
defer db.Close()
userCache := auth.NewUserCache(db)
service := NewUserService(db, userCache)
gin.SetMode(gin.TestMode)
w := httptest.NewRecorder()
c, _ := gin.CreateTestContext(w)
// Set up context with admin permissions
c.Set("uid", uint64(1))
c.Set("rid", uint64(1))
c.Set("uhash", "testhash1")
c.Set("prm", []string{"users.create"})
// Create request body
userRequest := models.UserPassword{
User: models.User{
Mail: "newuser@test.com",
Name: "New User",
RoleID: 2,
Status: models.UserStatusActive,
Type: models.UserTypeLocal,
},
Password: "SecurePass123!",
}
body, err := json.Marshal(userRequest)
require.NoError(t, err)
c.Request, _ = http.NewRequest("POST", "/users/", bytes.NewBuffer(body))
c.Request.Header.Set("Content-Type", "application/json")
// Call the handler
service.CreateUser(c)
// Check response status
assert.Equal(t, http.StatusCreated, w.Code, "Expected HTTP 201 Created")
// Verify user was created
var createdUser models.User
err = db.Where("mail = ?", "newuser@test.com").First(&createdUser).Error
require.NoError(t, err, "User should be created in database")
assert.Equal(t, "New User", createdUser.Name)
assert.Equal(t, uint64(2), createdUser.RoleID)
// Verify user_preferences was created
var userPrefs models.UserPreferences
err = db.Where("user_id = ?", createdUser.ID).First(&userPrefs).Error
require.NoError(t, err, "User preferences should be created in database")
assert.Equal(t, createdUser.ID, userPrefs.UserID)
assert.NotNil(t, userPrefs.Preferences.FavoriteFlows)
assert.Equal(t, 0, len(userPrefs.Preferences.FavoriteFlows), "FavoriteFlows should be empty array")
}
func TestCreateUser_RollbackOnPreferencesError(t *testing.T) {
db := setupTestDB(t)
defer db.Close()
// Drop user_preferences table to simulate error
db.Exec("DROP TABLE user_preferences")
userCache := auth.NewUserCache(db)
service := NewUserService(db, userCache)
gin.SetMode(gin.TestMode)
w := httptest.NewRecorder()
c, _ := gin.CreateTestContext(w)
c.Set("uid", uint64(1))
c.Set("rid", uint64(1))
c.Set("uhash", "testhash1")
c.Set("prm", []string{"users.create"})
userRequest := models.UserPassword{
User: models.User{
Mail: "failuser@test.com",
Name: "Fail User",
RoleID: 2,
Status: models.UserStatusActive,
Type: models.UserTypeLocal,
},
Password: "SecurePass123!",
}
body, err := json.Marshal(userRequest)
require.NoError(t, err)
c.Request, _ = http.NewRequest("POST", "/users/", bytes.NewBuffer(body))
c.Request.Header.Set("Content-Type", "application/json")
service.CreateUser(c)
// Should return error
assert.Equal(t, http.StatusInternalServerError, w.Code, "Expected HTTP 500 on preferences creation error")
// Verify user was NOT created (transaction rolled back)
var user models.User
err = db.Where("mail = ?", "failuser@test.com").First(&user).Error
assert.Error(t, err, "User should not exist due to transaction rollback")
assert.Equal(t, gorm.ErrRecordNotFound, err)
}
func TestCreateUser_InvalidPermissions(t *testing.T) {
db := setupTestDB(t)
defer db.Close()
userCache := auth.NewUserCache(db)
service := NewUserService(db, userCache)
gin.SetMode(gin.TestMode)
w := httptest.NewRecorder()
c, _ := gin.CreateTestContext(w)
// Set up context WITHOUT users.create permission
c.Set("uid", uint64(2))
c.Set("rid", uint64(2))
c.Set("uhash", "testhash2")
c.Set("prm", []string{"flows.view"})
userRequest := models.UserPassword{
User: models.User{
Mail: "unauthorized@test.com",
Name: "Unauthorized User",
RoleID: 2,
Status: models.UserStatusActive,
Type: models.UserTypeLocal,
},
Password: "SecurePass123!",
}
body, err := json.Marshal(userRequest)
require.NoError(t, err)
c.Request, _ = http.NewRequest("POST", "/users/", bytes.NewBuffer(body))
c.Request.Header.Set("Content-Type", "application/json")
service.CreateUser(c)
// Should return forbidden
assert.Equal(t, http.StatusForbidden, w.Code, "Expected HTTP 403 Forbidden")
// Verify user was NOT created
var user models.User
err = db.Where("mail = ?", "unauthorized@test.com").First(&user).Error
assert.Error(t, err, "User should not be created")
}
func TestCreateUser_MultipleUsers(t *testing.T) {
db := setupTestDB(t)
defer db.Close()
userCache := auth.NewUserCache(db)
service := NewUserService(db, userCache)
testCases := []struct {
name string
mail string
username string
roleID uint64
}{
{
name: "create first user",
mail: "newuser1@test.com",
username: "User One",
roleID: 2,
},
{
name: "create second user",
mail: "newuser2@test.com",
username: "User Two",
roleID: 2,
},
{
name: "create third user",
mail: "newuser3@test.com",
username: "User Three",
roleID: 2,
},
}
for _, tc := range testCases {
t.Run(tc.name, func(t *testing.T) {
gin.SetMode(gin.TestMode)
w := httptest.NewRecorder()
c, _ := gin.CreateTestContext(w)
c.Set("uid", uint64(1))
c.Set("rid", uint64(1))
c.Set("uhash", "testhash1")
c.Set("prm", []string{"users.create"})
userRequest := models.UserPassword{
User: models.User{
Mail: tc.mail,
Name: tc.username,
RoleID: tc.roleID,
Status: models.UserStatusActive,
Type: models.UserTypeLocal,
},
Password: "SecurePass123!",
}
body, err := json.Marshal(userRequest)
require.NoError(t, err)
c.Request, _ = http.NewRequest("POST", "/users/", bytes.NewBuffer(body))
c.Request.Header.Set("Content-Type", "application/json")
service.CreateUser(c)
assert.Equal(t, http.StatusCreated, w.Code, "Expected HTTP 201 Created")
// Verify both user and preferences were created
var user models.User
err = db.Where("mail = ?", tc.mail).First(&user).Error
require.NoError(t, err)
var prefs models.UserPreferences
err = db.Where("user_id = ?", user.ID).First(&prefs).Error
require.NoError(t, err)
assert.Equal(t, user.ID, prefs.UserID)
})
}
// Verify all users and preferences exist
var userCount int
db.Model(&models.User{}).Where("mail LIKE ?", "newuser%@test.com").Count(&userCount)
assert.Equal(t, 3, userCount, "Should have 3 newly created users")
var prefsCount int
db.Model(&models.UserPreferences{}).Count(&prefsCount)
assert.Equal(t, 5, prefsCount, "Should have 5 user preferences total (2 initial + 3 created)")
}
func TestCreateUser_InvalidJSON(t *testing.T) {
db := setupTestDB(t)
defer db.Close()
userCache := auth.NewUserCache(db)
service := NewUserService(db, userCache)
gin.SetMode(gin.TestMode)
w := httptest.NewRecorder()
c, _ := gin.CreateTestContext(w)
c.Set("uid", uint64(1))
c.Set("rid", uint64(1))
c.Set("uhash", "testhash1")
c.Set("prm", []string{"users.create"})
// Invalid JSON
c.Request, _ = http.NewRequest("POST", "/users/", bytes.NewBufferString("{invalid json"))
c.Request.Header.Set("Content-Type", "application/json")
service.CreateUser(c)
assert.Equal(t, http.StatusBadRequest, w.Code, "Expected HTTP 400 Bad Request")
}
func TestCreateUser_DuplicateEmail(t *testing.T) {
db := setupTestDB(t)
defer db.Close()
userCache := auth.NewUserCache(db)
service := NewUserService(db, userCache)
// Create first user
gin.SetMode(gin.TestMode)
w1 := httptest.NewRecorder()
c1, _ := gin.CreateTestContext(w1)
c1.Set("uid", uint64(1))
c1.Set("rid", uint64(1))
c1.Set("uhash", "testhash1")
c1.Set("prm", []string{"users.create"})
userRequest := models.UserPassword{
User: models.User{
Mail: "duplicate@test.com",
Name: "First User",
RoleID: 2,
Status: models.UserStatusActive,
Type: models.UserTypeLocal,
},
Password: "SecurePass123!",
}
body, err := json.Marshal(userRequest)
require.NoError(t, err)
c1.Request, _ = http.NewRequest("POST", "/users/", bytes.NewBuffer(body))
c1.Request.Header.Set("Content-Type", "application/json")
service.CreateUser(c1)
assert.Equal(t, http.StatusCreated, w1.Code)
// Try to create second user with same email
w2 := httptest.NewRecorder()
c2, _ := gin.CreateTestContext(w2)
c2.Set("uid", uint64(1))
c2.Set("rid", uint64(1))
c2.Set("uhash", "testhash1")
c2.Set("prm", []string{"users.create"})
userRequest2 := models.UserPassword{
User: models.User{
Mail: "duplicate@test.com", // Same email
Name: "Second User",
RoleID: 2,
Status: models.UserStatusActive,
Type: models.UserTypeLocal,
},
Password: "AnotherPass456!",
}
body2, err := json.Marshal(userRequest2)
require.NoError(t, err)
c2.Request, _ = http.NewRequest("POST", "/users/", bytes.NewBuffer(body2))
c2.Request.Header.Set("Content-Type", "application/json")
service.CreateUser(c2)
// Should fail due to unique constraint
assert.Equal(t, http.StatusInternalServerError, w2.Code, "Expected error on duplicate email")
// Verify only one user exists
var count int
db.Model(&models.User{}).Where("mail = ?", "duplicate@test.com").Count(&count)
assert.Equal(t, 1, count, "Should have only one user with this email")
}
+222
View File
@@ -0,0 +1,222 @@
package services
import (
"errors"
"net/http"
"slices"
"strconv"
"pentagi/pkg/server/logger"
"pentagi/pkg/server/models"
"pentagi/pkg/server/rdb"
"pentagi/pkg/server/response"
"github.com/gin-gonic/gin"
"github.com/jinzhu/gorm"
)
type vecstorelogs struct {
VecstoreLogs []models.Vecstorelog `json:"vecstorelogs"`
Total uint64 `json:"total"`
}
type vecstorelogsGrouped struct {
Grouped []string `json:"grouped"`
Total uint64 `json:"total"`
}
var vecstorelogsSQLMappers = map[string]any{
"id": "{{table}}.id",
"initiator": "{{table}}.initiator",
"executor": "{{table}}.executor",
"filter": "{{table}}.filter",
"query": "{{table}}.query",
"action": "{{table}}.action",
"result": "{{table}}.result",
"flow_id": "{{table}}.flow_id",
"task_id": "{{table}}.task_id",
"subtask_id": "{{table}}.subtask_id",
"created_at": "{{table}}.created_at",
"data": "({{table}}.filter || ' ' || {{table}}.query || ' ' || {{table}}.result)",
}
type VecstorelogService struct {
db *gorm.DB
}
func NewVecstorelogService(db *gorm.DB) *VecstorelogService {
return &VecstorelogService{
db: db,
}
}
// GetVecstorelogs is a function to return vecstorelogs list
// @Summary Retrieve vecstorelogs list
// @Tags Vecstorelogs
// @Produce json
// @Security BearerAuth
// @Param request query rdb.TableQuery true "query table params"
// @Success 200 {object} response.successResp{data=vecstorelogs} "vecstorelogs list received successful"
// @Failure 400 {object} response.errorResp "invalid query request data"
// @Failure 403 {object} response.errorResp "getting vecstorelogs not permitted"
// @Failure 500 {object} response.errorResp "internal error on getting vecstorelogs"
// @Router /vecstorelogs/ [get]
func (s *VecstorelogService) GetVecstorelogs(c *gin.Context) {
var (
err error
query rdb.TableQuery
resp vecstorelogs
)
if err = c.ShouldBindQuery(&query); err != nil {
logger.FromContext(c).WithError(err).Errorf("error binding query")
response.Error(c, response.ErrVecstorelogsInvalidRequest, err)
return
}
uid := c.GetUint64("uid")
privs := c.GetStringSlice("prm")
var scope func(db *gorm.DB) *gorm.DB
if slices.Contains(privs, "vecstorelogs.admin") {
scope = func(db *gorm.DB) *gorm.DB {
return db.
Joins("INNER JOIN flows f ON f.id = flow_id")
}
} else if slices.Contains(privs, "vecstorelogs.view") {
scope = func(db *gorm.DB) *gorm.DB {
return db.
Joins("INNER JOIN flows f ON f.id = flow_id").
Where("f.user_id = ?", uid)
}
} else {
logger.FromContext(c).Errorf("error filtering user role permissions: permission not found")
response.Error(c, response.ErrNotPermitted, nil)
return
}
query.Init("vecstorelogs", vecstorelogsSQLMappers)
if query.Group != "" {
if _, ok := vecstorelogsSQLMappers[query.Group]; !ok {
logger.FromContext(c).Errorf("error finding vecstorelogs grouped: group field not found")
response.Error(c, response.ErrVecstorelogsInvalidRequest, errors.New("group field not found"))
return
}
var respGrouped vecstorelogsGrouped
if respGrouped.Total, err = query.QueryGrouped(s.db, &respGrouped.Grouped, scope); err != nil {
logger.FromContext(c).WithError(err).Errorf("error finding vecstorelogs grouped")
response.Error(c, response.ErrInternal, err)
return
}
response.Success(c, http.StatusOK, respGrouped)
return
}
if resp.Total, err = query.Query(s.db, &resp.VecstoreLogs, scope); err != nil {
logger.FromContext(c).WithError(err).Errorf("error finding vecstorelogs")
response.Error(c, response.ErrInternal, err)
return
}
for i := 0; i < len(resp.VecstoreLogs); i++ {
if err = resp.VecstoreLogs[i].Valid(); err != nil {
logger.FromContext(c).WithError(err).Errorf("error validating vecstorelog data '%d'", resp.VecstoreLogs[i].ID)
response.Error(c, response.ErrVecstorelogsInvalidData, err)
return
}
}
response.Success(c, http.StatusOK, resp)
}
// GetFlowVecstorelogs is a function to return vecstorelogs list by flow id
// @Summary Retrieve vecstorelogs list by flow id
// @Tags Vecstorelogs
// @Produce json
// @Security BearerAuth
// @Param flowID path int true "flow id" minimum(0)
// @Param request query rdb.TableQuery true "query table params"
// @Success 200 {object} response.successResp{data=vecstorelogs} "vecstorelogs list received successful"
// @Failure 400 {object} response.errorResp "invalid query request data"
// @Failure 403 {object} response.errorResp "getting vecstorelogs not permitted"
// @Failure 500 {object} response.errorResp "internal error on getting vecstorelogs"
// @Router /flows/{flowID}/vecstorelogs/ [get]
func (s *VecstorelogService) GetFlowVecstorelogs(c *gin.Context) {
var (
err error
flowID uint64
query rdb.TableQuery
resp vecstorelogs
)
if flowID, err = strconv.ParseUint(c.Param("flowID"), 10, 64); err != nil {
logger.FromContext(c).WithError(err).Errorf("error parsing flow id")
response.Error(c, response.ErrVecstorelogsInvalidRequest, err)
return
}
if err = c.ShouldBindQuery(&query); err != nil {
logger.FromContext(c).WithError(err).Errorf("error binding query")
response.Error(c, response.ErrVecstorelogsInvalidRequest, err)
return
}
uid := c.GetUint64("uid")
privs := c.GetStringSlice("prm")
var scope func(db *gorm.DB) *gorm.DB
if slices.Contains(privs, "vecstorelogs.admin") {
scope = func(db *gorm.DB) *gorm.DB {
return db.
Joins("INNER JOIN flows f ON f.id = flow_id").
Where("f.id = ?", flowID)
}
} else if slices.Contains(privs, "vecstorelogs.view") {
scope = func(db *gorm.DB) *gorm.DB {
return db.
Joins("INNER JOIN flows f ON f.id = flow_id").
Where("f.id = ? AND f.user_id = ?", flowID, uid)
}
} else {
logger.FromContext(c).Errorf("error filtering user role permissions: permission not found")
response.Error(c, response.ErrNotPermitted, nil)
return
}
query.Init("vecstorelogs", vecstorelogsSQLMappers)
if query.Group != "" {
if _, ok := vecstorelogsSQLMappers[query.Group]; !ok {
logger.FromContext(c).Errorf("error finding vecstorelogs grouped: group field not found")
response.Error(c, response.ErrVecstorelogsInvalidRequest, errors.New("group field not found"))
return
}
var respGrouped vecstorelogsGrouped
if respGrouped.Total, err = query.QueryGrouped(s.db, &respGrouped.Grouped, scope); err != nil {
logger.FromContext(c).WithError(err).Errorf("error finding vecstorelogs grouped")
response.Error(c, response.ErrInternal, err)
return
}
response.Success(c, http.StatusOK, respGrouped)
return
}
if resp.Total, err = query.Query(s.db, &resp.VecstoreLogs, scope); err != nil {
logger.FromContext(c).WithError(err).Errorf("error finding vecstorelogs")
response.Error(c, response.ErrInternal, err)
return
}
for i := 0; i < len(resp.VecstoreLogs); i++ {
if err = resp.VecstoreLogs[i].Valid(); err != nil {
logger.FromContext(c).WithError(err).Errorf("error validating vecstorelog data '%d'", resp.VecstoreLogs[i].ID)
response.Error(c, response.ErrVecstorelogsInvalidData, err)
return
}
}
response.Success(c, http.StatusOK, resp)
}