first commit
This commit is contained in:
Beyhan Oğur
252
framework/configstore/tables/mcp.go
Normal file
252
framework/configstore/tables/mcp.go
Normal file
@@ -0,0 +1,252 @@
|
||||
package tables
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"time"
|
||||
|
||||
"github.com/bytedance/sonic"
|
||||
"github.com/maximhq/bifrost/core/schemas"
|
||||
"github.com/maximhq/bifrost/framework/encrypt"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
// TableMCPClient represents an MCP client configuration in the database
|
||||
type TableMCPClient struct {
|
||||
ID uint `gorm:"primaryKey;autoIncrement" json:"id"` // ID is used as the internal primary key and is also accessed by public methods, so it must be present.
|
||||
ClientID string `gorm:"type:varchar(255);uniqueIndex;not null" json:"client_id"`
|
||||
Name string `gorm:"type:varchar(255);uniqueIndex;not null" json:"name"`
|
||||
IsCodeModeClient bool `gorm:"default:false" json:"is_code_mode_client"` // Whether the client is a code mode client
|
||||
ConnectionType string `gorm:"type:varchar(20);not null" json:"connection_type"` // schemas.MCPConnectionType
|
||||
ConnectionString *schemas.EnvVar `gorm:"type:text" json:"connection_string,omitempty"`
|
||||
StdioConfigJSON *string `gorm:"type:text" json:"-"` // JSON serialized schemas.MCPStdioConfig
|
||||
ToolsToExecuteJSON string `gorm:"type:text" json:"-"` // JSON serialized []string
|
||||
ToolsToAutoExecuteJSON string `gorm:"type:text" json:"-"` // JSON serialized []string
|
||||
HeadersJSON string `gorm:"type:text" json:"-"` // JSON serialized map[string]string
|
||||
AllowedExtraHeadersJSON string `gorm:"type:text" json:"-"` // JSON serialized []string
|
||||
IsPingAvailable *bool `gorm:"default:true" json:"is_ping_available,omitempty"` // Whether the MCP server supports ping for health checks
|
||||
ToolPricingJSON string `gorm:"type:text" json:"-"` // JSON serialized map[string]float64
|
||||
ToolSyncInterval int `gorm:"default:0" json:"tool_sync_interval"` // Per-client tool sync interval in minutes (0 = use global, -1 = disabled)
|
||||
|
||||
// Per-user OAuth: discovered tools persisted so they survive restart
|
||||
DiscoveredToolsJSON string `gorm:"type:text" json:"-"` // JSON serialized map[string]schemas.ChatTool
|
||||
ToolNameMappingJSON string `gorm:"type:text" json:"-"` // JSON serialized map[string]string
|
||||
|
||||
// OAuth authentication fields
|
||||
AuthType string `gorm:"type:varchar(20);default:'headers'" json:"auth_type"` // "none", "headers", "oauth"
|
||||
OauthConfigID *string `gorm:"type:varchar(255);index;constraint:OnDelete:CASCADE" json:"oauth_config_id"` // Foreign key to oauth_configs.ID with CASCADE delete
|
||||
OauthConfig *TableOauthConfig `gorm:"foreignKey:OauthConfigID;references:ID;constraint:OnDelete:CASCADE" json:"-"` // Gorm relationship
|
||||
|
||||
AllowOnAllVirtualKeys bool `gorm:"default:false" json:"allow_on_all_virtual_keys"` // Whether to allow the MCP client to run on all virtual keys
|
||||
|
||||
// Config hash is used to detect the changes synced from config.json file
|
||||
// Every time we sync the config.json file, we will update the config hash
|
||||
ConfigHash string `gorm:"type:varchar(255);null" json:"config_hash"`
|
||||
|
||||
EncryptionStatus string `gorm:"type:varchar(20);default:'plain_text'" json:"-"`
|
||||
|
||||
CreatedAt time.Time `gorm:"index;not null" json:"created_at"`
|
||||
UpdatedAt time.Time `gorm:"index;not null" json:"updated_at"`
|
||||
|
||||
// Virtual fields for runtime use (not stored in DB)
|
||||
StdioConfig *schemas.MCPStdioConfig `gorm:"-" json:"stdio_config,omitempty"`
|
||||
ToolsToExecute schemas.WhiteList `gorm:"-" json:"tools_to_execute"`
|
||||
ToolsToAutoExecute schemas.WhiteList `gorm:"-" json:"tools_to_auto_execute"`
|
||||
Headers map[string]schemas.EnvVar `gorm:"-" json:"headers"`
|
||||
AllowedExtraHeaders schemas.WhiteList `gorm:"-" json:"allowed_extra_headers"`
|
||||
ToolPricing map[string]float64 `gorm:"-" json:"tool_pricing"`
|
||||
DiscoveredTools map[string]schemas.ChatTool `gorm:"-" json:"-"`
|
||||
DiscoveredToolNameMapping map[string]string `gorm:"-" json:"-"`
|
||||
}
|
||||
|
||||
// TableName sets the table name for each model
|
||||
func (TableMCPClient) TableName() string { return "config_mcp_clients" }
|
||||
|
||||
// BeforeSave is a GORM hook that serializes runtime fields (stdio config, tools, headers,
|
||||
// pricing) into JSON columns and encrypts the connection string and headers before writing
|
||||
// to the database. Environment-variable-backed connection strings are not encrypted.
|
||||
func (c *TableMCPClient) BeforeSave(tx *gorm.DB) error {
|
||||
if c.StdioConfig != nil {
|
||||
data, err := json.Marshal(c.StdioConfig)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
config := string(data)
|
||||
c.StdioConfigJSON = &config
|
||||
} else {
|
||||
c.StdioConfigJSON = nil
|
||||
}
|
||||
|
||||
if c.ToolsToExecute != nil {
|
||||
if err := c.ToolsToExecute.Validate(); err != nil {
|
||||
return fmt.Errorf("invalid tools_to_execute: %w", err)
|
||||
}
|
||||
data, err := json.Marshal(c.ToolsToExecute)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
c.ToolsToExecuteJSON = string(data)
|
||||
} else {
|
||||
c.ToolsToExecuteJSON = "[]"
|
||||
}
|
||||
|
||||
if c.ToolsToAutoExecute != nil {
|
||||
if err := c.ToolsToAutoExecute.Validate(); err != nil {
|
||||
return fmt.Errorf("invalid tools_to_auto_execute: %w", err)
|
||||
}
|
||||
data, err := json.Marshal(c.ToolsToAutoExecute)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
c.ToolsToAutoExecuteJSON = string(data)
|
||||
} else {
|
||||
c.ToolsToAutoExecuteJSON = "[]"
|
||||
}
|
||||
|
||||
if c.Headers != nil {
|
||||
headersToSerialize := make(map[string]string, len(c.Headers))
|
||||
for key, value := range c.Headers {
|
||||
if value.IsFromEnv() {
|
||||
headersToSerialize[key] = value.EnvVar
|
||||
} else {
|
||||
headersToSerialize[key] = value.GetValue()
|
||||
}
|
||||
}
|
||||
data, err := json.Marshal(headersToSerialize)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
c.HeadersJSON = string(data)
|
||||
} else {
|
||||
c.HeadersJSON = "{}"
|
||||
}
|
||||
|
||||
if c.AllowedExtraHeaders != nil {
|
||||
if err := c.AllowedExtraHeaders.Validate(); err != nil {
|
||||
return fmt.Errorf("invalid allowed_extra_headers: %w", err)
|
||||
}
|
||||
data, err := json.Marshal(c.AllowedExtraHeaders)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
c.AllowedExtraHeadersJSON = string(data)
|
||||
} else {
|
||||
c.AllowedExtraHeadersJSON = "[]"
|
||||
}
|
||||
|
||||
if c.ToolPricing != nil {
|
||||
data, err := json.Marshal(c.ToolPricing)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
c.ToolPricingJSON = string(data)
|
||||
} else {
|
||||
c.ToolPricingJSON = "{}"
|
||||
}
|
||||
|
||||
if c.DiscoveredTools != nil {
|
||||
data, err := json.Marshal(c.DiscoveredTools)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
c.DiscoveredToolsJSON = string(data)
|
||||
}
|
||||
|
||||
if c.DiscoveredToolNameMapping != nil {
|
||||
data, err := json.Marshal(c.DiscoveredToolNameMapping)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
c.ToolNameMappingJSON = string(data)
|
||||
}
|
||||
|
||||
// Encrypt sensitive fields after serialization.
|
||||
// Always set EncryptionStatus when encryption is enabled so the startup
|
||||
// batch pass does not re-process this row indefinitely.
|
||||
if encrypt.IsEnabled() {
|
||||
if c.ConnectionString != nil && !c.ConnectionString.IsFromEnv() && c.ConnectionString.GetValue() != "" {
|
||||
// Copy to avoid encrypting the shared ConnectionString through the pointer
|
||||
cs := *c.ConnectionString
|
||||
enc, err := encrypt.Encrypt(cs.Val)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to encrypt mcp connection string: %w", err)
|
||||
}
|
||||
cs.Val = enc
|
||||
c.ConnectionString = &cs
|
||||
}
|
||||
if c.HeadersJSON != "" && c.HeadersJSON != "{}" {
|
||||
enc, err := encrypt.Encrypt(c.HeadersJSON)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to encrypt mcp headers: %w", err)
|
||||
}
|
||||
c.HeadersJSON = enc
|
||||
}
|
||||
c.EncryptionStatus = EncryptionStatusEncrypted
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// AfterFind is a GORM hook that decrypts the connection string and headers (if encrypted)
|
||||
// and deserializes JSON columns back into runtime structs after reading from the database.
|
||||
func (c *TableMCPClient) AfterFind(tx *gorm.DB) error {
|
||||
if c.EncryptionStatus == "encrypted" {
|
||||
if c.HeadersJSON != "" && c.HeadersJSON != "{}" {
|
||||
decrypted, err := encrypt.Decrypt(c.HeadersJSON)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to decrypt mcp headers: %w", err)
|
||||
}
|
||||
c.HeadersJSON = decrypted
|
||||
}
|
||||
if c.ConnectionString != nil && !c.ConnectionString.IsFromEnv() && c.ConnectionString.GetValue() != "" {
|
||||
decrypted, err := encrypt.Decrypt(c.ConnectionString.Val)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to decrypt mcp connection string: %w", err)
|
||||
}
|
||||
c.ConnectionString.Val = decrypted
|
||||
}
|
||||
}
|
||||
if c.StdioConfigJSON != nil {
|
||||
var config schemas.MCPStdioConfig
|
||||
if err := sonic.Unmarshal([]byte(*c.StdioConfigJSON), &config); err != nil {
|
||||
return err
|
||||
}
|
||||
c.StdioConfig = &config
|
||||
}
|
||||
if c.ToolsToExecuteJSON != "" {
|
||||
if err := sonic.Unmarshal([]byte(c.ToolsToExecuteJSON), &c.ToolsToExecute); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
if c.ToolsToAutoExecuteJSON != "" {
|
||||
if err := sonic.Unmarshal([]byte(c.ToolsToAutoExecuteJSON), &c.ToolsToAutoExecute); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
if c.HeadersJSON != "" {
|
||||
if err := sonic.Unmarshal([]byte(c.HeadersJSON), &c.Headers); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
if c.AllowedExtraHeadersJSON != "" {
|
||||
if err := sonic.Unmarshal([]byte(c.AllowedExtraHeadersJSON), &c.AllowedExtraHeaders); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
if c.ToolPricingJSON != "" {
|
||||
if err := json.Unmarshal([]byte(c.ToolPricingJSON), &c.ToolPricing); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
if c.DiscoveredToolsJSON != "" {
|
||||
if err := sonic.Unmarshal([]byte(c.DiscoveredToolsJSON), &c.DiscoveredTools); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
if c.ToolNameMappingJSON != "" {
|
||||
if err := sonic.Unmarshal([]byte(c.ToolNameMappingJSON), &c.DiscoveredToolNameMapping); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
Reference in New Issue
Block a user