docs: documentation fixes and updates

Author: Pratham-Mishra04

core/schemas/chatcompletions.go

@@ -16,6 +16,7 @@ type BifrostChatRequest struct {
 	Fallbacks []Fallback      `json:"fallbacks,omitempty"`
 }
 
+// BifrostChatResponse represents the complete result from a chat completion request.
 type BifrostChatResponse struct {
 	ID                string                     `json:"id"`
 	Choices           []BifrostResponseChoice    `json:"choices"`
@@ -203,12 +204,12 @@ type ChatToolFunction struct {
 
 // ToolFunctionParameters represents the parameters for a function definition.
 type ToolFunctionParameters struct {
-	Type                 string                 `json:"type"`                           // Type of the parameters
-	Description          *string                `json:"description,omitempty"`          // Description of the parameters
-	Required             []string               `json:"required,omitempty"`             // Required parameter names
+	Type                 string                  `json:"type"`                           // Type of the parameters
+	Description          *string                 `json:"description,omitempty"`          // Description of the parameters
+	Required             []string                `json:"required,omitempty"`             // Required parameter names
 	Properties           *map[string]interface{} `json:"properties,omitempty"`           // Parameter properties
-	Enum                 []string               `json:"enum,omitempty"`                 // Enum values for the parameters
-	AdditionalProperties *bool                  `json:"additionalProperties,omitempty"` // Whether to allow additional properties
+	Enum                 []string                `json:"enum,omitempty"`                 // Enum values for the parameters
+	AdditionalProperties *bool                   `json:"additionalProperties,omitempty"` // Whether to allow additional properties
 }
 
 type ChatToolCustom struct {

docs/architecture/framework/pricing.mdx

@@ -109,9 +109,6 @@ Calculate costs from a Bifrost response:
 // Calculate cost for a completed request
 cost := pricingManager.CalculateCost(
     result,                           // *schemas.BifrostResponse
-    schemas.OpenAI,                   // provider
-    "gpt-4",                          // model
-    schemas.ChatCompletionRequest,    // request type
 )
 
 logger.Info("Request cost: $%.6f", cost)
@@ -147,9 +144,6 @@ For workflows that implement semantic caching, use cache-aware cost calculation:
 // This automatically handles cache hits/misses and embedding costs
 cost := pricingManager.CalculateCostWithCacheDebug(
     result,                          // *schemas.BifrostResponse with cache debug info
-    schemas.Anthropic,               // provider
-    "claude-3-sonnet",               // model
-    schemas.ChatCompletionRequest,   // request type
 )
 
 // Cache hits return 0 for direct hits, embedding cost for semantic matches

docs/docs.json

@@ -105,9 +105,18 @@
                   }
                 ]
               },
-              "features/governance",
               "features/semantic-caching",
               "features/custom-providers",
