maximhq · akshaydeo · Jan 7, 2026 · Jan 6, 2026
diff --git a/.gitignore b/.gitignore
@@ -54,4 +54,5 @@ test-reports
 
 # Cursor specific
 .cursor/
-build/
+build/
+target/
diff --git a/core/bifrost.go b/core/bifrost.go
diff --git a/core/bifrost_test.go b/core/bifrost_test.go
@@ -51,7 +51,7 @@ func createBifrostError(message string, statusCode *int, errorType *string, isBi
 // Test executeRequestWithRetries - success scenarios
 func TestExecuteRequestWithRetries_SuccessScenarios(t *testing.T) {
 	config := createTestConfig(3, 100*time.Millisecond, 1*time.Second)
-	ctx := context.Background()
+	ctx := schemas.NewBifrostContext(context.Background(), schemas.NoDeadline)
 
 	// Test immediate success
 	t.Run("ImmediateSuccess", func(t *testing.T) {
@@ -62,7 +62,7 @@ func TestExecuteRequestWithRetries_SuccessScenarios(t *testing.T) {
 		}
 
 		result, err := executeRequestWithRetries(
-			&ctx,
+			ctx,
 			config,
 			handler,
 			schemas.ChatCompletionRequest,
@@ -96,7 +96,7 @@ func TestExecuteRequestWithRetries_SuccessScenarios(t *testing.T) {
 		}
 
 		result, err := executeRequestWithRetries(
-			&ctx,
+			ctx,
 			config,
 			handler,
 			schemas.ChatCompletionRequest,
@@ -120,7 +120,7 @@ func TestExecuteRequestWithRetries_SuccessScenarios(t *testing.T) {
 // Test executeRequestWithRetries - retry limits
 func TestExecuteRequestWithRetries_RetryLimits(t *testing.T) {
 	config := createTestConfig(2, 100*time.Millisecond, 1*time.Second)
-	ctx := context.Background()
+	ctx := schemas.NewBifrostContext(context.Background(), schemas.NoDeadline)
 	t.Run("ExceedsMaxRetries", func(t *testing.T) {
 		callCount := 0
 		handler := func() (string, *schemas.BifrostError) {
@@ -130,7 +130,7 @@ func TestExecuteRequestWithRetries_RetryLimits(t *testing.T) {
 		}
 
 		result, err := executeRequestWithRetries(
-			&ctx,
+			ctx,
 			config,
 			handler,
 			schemas.ChatCompletionRequest,
@@ -161,7 +161,7 @@ func TestExecuteRequestWithRetries_RetryLimits(t *testing.T) {
 // Test executeRequestWithRetries - non-retryable errors
 func TestExecuteRequestWithRetries_NonRetryableErrors(t *testing.T) {
 	config := createTestConfig(3, 100*time.Millisecond, 1*time.Second)
-	ctx := context.Background()
+	ctx := schemas.NewBifrostContext(context.Background(), schemas.NoDeadline)
 	testCases := []struct {
 		name  string
 		error *schemas.BifrostError
@@ -193,7 +193,7 @@ func TestExecuteRequestWithRetries_NonRetryableErrors(t *testing.T) {
 			}
 
 			result, err := executeRequestWithRetries(
-				&ctx,
+				ctx,
 				config,
 				handler,
 				schemas.ChatCompletionRequest,
@@ -218,7 +218,7 @@ func TestExecuteRequestWithRetries_NonRetryableErrors(t *testing.T) {
 // Test executeRequestWithRetries - retryable conditions
 func TestExecuteRequestWithRetries_RetryableConditions(t *testing.T) {
 	config := createTestConfig(1, 100*time.Millisecond, 1*time.Second)
-	ctx := context.Background()
+	ctx := schemas.NewBifrostContext(context.Background(), schemas.NoDeadline)
 	testCases := []struct {
 		name  string
 		error *schemas.BifrostError
@@ -266,7 +266,7 @@ func TestExecuteRequestWithRetries_RetryableConditions(t *testing.T) {
 			}
 
 			result, err := executeRequestWithRetries(
-				&ctx,
+				ctx,
 				config,
 				handler,
 				schemas.ChatCompletionRequest,
@@ -476,7 +476,7 @@ func TestIsRateLimitError_EdgeCases(t *testing.T) {
 // Test retry logging and attempt counting
 func TestExecuteRequestWithRetries_LoggingAndCounting(t *testing.T) {
 	config := createTestConfig(2, 50*time.Millisecond, 1*time.Second)
-	ctx := context.Background()
+	ctx := schemas.NewBifrostContext(context.Background(), schemas.NoDeadline)
 
 	// Capture calls and timing for verification
 	var attemptCounts []int
@@ -495,7 +495,7 @@ func TestExecuteRequestWithRetries_LoggingAndCounting(t *testing.T) {
 	}
 
 	result, err := executeRequestWithRetries(
-		&ctx,
+		ctx,
 		config,
 		handler,
 		schemas.ChatCompletionRequest,
@@ -638,7 +638,7 @@ func (ma *MockAccount) GetConfigForProvider(provider schemas.ModelProvider) (*sc
 	return nil, fmt.Errorf("provider %s not configured", provider)
 }
 
-func (ma *MockAccount) GetKeysForProvider(ctx *context.Context, provider schemas.ModelProvider) ([]schemas.Key, error) {
+func (ma *MockAccount) GetKeysForProvider(ctx *schemas.BifrostContext, provider schemas.ModelProvider) ([]schemas.Key, error) {
 	if keys, exists := ma.keys[provider]; exists {
 		return keys, nil
 	}
@@ -653,7 +653,7 @@ func TestUpdateProvider(t *testing.T) {
 		account.AddProvider(schemas.OpenAI, 5, 1000)
 
 		// Initialize Bifrost
-		ctx := context.Background()
+		ctx := schemas.NewBifrostContext(context.Background(), schemas.NoDeadline)
 		bifrost, err := Init(ctx, schemas.BifrostConfig{
 			Account: account,
 			Logger:  NewDefaultLogger(schemas.LogLevelError), // Keep tests quiet

diff --git a/core/chatbot_test.go b/core/chatbot_test.go
@@ -56,7 +56,7 @@ func (account *ComprehensiveTestAccount) GetConfiguredProviders() ([]schemas.Mod
 }
 
 // GetKeysForProvider returns the API keys and associated models for a given provider.
-func (account *ComprehensiveTestAccount) GetKeysForProvider(ctx *context.Context, providerKey schemas.ModelProvider) ([]schemas.Key, error) {
+func (account *ComprehensiveTestAccount) GetKeysForProvider(ctx *schemas.BifrostContext, providerKey schemas.ModelProvider) ([]schemas.Key, error) {
 	switch providerKey {
 	case schemas.OpenAI:
 		return []schemas.Key{
@@ -306,8 +306,8 @@ func (s *ChatSession) getAvailableProviders() []schemas.ModelProvider {
 			availableProviders = append(availableProviders, provider)
 			continue
 		}
-		ctx := context.Background()
-		keys, err := s.account.GetKeysForProvider(&ctx, provider)
+		ctx := schemas.NewBifrostContext(context.Background(), schemas.NoDeadline)
+		keys, err := s.account.GetKeysForProvider(ctx, provider)
 		if err == nil && len(keys) > 0 && keys[0].Value != "" {
 			availableProviders = append(availableProviders, provider)
 		}
@@ -317,8 +317,8 @@ func (s *ChatSession) getAvailableProviders() []schemas.ModelProvider {
 
 // getAvailableModels returns available models for a given provider
 func (s *ChatSession) getAvailableModels(provider schemas.ModelProvider) []string {
-	ctx := context.Background()
-	keys, err := s.account.GetKeysForProvider(&ctx, provider)
+	ctx := schemas.NewBifrostContext(context.Background(), schemas.NoDeadline)
+	keys, err := s.account.GetKeysForProvider(ctx, provider)
 	if err != nil || len(keys) == 0 {
 		return []string{}
 	}
@@ -486,7 +486,7 @@ func (s *ChatSession) SendMessage(message string) (string, error) {
 	stopChan, wg := startLoader()
 
 	// Send request
-	response, err := s.client.ChatCompletionRequest(context.Background(), request)
+	response, err := s.client.ChatCompletionRequest(schemas.NewBifrostContext(context.Background(), schemas.NoDeadline), request)
 
 	// Stop loading animation
 	stopLoader(stopChan, wg)
@@ -563,7 +563,7 @@ func (s *ChatSession) handleToolCalls(assistantMessage schemas.ChatMessage) (str
 		stopChan, wg := startLoader()
 
 		// Execute the tool using Bifrost's integrated MCP functionality
-		toolResult, err := s.client.ExecuteChatMCPTool(context.Background(), toolCall)
+		toolResult, err := s.client.ExecuteChatMCPTool(schemas.NewBifrostContext(context.Background(), schemas.NoDeadline), toolCall)
 
 		// Stop loading animation
 		stopLoader(stopChan, wg)
@@ -638,7 +638,7 @@ func (s *ChatSession) synthesizeToolResults() (string, error) {
 	stopChan, wg := startLoader()
 
 	// Send synthesis request
-	synthesisResponse, err := s.client.ChatCompletionRequest(context.Background(), synthesisRequest)
+	synthesisResponse, err := s.client.ChatCompletionRequest(schemas.NewBifrostContext(context.Background(), schemas.NoDeadline), synthesisRequest)
 
 	// Stop loading animation
 	stopLoader(stopChan, wg)

diff --git a/core/internal/testutil/account.go b/core/internal/testutil/account.go
@@ -4,7 +4,6 @@
 package testutil
 
 import (
-	"context"
 	"fmt"
 	"os"
 	"time"
@@ -123,7 +122,7 @@ func (account *ComprehensiveTestAccount) GetConfiguredProviders() ([]schemas.Mod
 }
 
 // GetKeysForProvider returns the API keys and associated models for a given provider.
-func (account *ComprehensiveTestAccount) GetKeysForProvider(ctx *context.Context, providerKey schemas.ModelProvider) ([]schemas.Key, error) {
+func (account *ComprehensiveTestAccount) GetKeysForProvider(ctx *schemas.BifrostContext, providerKey schemas.ModelProvider) ([]schemas.Key, error) {
 	switch providerKey {
 	case schemas.OpenAI:
 		return []schemas.Key{

diff --git a/core/internal/testutil/automatic_function_calling.go b/core/internal/testutil/automatic_function_calling.go
@@ -1,7 +1,6 @@
 package testutil
 
 import (
-	"context"
 	"os"
 	"strings"
 	"testing"
@@ -11,7 +10,7 @@ import (
 )
 
 // RunAutomaticFunctionCallingTest executes the automatic function calling test scenario using dual API testing framework
-func RunAutomaticFunctionCallingTest(t *testing.T, client *bifrost.Bifrost, ctx context.Context, testConfig ComprehensiveTestConfig) {
+func RunAutomaticFunctionCallingTest(t *testing.T, client *bifrost.Bifrost, ctx *schemas.BifrostContext, testConfig ComprehensiveTestConfig) {
 	if !testConfig.Scenarios.AutomaticFunctionCall {
 		t.Logf("Automatic function calling not supported for provider %s", testConfig.Provider)
 		return

diff --git a/core/internal/testutil/batch.go b/core/internal/testutil/batch.go
@@ -2,15 +2,14 @@
 package testutil
 
 import (
-	"context"
 	"testing"
 
 	bifrost "github.com/maximhq/bifrost/core"
 	"github.com/maximhq/bifrost/core/schemas"
 )
 
 // RunBatchCreateTest tests the batch create functionality
-func RunBatchCreateTest(t *testing.T, client *bifrost.Bifrost, ctx context.Context, testConfig ComprehensiveTestConfig) {
+func RunBatchCreateTest(t *testing.T, client *bifrost.Bifrost, ctx *schemas.BifrostContext, testConfig ComprehensiveTestConfig) {
 	if !testConfig.Scenarios.BatchCreate {
 		t.Logf("[SKIPPED] Batch Create: Not supported by provider %s", testConfig.Provider)
 		return
@@ -65,7 +64,7 @@ func RunBatchCreateTest(t *testing.T, client *bifrost.Bifrost, ctx context.Conte
 }
 
 // RunBatchListTest tests the batch list functionality
-func RunBatchListTest(t *testing.T, client *bifrost.Bifrost, ctx context.Context, testConfig ComprehensiveTestConfig) {
+func RunBatchListTest(t *testing.T, client *bifrost.Bifrost, ctx *schemas.BifrostContext, testConfig ComprehensiveTestConfig) {
 	if !testConfig.Scenarios.BatchList {
 		t.Logf("[SKIPPED] Batch List: Not supported by provider %s", testConfig.Provider)
 		return
@@ -100,7 +99,7 @@ func RunBatchListTest(t *testing.T, client *bifrost.Bifrost, ctx context.Context
 }
 
 // RunBatchRetrieveTest tests the batch retrieve functionality
-func RunBatchRetrieveTest(t *testing.T, client *bifrost.Bifrost, ctx context.Context, testConfig ComprehensiveTestConfig) {
+func RunBatchRetrieveTest(t *testing.T, client *bifrost.Bifrost, ctx *schemas.BifrostContext, testConfig ComprehensiveTestConfig) {
 	if !testConfig.Scenarios.BatchRetrieve {
 		t.Logf("[SKIPPED] Batch Retrieve: Not supported by provider %s", testConfig.Provider)
 		return
@@ -173,7 +172,7 @@ func RunBatchRetrieveTest(t *testing.T, client *bifrost.Bifrost, ctx context.Con
 }
 
 // RunBatchCancelTest tests the batch cancel functionality
-func RunBatchCancelTest(t *testing.T, client *bifrost.Bifrost, ctx context.Context, testConfig ComprehensiveTestConfig) {
+func RunBatchCancelTest(t *testing.T, client *bifrost.Bifrost, ctx *schemas.BifrostContext, testConfig ComprehensiveTestConfig) {
 	if !testConfig.Scenarios.BatchCancel {
 		t.Logf("[SKIPPED] Batch Cancel: Not supported by provider %s", testConfig.Provider)
 		return
@@ -241,7 +240,7 @@ func RunBatchCancelTest(t *testing.T, client *bifrost.Bifrost, ctx context.Conte
 }
 
 // RunBatchResultsTest tests the batch results functionality
-func RunBatchResultsTest(t *testing.T, client *bifrost.Bifrost, ctx context.Context, testConfig ComprehensiveTestConfig) {
+func RunBatchResultsTest(t *testing.T, client *bifrost.Bifrost, ctx *schemas.BifrostContext, testConfig ComprehensiveTestConfig) {
 	if !testConfig.Scenarios.BatchResults {
 		t.Logf("[SKIPPED] Batch Results: Not supported by provider %s", testConfig.Provider)
 		return
@@ -274,7 +273,7 @@ func RunBatchResultsTest(t *testing.T, client *bifrost.Bifrost, ctx context.Cont
 }
 
 // RunBatchUnsupportedTest tests that unsupported providers return appropriate errors
-func RunBatchUnsupportedTest(t *testing.T, client *bifrost.Bifrost, ctx context.Context, testConfig ComprehensiveTestConfig) {
+func RunBatchUnsupportedTest(t *testing.T, client *bifrost.Bifrost, ctx *schemas.BifrostContext, testConfig ComprehensiveTestConfig) {
 	// Only run this test for providers that don't support batch
 	if testConfig.Scenarios.BatchCreate || testConfig.Scenarios.BatchList ||
 		testConfig.Scenarios.BatchRetrieve || testConfig.Scenarios.BatchCancel ||
@@ -331,7 +330,7 @@ func RunBatchUnsupportedTest(t *testing.T, client *bifrost.Bifrost, ctx context.
 // ============================================================================
 
 // RunFileUploadTest tests the file upload functionality
-func RunFileUploadTest(t *testing.T, client *bifrost.Bifrost, ctx context.Context, testConfig ComprehensiveTestConfig) {
+func RunFileUploadTest(t *testing.T, client *bifrost.Bifrost, ctx *schemas.BifrostContext, testConfig ComprehensiveTestConfig) {
 	if !testConfig.Scenarios.FileUpload {
 		t.Logf("[SKIPPED] File Upload: Not supported by provider %s", testConfig.Provider)
 		return
@@ -378,7 +377,7 @@ func RunFileUploadTest(t *testing.T, client *bifrost.Bifrost, ctx context.Contex
 }
 
 // RunFileListTest tests the file list functionality
-func RunFileListTest(t *testing.T, client *bifrost.Bifrost, ctx context.Context, testConfig ComprehensiveTestConfig) {
+func RunFileListTest(t *testing.T, client *bifrost.Bifrost, ctx *schemas.BifrostContext, testConfig ComprehensiveTestConfig) {
 	if !testConfig.Scenarios.FileList {
 		t.Logf("[SKIPPED] File List: Not supported by provider %s", testConfig.Provider)
 		return
@@ -414,7 +413,7 @@ func RunFileListTest(t *testing.T, client *bifrost.Bifrost, ctx context.Context,
 }
 
 // RunFileRetrieveTest tests the file retrieve functionality
-func RunFileRetrieveTest(t *testing.T, client *bifrost.Bifrost, ctx context.Context, testConfig ComprehensiveTestConfig) {
+func RunFileRetrieveTest(t *testing.T, client *bifrost.Bifrost, ctx *schemas.BifrostContext, testConfig ComprehensiveTestConfig) {
 	if !testConfig.Scenarios.FileRetrieve {
 		t.Logf("[SKIPPED] File Retrieve: Not supported by provider %s", testConfig.Provider)
 		return
@@ -477,7 +476,7 @@ func RunFileRetrieveTest(t *testing.T, client *bifrost.Bifrost, ctx context.Cont
 }
 
 // RunFileDeleteTest tests the file delete functionality
-func RunFileDeleteTest(t *testing.T, client *bifrost.Bifrost, ctx context.Context, testConfig ComprehensiveTestConfig) {
+func RunFileDeleteTest(t *testing.T, client *bifrost.Bifrost, ctx *schemas.BifrostContext, testConfig ComprehensiveTestConfig) {
 	if !testConfig.Scenarios.FileDelete {
 		t.Logf("[SKIPPED] File Delete: Not supported by provider %s", testConfig.Provider)
 		return
@@ -540,7 +539,7 @@ func RunFileDeleteTest(t *testing.T, client *bifrost.Bifrost, ctx context.Contex
 }
 
 // RunFileContentTest tests the file content download functionality
-func RunFileContentTest(t *testing.T, client *bifrost.Bifrost, ctx context.Context, testConfig ComprehensiveTestConfig) {
+func RunFileContentTest(t *testing.T, client *bifrost.Bifrost, ctx *schemas.BifrostContext, testConfig ComprehensiveTestConfig) {
 	if !testConfig.Scenarios.FileContent {
 		t.Logf("[SKIPPED] File Content: Not supported by provider %s", testConfig.Provider)
 		return
@@ -604,7 +603,7 @@ func RunFileContentTest(t *testing.T, client *bifrost.Bifrost, ctx context.Conte
 }
 
 // RunFileUnsupportedTest tests that unsupported providers return appropriate errors for file operations
-func RunFileUnsupportedTest(t *testing.T, client *bifrost.Bifrost, ctx context.Context, testConfig ComprehensiveTestConfig) {
+func RunFileUnsupportedTest(t *testing.T, client *bifrost.Bifrost, ctx *schemas.BifrostContext, testConfig ComprehensiveTestConfig) {
 	// Only run this test for providers that don't support any file operations
 	if testConfig.Scenarios.FileUpload || testConfig.Scenarios.FileList ||
 		testConfig.Scenarios.FileRetrieve || testConfig.Scenarios.FileDelete ||
@@ -647,7 +646,7 @@ func RunFileUnsupportedTest(t *testing.T, client *bifrost.Bifrost, ctx context.C
 }
 
 // RunFileAndBatchIntegrationTest tests the integration between file upload and batch create
-func RunFileAndBatchIntegrationTest(t *testing.T, client *bifrost.Bifrost, ctx context.Context, testConfig ComprehensiveTestConfig) {
+func RunFileAndBatchIntegrationTest(t *testing.T, client *bifrost.Bifrost, ctx *schemas.BifrostContext, testConfig ComprehensiveTestConfig) {
 	// Skip if file-based batch input is not supported
 	if !testConfig.Scenarios.FileBatchInput {
 		t.Logf("[SKIPPED] File and Batch Integration: FileBatchInput=%v for provider %s",

diff --git a/core/internal/testutil/chat_audio.go b/core/internal/testutil/chat_audio.go
@@ -1,7 +1,6 @@
 package testutil
 
 import (
-	"context"
 	"os"
 	"strings"
 	"testing"
@@ -11,7 +10,7 @@ import (
 )
 
 // RunChatAudioTest executes the chat audio test scenario
-func RunChatAudioTest(t *testing.T, client *bifrost.Bifrost, ctx context.Context, testConfig ComprehensiveTestConfig) {
+func RunChatAudioTest(t *testing.T, client *bifrost.Bifrost, ctx *schemas.BifrostContext, testConfig ComprehensiveTestConfig) {
 	if !testConfig.Scenarios.ChatAudio || testConfig.ChatAudioModel == "" {
 		t.Logf("Chat audio not supported for provider %s", testConfig.Provider)
 		return
@@ -149,7 +148,7 @@ func RunChatAudioTest(t *testing.T, client *bifrost.Bifrost, ctx context.Context
 }
 
 // RunChatAudioStreamTest executes the chat audio streaming test scenario
-func RunChatAudioStreamTest(t *testing.T, client *bifrost.Bifrost, ctx context.Context, testConfig ComprehensiveTestConfig) {
+func RunChatAudioStreamTest(t *testing.T, client *bifrost.Bifrost, ctx *schemas.BifrostContext, testConfig ComprehensiveTestConfig) {
 	if !testConfig.Scenarios.ChatAudio || testConfig.ChatAudioModel == "" {
 		t.Logf("Chat audio streaming not supported for provider %s", testConfig.Provider)
 		return

diff --git a/core/internal/testutil/chat_completion_stream.go b/core/internal/testutil/chat_completion_stream.go
@@ -13,7 +13,7 @@ import (
 )
 
 // RunChatCompletionStreamTest executes the chat completion stream test scenario
-func RunChatCompletionStreamTest(t *testing.T, client *bifrost.Bifrost, ctx context.Context, testConfig ComprehensiveTestConfig) {
+func RunChatCompletionStreamTest(t *testing.T, client *bifrost.Bifrost, ctx *schemas.BifrostContext, testConfig ComprehensiveTestConfig) {
 	if !testConfig.Scenarios.CompletionStream {
 		t.Logf("Chat completion stream not supported for provider %s", testConfig.Provider)
 		return

diff --git a/core/internal/testutil/complete_end_to_end.go b/core/internal/testutil/complete_end_to_end.go
@@ -1,18 +1,16 @@
 package testutil
 
 import (
-	"context"
 	"os"
 	"strings"
 	"testing"
 
-
 	bifrost "github.com/maximhq/bifrost/core"
 	"github.com/maximhq/bifrost/core/schemas"
 )
 
 // RunCompleteEnd2EndTest executes the complete end-to-end test scenario
-func RunCompleteEnd2EndTest(t *testing.T, client *bifrost.Bifrost, ctx context.Context, testConfig ComprehensiveTestConfig) {
+func RunCompleteEnd2EndTest(t *testing.T, client *bifrost.Bifrost, ctx *schemas.BifrostContext, testConfig ComprehensiveTestConfig) {
 	if !testConfig.Scenarios.CompleteEnd2End {
 		t.Logf("Complete end-to-end not supported for provider %s", testConfig.Provider)
 		return
-Original file line number
+Diff line change
@@ Expand Up / @@ -54,4 +54,5 @@ test-reports @@
     # Cursor specific
     .cursor/
-    build/
+    build/
+    target/