From 65bdd6e648917a2693b257631f272453bade34b4 Mon Sep 17 00:00:00 2001
From: Pratham-Mishra04 <pratham@getmaxim.ai>
Date: Fri, 29 May 2026 17:46:06 +0530
Subject: [PATCH] fix: test fixes

---
 framework/configstore/migrations.go        | 12 ++++++++---
 transports/bifrost-http/lib/config_test.go |  3 +++
 transports/config.schema.json              | 25 +++++++++++++++++++++-
 3 files changed, 36 insertions(+), 4 deletions(-)

diff --git a/framework/configstore/migrations.go b/framework/configstore/migrations.go
index 89dcb67fd57..7392a56b94a 100644
--- a/framework/configstore/migrations.go
+++ b/framework/configstore/migrations.go
@@ -801,6 +801,15 @@ func triggerMigrations(ctx context.Context, db *gorm.DB) error {
 	if err := migrationAddCreatedByUserIDColumnForVirtualKeys(ctx, db); err != nil {
 		return err
 	}
+	// Must run before migrationRefreshConfigHashAfterMCPExternalServerURLRemoval:
+	// that migration SELECTs config_client using the column list derived from
+	// the TableClientConfig struct, which still declares allow_direct_keys.
+	// Without re-adding the column first, the SELECT fails with
+	// "no such column: allow_direct_keys" on any DB where the earlier
+	// drop_allow_direct_keys_column_ddl migration has run.
+	if err := migrationReAddAllowDirectKeysColumn(ctx, db); err != nil {
+		return err
+	}
 	if err := migrationRefreshConfigHashAfterMCPExternalServerURLRemoval(ctx, db); err != nil {
 		return err
 	}
@@ -813,9 +822,6 @@ func triggerMigrations(ctx context.Context, db *gorm.DB) error {
 	if err := migrationAddPerUserHeadersFlowsTable(ctx, db); err != nil {
 		return err
 	}
-	if err := migrationReAddAllowDirectKeysColumn(ctx, db); err != nil {
-		return err
-	}
 	if err := migrationAddMCPClientTLSConfigColumn(ctx, db); err != nil {
 		return err
 	}
diff --git a/transports/bifrost-http/lib/config_test.go b/transports/bifrost-http/lib/config_test.go
index 08381a4bae2..0ba5d488303 100644
--- a/transports/bifrost-http/lib/config_test.go
+++ b/transports/bifrost-http/lib/config_test.go
@@ -15766,6 +15766,9 @@ var excludedSchemaFields = map[string]map[string]bool{
 		"budget_id":        true, // Replaced by budgets[] relationship with team_id FK on TableBudget
 		"business_unit_id": true, // Enterprise feature; not in OSS TableTeam
 	},
+	"governance.virtual_keys": {
+		"access_profile_id": true, // Enterprise access-profile assignment; not on OSS TableVirtualKey
+	},
 	"governance.virtual_keys.provider_configs": {
 		"keys":    true, // Complex nested type, validated separately
 		"key_ids": true, // Config-file format; handled via custom UnmarshalJSON into allow_all_keys/keys
diff --git a/transports/config.schema.json b/transports/config.schema.json
index df1ed060e44..914caa5d24f 100644
--- a/transports/config.schema.json
+++ b/transports/config.schema.json
@@ -190,6 +190,11 @@
           },
           "description": "Additional allowed headers for CORS and WebSocket"
         },
+        "allow_direct_keys": {
+          "type": "boolean",
+          "description": "Allow callers to bypass the registered key pool by supplying x-bf-direct-key: true and an Authorization header carrying the provider's raw API key.",
+          "default": false
+        },
         "mcp_agent_depth": {
           "type": "integer",
           "minimum": 1,
@@ -485,6 +490,10 @@
                 "type": "string",
                 "description": "Team name"
               },
+              "source_id": {
+                "type": "string",
+                "description": "Optional external source identifier (e.g. SCIM group ID) used to map this team to an upstream system. Unique when set."
+              },
               "customer_id": {
                 "type": "string",
                 "description": "Associated customer ID"
@@ -1908,6 +1917,13 @@
             "type": "string"
           }
         },
+        "blacklisted_models": {
+          "type": "array",
+          "description": "Models blocked for this provider config even if matched by allowed_models. Use [\"*\"] to block all; empty array blocks none.",
+          "items": {
+            "type": "string"
+          }
+        },
         "rate_limit_id": {
           "type": "string",
           "description": "Associated rate limit ID"
@@ -2816,7 +2832,7 @@
         },
         "auth_type": {
           "type": "string",
-          "enum": ["none", "headers", "oauth", "per_user_oauth"],
+          "enum": ["none", "headers", "oauth", "per_user_oauth", "per_user_headers"],
           "description": "Authentication type for MCP connection"
         },
         "oauth_config_id": {
@@ -2830,6 +2846,13 @@
             "type": "string"
           }
         },
+        "per_user_header_keys": {
+          "type": "array",
+          "description": "Header names each caller must supply when auth_type is 'per_user_headers'. Required (non-empty) for that auth type.",
+          "items": {
+            "type": "string"
+          }
+        },
         "stdio_config": {
           "type": "object",
           "properties": {