Skip to content

Commit ab685cf

Browse files
akshaydeoakshaydeo
authored
fixes race condition while loading mcp servers from config file (#1470)
## Support for MCP Client Names in Virtual Key Configurations This PR adds support for referencing MCP clients by name in virtual key configurations, making it easier to define MCP access permissions in config files. ## Issue Closes #1462 ## Changes - Added support for `mcp_client_name` in virtual key MCP configurations, allowing config files to reference MCP clients by their human-readable names instead of database IDs - Implemented a resolution mechanism that converts MCP client names to their corresponding database IDs during config loading - Added graceful handling for unresolvable MCP client names, with appropriate warning logs - Created a new example configuration demonstrating the feature with PostgreSQL and MCP clients - Added comprehensive tests to verify the name resolution works correctly ## Type of change - [x] Feature - [x] Bug fix ## Affected areas - [x] Core (Go) - [x] Transports (HTTP) ## How to test ```sh # Run the tests for the new functionality go test ./transports/bifrost-http/lib -run TestSQLite_VKMCPConfig_MCPClientNameResolution go test ./transports/bifrost-http/lib -run TestSQLite_VKMCPConfig_MCPClientNameNotFound # Try the example config cp examples/configs/withpostgresmcpclientsinconfig/config.json ./config.json # Set required environment variables # Start Bifrost ``` ## Breaking changes - [x] No ## Related issues Fixes an issue where virtual keys with MCP configurations in config files would fail with foreign key constraint violations when using client names instead of IDs. ## Security considerations No security implications - this is a usability improvement that maintains the same access control model. ## Checklist - [x] I added/updated tests where appropriate - [x] I added example configuration
2 parents f5f0f02 + 8f6f420 commit ab685cf

6 files changed

Lines changed: 655 additions & 0 deletions

File tree

‎.github/workflows/configs/withpostgresmcpclientsinconfig/config.json‎

Lines changed: 139 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,139 @@
1+
{
2+
"$schema": "https://www.getbifrost.ai/schema",
3+
"client": {
4+
"allow_direct_keys": false,
5+
"allowed_origins": [
6+
"*"
7+
],
8+
"disable_content_logging": false,
9+
"drop_excess_requests": false,
10+
"enable_governance": true,
11+
"enable_litellm_fallbacks": false,
12+
"enable_logging": true,
13+
"enforce_governance_header": true,
14+
"initial_pool_size": 300,
15+
"log_retention_days": 365,
16+
"max_request_body_size_mb": 100
17+
},
18+
"config_store": {
19+
"enabled": true,
20+
"type": "postgres",
21+
"config": {
22+
"host": "localhost",
23+
"port": "5432",
24+
"user": "bifrost",
25+
"password": "bifrost_password",
26+
"db_name": "bifrost",
27+
"ssl_mode": "disable"
28+
}
29+
},
30+
"logs_store": {
31+
"enabled": true,
32+
"type": "postgres",
33+
"config": {
34+
"host": "localhost",
35+
"port": "5432",
36+
"user": "bifrost",
37+
"password": "bifrost_password",
38+
"db_name": "bifrost",
39+
"ssl_mode": "disable"
40+
}
41+
},
42+
"mcp": {
43+
"client_configs": [
44+
{
45+
"id": "weather-mcp-server",
46+
"name": "WeatherService",
47+
"connection_type": "http",
48+
"http_url": "http://localhost:8080/mcp",
49+
"is_enabled": true
50+
},
51+
{
52+
"id": "calendar-mcp-server",
53+
"name": "CalendarService",
54+
"connection_type": "http",
55+
"http_url": "http://localhost:8081/mcp",
56+
"is_enabled": true
57+
}
58+
]
59+
},
60+
"governance": {
61+
"auth_config": {
62+
"admin_password": "env.BIFROST_ADMIN_PASSWORD",
63+
"admin_username": "env.BIFROST_ADMIN_USERNAME",
64+
"disable_auth_on_inference": true,
65+
"is_enabled": false
66+
},
67+
"virtual_keys": [
68+
{
69+
"id": "vk-ai-portal-prod",
70+
"is_active": true,
71+
"name": "ai-portal-production-key",
72+
"description": "Virtual key for AI portal with MCP access to weather and calendar services",
73+
"value": "env.BIFROST_VK_AI_PORTAL",
74+
"mcp_configs": [
75+
{
76+
"mcp_client_name": "WeatherService",
77+
"tools_to_execute": [
78+
"*"
79+
]
80+
},
81+
{
82+
"mcp_client_name": "CalendarService",
83+
"tools_to_execute": [
84+
"get_events",
85+
"create_event"
86+
]
87+
}
88+
],
89+
"provider_configs": [
90+
{
91+
"provider": "openai",
92+
"weight": 1.0
93+
}
94+
]
95+
},
96+
{
97+
"id": "vk-internal-tools",
98+
"is_active": true,
99+
"name": "internal-tools-key",
100+
"description": "Virtual key for internal tools with limited MCP access",
101+
"value": "env.BIFROST_VK_INTERNAL",
102+
"mcp_configs": [
103+
{
104+
"mcp_client_name": "WeatherService",
105+
"tools_to_execute": [
106+
"get_current_weather"
107+
]
108+
}
109+
],
110+
"provider_configs": [
111+
{
112+
"provider": "openai",
113+
"weight": 1.0
114+
}
115+
]
116+
}
117+
]
118+
},
119+
"plugins": [
120+
{
121+
"config": {
122+
"is_vk_mandatory": true
123+
},
124+
"enabled": true,
125+
"name": "governance"
126+
}
127+
],
128+
"providers": {
129+
"openai": {
130+
"keys": [
131+
{
132+
"name": "openai-primary",
133+
"value": "env.OPENAI_API_KEY",
134+
"weight": 1
135+
}
136+
]
137+
}
138+
}
139+
}

‎.github/workflows/scripts/release-bifrost-http.sh‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -151,6 +151,7 @@ CONFIGS_TO_TEST=(
151151
"withdynamicplugin"
152152
"withobservability"
153153
"withsemanticcache"
154+
"withpostgresmcpclientsinconfig"
154155
)
155156

156157
TEST_BINARY="../tmp/bifrost-http"

‎examples/configs/withpostgresmcpclientsinconfig/config.json‎

Lines changed: 139 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,139 @@
1+
{
2+
"$schema": "https://www.getbifrost.ai/schema",
3+
"client": {
4+
"allow_direct_keys": false,
5+
"allowed_origins": [
6+
"*"
7+
],
8+
"disable_content_logging": false,
9+
"drop_excess_requests": false,
10+
"enable_governance": true,
11+
"enable_litellm_fallbacks": false,
12+
"enable_logging": true,
13+
"enforce_governance_header": true,
14+
"initial_pool_size": 300,
15+
"log_retention_days": 365,
16+
"max_request_body_size_mb": 100
17+
},
18+
"config_store": {
19+
"enabled": true,
20+
"type": "postgres",
21+
"config": {
22+
"host": "localhost",
23+
"port": "5432",
24+
"user": "bifrost",
25+
"password": "bifrost_password",
26+
"db_name": "bifrost",
27+
"ssl_mode": "disable"
28+
}
29+
},
30+
"logs_store": {
31+
"enabled": true,
32+
"type": "postgres",
33+
"config": {
34+
"host": "localhost",
35+
"port": "5432",
36+
"user": "bifrost",
37+
"password": "bifrost_password",
38+
"db_name": "bifrost",
39+
"ssl_mode": "disable"
40+
}
41+
},
42+
"mcp": {
43+
"client_configs": [
44+
{
45+
"id": "weather-mcp-server",
46+
"name": "WeatherService",
47+
"connection_type": "http",
48+
"http_url": "http://localhost:8080/mcp",
49+
"is_enabled": true
50+
},
51+
{
52+
"id": "calendar-mcp-server",
53+
"name": "CalendarService",
54+
"connection_type": "http",
55+
"http_url": "http://localhost:8081/mcp",
56+
"is_enabled": true
57+
}
58+
]
59+
},
60+
"governance": {
61+
"auth_config": {
62+
"admin_password": "env.BIFROST_ADMIN_PASSWORD",
63+
"admin_username": "env.BIFROST_ADMIN_USERNAME",
64+
"disable_auth_on_inference": true,
65+
"is_enabled": false
66+
},
67+
"virtual_keys": [
68+
{
69+
"id": "vk-ai-portal-prod",
70+
"is_active": true,
71+
"name": "ai-portal-production-key",
72+
"description": "Virtual key for AI portal with MCP access to weather and calendar services",
73+
"value": "env.BIFROST_VK_AI_PORTAL",
74+
"mcp_configs": [
75+
{
76+
"mcp_client_name": "WeatherService",
77+
"tools_to_execute": [
78+
"*"
79+
]
80+
},
81+
{
82+
"mcp_client_name": "CalendarService",
83+
"tools_to_execute": [
84+
"get_events",
85+
"create_event"
86+
]
87+
}
88+
],
89+
"provider_configs": [
90+
{
91+
"provider": "openai",
92+
"weight": 1.0
93+
}
94+
]
95+
},
96+
{
97+
"id": "vk-internal-tools",
98+
"is_active": true,
99+
"name": "internal-tools-key",
100+
"description": "Virtual key for internal tools with limited MCP access",
101+
"value": "env.BIFROST_VK_INTERNAL",
102+
"mcp_configs": [
103+
{
104+
"mcp_client_name": "WeatherService",
105+
"tools_to_execute": [
106+
"get_current_weather"
107+
]
108+
}
109+
],
110+
"provider_configs": [
111+
{
112+
"provider": "openai",
113+
"weight": 1.0
114+
}
115+
]
116+
}
117+
]
118+
},
119+
"plugins": [
120+
{
121+
"config": {
122+
"is_vk_mandatory": true
123+
},
124+
"enabled": true,
125+
"name": "governance"
126+
}
127+
],
128+
"providers": {
129+
"openai": {
130+
"keys": [
131+
{
132+
"name": "openai-primary",
133+
"value": "env.OPENAI_API_KEY",
134+
"weight": 1
135+
}
136+
]
137+
}
138+
}
139+
}

‎framework/configstore/tables/virtualkey.go‎

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -137,13 +137,44 @@ type TableVirtualKeyMCPConfig struct {
137137
MCPClientID uint `gorm:"not null;uniqueIndex:idx_vk_mcpclient" json:"mcp_client_id"`
138138
MCPClient TableMCPClient `gorm:"foreignKey:MCPClientID" json:"mcp_client"`
139139
ToolsToExecute []string `gorm:"type:text;serializer:json" json:"tools_to_execute"`
140+
141+
// MCPClientName is used during config file parsing to resolve the MCP client by name.
142+
// This field is not persisted to the database - it's only used to capture
143+
// "mcp_client_name" from config.json and then resolve it to MCPClientID.
144+
MCPClientName string `gorm:"-" json:"-"`
140145
}
141146

142147
// TableName sets the table name for each model
143148
func (TableVirtualKeyMCPConfig) TableName() string {
144149
return "governance_virtual_key_mcp_configs"
145150
}
146151

152+
// UnmarshalJSON custom unmarshaller to handle both "mcp_client_id" (database format)
153+
// and "mcp_client_name" (config file format) for MCP client references.
154+
func (mc *TableVirtualKeyMCPConfig) UnmarshalJSON(data []byte) error {
155+
// Temporary struct to capture all fields including mcp_client_name
156+
type Alias TableVirtualKeyMCPConfig
157+
type TempMCPConfig struct {
158+
Alias
159+
MCPClientName string `json:"mcp_client_name"` // Config file format: MCP client name
160+
}
161+
162+
var temp TempMCPConfig
163+
if err := json.Unmarshal(data, &temp); err != nil {
164+
return err
165+
}
166+
167+
// Copy all standard fields
168+
*mc = TableVirtualKeyMCPConfig(temp.Alias)
169+
170+
// Capture mcp_client_name for later resolution to MCPClientID
171+
if temp.MCPClientName != "" {
172+
mc.MCPClientName = temp.MCPClientName
173+
}
174+
175+
return nil
176+
}
177+
147178
// TableVirtualKey represents a virtual key with budget, rate limits, and team/customer association
148179
type TableVirtualKey struct {
149180
ID string `gorm:"primaryKey;type:varchar(255)" json:"id"`

0 commit comments

Comments
 (0)