+              {
+                "group": "Governance",
+                "icon": "user-lock",
+                "pages": [
+                  "features/governance/virtual-keys",
+                  "features/governance/budget-and-limits",
+                  "features/governance/routing",
+                  "features/governance/mcp-tools"
+                ]
+              },
               {
                 "group": "Plugins",
                 "icon": "puzzle-piece",
@@ -124,7 +133,7 @@
             "pages": [
               "enterprise/guardrails",
               "enterprise/clustering",
-              "enterprise/governance",
+              "enterprise/advanced-governance",
               "enterprise/mcp-with-fa",
               "enterprise/vault-support",
               "enterprise/invpc-deployments",

docs/enterprise/governance.mdx renamed to docs/enterprise/advanced-governance.mdx

@@ -1,5 +1,5 @@
 ---
-title: "Governance"
+title: "Advanced Governance"
 description: "Advanced governance features with enhanced security, compliance reporting, audit trails, and enterprise-grade access controls for large-scale deployments."
 icon: "shield-check"
 ---

docs/features/governance/budget-and-limits.mdx

@@ -0,0 +1,127 @@
+---
+title: "Budget and Limits"
+description: "Enterprise-grade budget management and cost control with hierarchical budget allocation through virtual keys, teams, and customers."
+icon: "money-bills"
+---
+
+## Overview
+
+Budgeting and rate limiting are a core feature of Bifrost's governance system managed through [Virtual Keys](./virtual-keys).
+
+Bifrost's budget management system provides comprehensive cost control and financial governance for enterprise AI deployments. It operates through a **hierarchical budget structure** that enables granular cost management, usage tracking, and financial oversight across your entire organization.
+
+**Core Hierarchy:**
+```
+Customer (has independent budget)
+    ↓ (one-to-many)
+Team (has independent budget) 
+    ↓ (one-to-many)
+Virtual Key (has independent budget + rate limits)
+
+OR
+
+Customer (has independent budget)
+    ↓ (direct attachment)
+Virtual Key (has independent budget + rate limits)
+
+OR
+
+Virtual Key (standalone - has independent budget + rate limits)
+```
+
+**Key Capabilities:**
+- **Virtual Keys** - Primary access control via `x-bf-vk` header (exclusive team OR customer attachment)
+- **Budget Management** - Independent budget limits at each hierarchy level with cumulative checking
+- **Rate Limiting** - Request and token-based throttling (VK-level only)
+- **Model/Provider Filtering** - Granular access control per virtual key  
+- **Usage Tracking** - Real-time monitoring and audit trails
+- **Audit Headers** - Optional team and customer identification
+
+---
+
+## Budget Management
+
+### Cost Calculation
+
+Bifrost automatically calculates costs based on:
+- **Provider Pricing** - Real-time model pricing data
+- **Token Usage** - Input + output tokens from API responses
+- **Request Type** - Different pricing for chat, text, embedding, speech, transcription
+- **Cache Status** - Reduced costs for cached responses
+- **Batch Operations** - Volume discounts for batch requests
+
+All cost calculation details are covered in [Architecture > Framework > Pricing](../../architecture/framework/pricing).
+
+### Budget Checking Flow
+
+When a request is made with a a virtual key, Bifrost checks **all applicable budgets independently** in the hierarchy. Each budget must have sufficient remaining balance for the request to proceed.
+
+**Checking Sequence:**
+
+**For VK → Team → Customer:**
+```
+1. ✓ VK Budget (if VK has budget)
+2. ✓ Team Budget (if VK's team has budget)  
+3. ✓ Customer Budget (if team's customer has budget)
+```
+
+**For VK → Customer (direct):**
+```
+1. ✓ VK Budget (if VK has budget)
+2. ✓ Customer Budget (if VK's customer has budget)
+```
+
+**For Standalone VK:**
+```
+1. ✓ VK Budget (if VK has budget)
+```
+
+**Important Notes:**
+- **All applicable budgets must pass** - any single budget failure blocks the request
+- **Budgets are independent** - each tracks its own usage and limits
+- **Costs are deducted from all applicable budgets** - same cost applied to each level
+- **Rate limits checked only at VK level** - teams and customers have no rate limits
+
+**Example:**
+- VK budget: $9/$10 remaining ✓
+- Team budget: $15/$20 remaining ✓  
+- Customer budget: $45/$50 remaining ✓
+- **Result: Allowed** (no budget is exceeded)
+- After request: 
+    - Request cost: $2 
+    - Updated VK=$11/$10, Team=$17/$20, Customer=$47/$50
+    - Then the next request will be blocked.
+
+## Rate Limiting
+
+Rate limits protect your system from abuse and manage traffic by setting thresholds on request frequency and token usage over a specific time window. Unlike budgets, rate limits are configured **exclusively at the Virtual Key level**.
+
+Bifrost supports two types of rate limits that work in parallel:
+- **Request Limits**: Control the maximum number of API calls that can be made within a set duration (e.g., 100 requests per minute).
+- **Token Limits**: Control the maximum number of tokens (prompt + completion) that can be processed within a set duration (e.g., 50,000 tokens per hour).
+
+For a request to be allowed, it must pass both the request limit and the token limit checks.
+
+## Reset Durations
+
+Budgets and rate limits support flexible reset durations:
+
+**Format Examples:**
+- `1m` - 1 minute
+- `5m` - 5 minutes  
+- `1h` - 1 hour
+- `1d` - 1 day
+- `1w` - 1 week
+- `1M` - 1 month
+
+**Common Patterns:**
+- **Rate Limits**: `1m`, `1h`, `1d` for request throttling
+- **Budgets**: `1d`, `1w`, `1M` for cost control
+
+---
+
+## Next Steps
+
+- **[Routing](./routing)** - Direct requests to specific AI models, providers, and keys using Virtual Keys.
+- **[MCP Tool Filtering](./mcp-tools)** - Manage MCP clients/tools for virtual keys.
+- **[Tracing](../tracing)** - Audit trails and request tracking

docs/features/governance/mcp-tools.mdx

@@ -0,0 +1,160 @@
+---
+title: "MCP Tool Filtering"
+description: "Control which MCP tools are available for each Virtual Key."
+icon: "grid-2"
+---
+
+## Overview
+
+MCP Tool Filtering allows you to control which tools are available to AI models on a per-request basis using Virtual Keys (VKs). By configuring a VirtualKey, you can create a strict allow-list of MCP clients and tools, ensuring that only approved tools can be executed.
+
+Make sure you have at least one MCP client set up. Read more about it [here](../mcp).
+
+## How It Works
+
+The filtering logic is determined by the Virtual Key's configuration:
+
+1.  **No MCP Configuration on Virtual Key (Default)**
+    - If a Virtual Key has no specific MCP configurations, all tools from all enabled MCP clients are available by default.
+    - In this state, a user can still manually filter tools for a single request by passing the `x-bf-mcp-include-tools` header.
+
+2.  **With MCP Configuration on Virtual Key**
+    - When you configure MCP clients on a Virtual Key, its settings take full precedence.
+    - Bifrost automatically generates an `x-bf-mcp-include-tools` header based on your VK configuration. This acts as a strict allow-list for the request.
+    - This generated header **overrides** any `x-bf-mcp-include-tools` header that might have been sent manually with the request.
+
+For each MCP client associated with a Virtual Key, you can specify the allowed tools:
+- **Select specific tools**: Only the chosen tools from that client will be available.
+- **Use `*` wildcard**: All available tools from that client will be permitted.
+- **Leave tool list empty**: All tools from that client will be **blocked**.
+- **Do not configure a client**: All tools from that client will be **blocked** (if other clients are configured).
+
+## Setting MCP Tool Restrictions
+
+<Tabs group="mcp-tool-restrictions">
+<Tab title="Web UI">
+
+You can configure which tools a Virtual Key has access to via the UI.
+
+1.  Go to **Virtual Keys** page.
+2.  Create/Edit virtual key
+![Virtual Key MCP Tool Restrictions](../../media/ui-virtual-key-mcp-filter.png)
+3.  In **MCP Client Configurations** section, add the MCP client you want to restrict the VK to
+4.  Add the tools you want to restrict the VK to, or leave it blank to allow all tools for this client
+5.  Click on the **Save** button
+
+</Tab>
+<Tab title="API">
+
+You can configure this via the REST API when creating (`POST`) or updating (`PUT`) a virtual key.
+
+**Create Virtual Key:**
+```bash
+curl -X POST http://localhost:8080/api/governance/virtual-keys \
+  -H "Content-Type: application/json" \
+  -d '{
+    "name": "vk-for-billing-support",
+    "mcp_configs": [
+      {
+        "mcp_client_name": "billing-client",
+        "tools_to_execute": ["check-status"]
+      },
+      {
+        "mcp_client_name": "support-client",
+        "tools_to_execute": ["*"]
+      }
+    ]
+  }'
+```
+
+**Update Virtual Key:**
+```bash
+curl -X PUT http://localhost:8080/api/governance/virtual-keys/{vk_id} \
+  -H "Content-Type: application/json" \
+  -d '{
+    "mcp_configs": [
+      {
+        "mcp_client_name": "billing-client",
+        "tools_to_execute": ["check-status"]
+      },
+      {
+        "mcp_client_name": "support-client",
+        "tools_to_execute": ["*"]
+      }
+    ]
+  }'
+```
+
+**Behavior:**
+- The virtual key can only access the `check-status` tool from `billing-client`.
+- It can access all tools from `support-client`.
+- Any other MCP client is implicitly blocked for this key.
+
+</Tab>
+
+<Tab title="config.json">
+
+You can also define MCP tool restrictions directly in your `config.json` file. The `mcp_configs` array under a virtual key should reference the MCP client by name.
+
+```json
+{
+  "governance": {
+    "virtual_keys": [
+      {
+        "id": "vk-billing-support-only",
+        "name": "VK for Billing and Support",
+        "mcp_configs": [
+          {
+            "mcp_client_name": "billing-client",
+            "tools_to_execute": ["check-status"]
+          },
+          {
+            "mcp_client_name": "support-client",
+            "tools_to_execute": ["*"]
+          }
+        ]
+      }
+    ]
+  }
+}
+```
+</Tab>
+</Tabs>
+
+## Example Scenario
+
+**Available MCP Clients & Tools:**
+- **`billing-client`**: with tools `[create-invoice, check-status]`
+- **`support-client`**: with tools `[create-ticket, get-faq]`
+
+<Tabs>
+<Tab title="VK with Full Access">
+**Configuration:**
+- `billing-client` -> Allowed Tools: `[*]` (wildcard)
+- `support-client` -> Allowed Tools: `[*]` (wildcard)
+
+**Result:**
+A request with this Virtual Key can access all four tools: `create-invoice`, `check-status`, `create-ticket`, and `get-faq`.
+
+</Tab>
+<Tab title="VK with Partial Access">
+**Configuration:**
+- `billing-client` -> Allowed Tools: `[check-status]`
+- `support-client` -> Not configured
+
+**Result:**
+A request with this Virtual Key can only access the `check-status` tool. All other tools are blocked.
+
+</Tab>
+<Tab title="VK with No Tools">
+**Configuration:**
+- `billing-client` -> Allowed Tools: `[]` (empty list)
+
+**Result:**
+A request with this Virtual Key cannot access any tools. All tools from all clients are blocked.
+</Tab>
+</Tabs>
+
+<Note>
+When a Virtual Key has MCP configurations, it dynamically generates the `x-bf-mcp-include-tools` header. This ensures that the VK's rules are always enforced and will override any manual header sent by the user. Though you can still use `x-bf-mcp-include-clients` header to filter the MCP clients per request.
+</Note>