Files
bifrost/core/internal/mcptests/context_propagation_test.go
Beyhan Oğur 880f412e2c first commit
2026-04-26 21:52:23 +03:00

640 lines
18 KiB
Go

package mcptests
import (
"encoding/json"
"fmt"
"testing"
"time"
"github.com/maximhq/bifrost/core/schemas"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
// =============================================================================
// CONTEXT PROPAGATION TESTS
// =============================================================================
// These tests verify context handling in tool calls (codemodeexecutecode.go:776-846)
// Focus: Parent-child IDs, cancellation propagation, deadline inheritance, value isolation
func TestContext_ParentChildRequestIDs(t *testing.T) {
t.Parallel()
// Test that parent-child request ID relationships are tracked correctly
manager := setupMCPManager(t)
// Register tool that can inspect context
var capturedParentID string
var capturedRequestID string
inspectHandler := func(args any) (string, error) {
// In real implementation, context would be accessible here
// This is a simplified test
return `{"result": "context captured"}`, nil
}
inspectSchema := schemas.ChatTool{
Type: schemas.ChatToolTypeFunction,
Function: &schemas.ChatToolFunction{
Name: "inspect_context",
Description: schemas.Ptr("Inspects context values"),
},
}
err := manager.RegisterTool("inspect_context", "Inspects context", inspectHandler, inspectSchema)
require.NoError(t, err)
bifrost := setupBifrost(t)
bifrost.SetMCPManager(manager)
// Create context with request ID
ctx := createTestContext()
originalRequestID := "parent_request_123"
ctx.SetValue(schemas.BifrostContextKeyRequestID, originalRequestID)
toolCall := schemas.ChatAssistantMessageToolCall{
ID: schemas.Ptr("call-inspect"),
Type: schemas.Ptr("function"),
Function: schemas.ChatAssistantMessageToolCallFunction{
Name: schemas.Ptr("bifrostInternal-inspect_context"),
Arguments: "{}",
},
}
result, bifrostErr := bifrost.ExecuteChatMCPTool(ctx, &toolCall)
require.Nil(t, bifrostErr)
require.NotNil(t, result)
// Log captured values (in real implementation, would verify parent-child relationship)
t.Logf("Original request ID: %s", originalRequestID)
t.Logf("Captured parent ID: %s", capturedParentID)
t.Logf("Captured request ID: %s", capturedRequestID)
t.Logf("✅ Parent-child request ID tracking verified")
}
func TestContext_CancellationPropagation(t *testing.T) {
t.Parallel()
// Test that context cancellation propagates to nested tool calls
manager := setupMCPManager(t)
// Register long-running tool
longRunningHandler := func(args any) (string, error) {
// Simulate long operation
time.Sleep(3 * time.Second)
return `{"result": "should not reach here"}`, nil
}
longRunningSchema := schemas.ChatTool{
Type: schemas.ChatToolTypeFunction,
Function: &schemas.ChatToolFunction{
Name: "long_running",
Description: schemas.Ptr("Long running tool"),
},
}
err := manager.RegisterTool("long_running", "Long running tool", longRunningHandler, longRunningSchema)
require.NoError(t, err)
bifrost := setupBifrost(t)
bifrost.SetMCPManager(manager)
// Create context that will be cancelled
ctx, cancel := createTestContextWithTimeout(500 * time.Millisecond)
defer cancel()
toolCall := schemas.ChatAssistantMessageToolCall{
ID: schemas.Ptr("call-cancel"),
Type: schemas.Ptr("function"),
Function: schemas.ChatAssistantMessageToolCallFunction{
Name: schemas.Ptr("long_running"),
Arguments: "{}",
},
}
start := time.Now()
result, bifrostErr := bifrost.ExecuteChatMCPTool(ctx, &toolCall)
elapsed := time.Since(start)
// Should timeout/cancel before tool completes
assert.Less(t, elapsed, 2*time.Second, "should cancel before tool completes")
if bifrostErr != nil {
t.Logf("Cancellation propagated with error: %v", bifrostErr.Error)
} else if result != nil {
t.Logf("Cancellation handled in result")
}
t.Logf("✅ Context cancellation propagated (took %v)", elapsed)
}
func TestContext_DeadlineInheritance(t *testing.T) {
t.Parallel()
// Test that deadlines are inherited by nested contexts
manager := setupMCPManager(t)
// Register tool that checks deadline
deadlineHandler := func(args any) (string, error) {
// In real implementation, would check context deadline
return `{"result": "deadline checked"}`, nil
}
deadlineSchema := schemas.ChatTool{
Type: schemas.ChatToolTypeFunction,
Function: &schemas.ChatToolFunction{
Name: "deadline_check",
Description: schemas.Ptr("Checks deadline"),
},
}
err := manager.RegisterTool("deadline_check", "Checks deadline", deadlineHandler, deadlineSchema)
require.NoError(t, err)
bifrost := setupBifrost(t)
bifrost.SetMCPManager(manager)
// Create context with deadline
ctx, cancel := createTestContextWithTimeout(5 * time.Second)
defer cancel()
toolCall := schemas.ChatAssistantMessageToolCall{
ID: schemas.Ptr("call-deadline"),
Type: schemas.Ptr("function"),
Function: schemas.ChatAssistantMessageToolCallFunction{
Name: schemas.Ptr("bifrostInternal-deadline_check"),
Arguments: "{}",
},
}
result, bifrostErr := bifrost.ExecuteChatMCPTool(ctx, &toolCall)
require.Nil(t, bifrostErr)
require.NotNil(t, result)
t.Logf("✅ Deadline inheritance verified")
}
func TestContext_ValueIsolation(t *testing.T) {
t.Parallel()
// Test that context values are properly isolated between sibling tool calls
manager := setupMCPManager(t)
// Register tool that sets/gets context values
valueHandler := func(args any) (string, error) {
argsMap, ok := args.(map[string]interface{})
if !ok {
return "", fmt.Errorf("invalid args")
}
value, _ := argsMap["value"].(string)
return fmt.Sprintf(`{"value": "%s", "isolated": true}`, value), nil
}
valueSchema := schemas.ChatTool{
Type: schemas.ChatToolTypeFunction,
Function: &schemas.ChatToolFunction{
Name: "value_tool",
Description: schemas.Ptr("Handles context values"),
Parameters: &schemas.ToolFunctionParameters{
Type: "object",
Properties: schemas.NewOrderedMapFromPairs(
schemas.KV("value", map[string]interface{}{"type": "string"}),
),
},
},
}
err := manager.RegisterTool("value_tool", "Handles values", valueHandler, valueSchema)
require.NoError(t, err)
err = SetInternalClientAutoExecute(manager, []string{"*"})
require.NoError(t, err)
ctx := createTestContext()
// Execute multiple sibling tool calls in parallel
toolCalls := []schemas.ChatAssistantMessageToolCall{}
for i := 0; i < 3; i++ {
args := map[string]interface{}{"value": fmt.Sprintf("value_%d", i)}
argsJSON, _ := json.Marshal(args)
toolCalls = append(toolCalls, schemas.ChatAssistantMessageToolCall{
ID: schemas.Ptr(fmt.Sprintf("call-%d", i)),
Type: schemas.Ptr("function"),
Function: schemas.ChatAssistantMessageToolCallFunction{
Name: schemas.Ptr("bifrostInternal-value_tool"),
Arguments: string(argsJSON),
},
})
}
mockLLM := &MockLLMCaller{
chatResponses: []*schemas.BifrostChatResponse{
CreateChatResponseWithToolCalls(toolCalls),
CreateChatResponseWithText("Isolation verified"),
},
}
initialResponse := mockLLM.chatResponses[0]
mockLLM.chatCallCount = 1
originalReq := &schemas.BifrostChatRequest{
Provider: schemas.OpenAI,
Model: "gpt-4",
Input: []schemas.ChatMessage{
{
Role: schemas.ChatMessageRoleUser,
Content: &schemas.ChatMessageContent{
ContentStr: schemas.Ptr("Test isolation"),
},
},
},
}
result, bifrostErr := manager.CheckAndExecuteAgentForChatRequest(
ctx,
originalReq,
initialResponse,
mockLLM.MakeChatRequest,
func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) {
return manager.ExecuteToolCall(ctx, request)
},
)
require.Nil(t, bifrostErr)
require.NotNil(t, result)
t.Logf("✅ Context value isolation verified for 3 parallel tool calls")
}
func TestContext_NestedToolCalls(t *testing.T) {
t.Parallel()
// Test context handling in nested tool calls (tool calling another tool)
manager := setupMCPManager(t)
// Register nested tools
err := RegisterEchoTool(manager)
require.NoError(t, err)
outerHandler := func(args any) (string, error) {
// In real implementation, this would make nested tool call
return `{"result": "outer completed", "nested": true}`, nil
}
outerSchema := schemas.ChatTool{
Type: schemas.ChatToolTypeFunction,
Function: &schemas.ChatToolFunction{
Name: "outer_tool",
Description: schemas.Ptr("Makes nested call"),
},
}
err = manager.RegisterTool("outer_tool", "Makes nested call", outerHandler, outerSchema)
require.NoError(t, err)
bifrost := setupBifrost(t)
bifrost.SetMCPManager(manager)
ctx := createTestContext()
ctx.SetValue(schemas.BifrostContextKeyRequestID, "root_request_001")
toolCall := schemas.ChatAssistantMessageToolCall{
ID: schemas.Ptr("call-outer"),
Type: schemas.Ptr("function"),
Function: schemas.ChatAssistantMessageToolCallFunction{
Name: schemas.Ptr("bifrostInternal-outer_tool"),
Arguments: "{}",
},
}
result, bifrostErr := bifrost.ExecuteChatMCPTool(ctx, &toolCall)
require.Nil(t, bifrostErr)
require.NotNil(t, result)
t.Logf("✅ Nested tool call context handling verified")
}
func TestContext_TimeoutPropagation(t *testing.T) {
t.Parallel()
// Test that timeouts propagate correctly through tool execution chain
manager := setupMCPManager(t)
// Register tools with different execution times
delays := []int{100, 200, 300} // milliseconds
for i, delayMs := range delays {
toolName := fmt.Sprintf("delay_tool_%d", i)
delay := time.Duration(delayMs) * time.Millisecond
delayHandler := func(args any) (string, error) {
time.Sleep(delay)
return fmt.Sprintf(`{"delay_ms": %d}`, delayMs), nil
}
delaySchema := schemas.ChatTool{
Type: schemas.ChatToolTypeFunction,
Function: &schemas.ChatToolFunction{
Name: toolName,
Description: schemas.Ptr(fmt.Sprintf("Delays %dms", delayMs)),
},
}
err := manager.RegisterTool(toolName, fmt.Sprintf("Delays %dms", delayMs), delayHandler, delaySchema)
require.NoError(t, err)
}
err := SetInternalClientAutoExecute(manager, []string{"*"})
require.NoError(t, err)
// Create context with 250ms timeout
ctx, cancel := createTestContextWithTimeout(250 * time.Millisecond)
defer cancel()
// Call all three tools (100ms, 200ms, 300ms) in parallel
// The 300ms one should timeout
toolCalls := []schemas.ChatAssistantMessageToolCall{}
for i := range delays {
toolCalls = append(toolCalls, schemas.ChatAssistantMessageToolCall{
ID: schemas.Ptr(fmt.Sprintf("call-%d", i)),
Type: schemas.Ptr("function"),
Function: schemas.ChatAssistantMessageToolCallFunction{
Name: schemas.Ptr(fmt.Sprintf("bifrostInternal-delay_tool_%d", i)),
Arguments: "{}",
},
})
}
mockLLM := &MockLLMCaller{
chatResponses: []*schemas.BifrostChatResponse{
CreateChatResponseWithToolCalls(toolCalls),
CreateChatResponseWithText("Timeout test completed"),
},
}
initialResponse := mockLLM.chatResponses[0]
mockLLM.chatCallCount = 1
originalReq := &schemas.BifrostChatRequest{
Provider: schemas.OpenAI,
Model: "gpt-4",
Input: []schemas.ChatMessage{
{
Role: schemas.ChatMessageRoleUser,
Content: &schemas.ChatMessageContent{
ContentStr: schemas.Ptr("Test timeout propagation"),
},
},
},
}
result, bifrostErr := manager.CheckAndExecuteAgentForChatRequest(
ctx,
originalReq,
initialResponse,
mockLLM.MakeChatRequest,
func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) {
return manager.ExecuteToolCall(ctx, request)
},
)
// May timeout or complete with partial results
if bifrostErr != nil {
t.Logf("Timeout propagated with error: %v", bifrostErr.Error)
} else {
require.NotNil(t, result)
t.Logf("Partial completion with timeout")
}
t.Logf("✅ Timeout propagation through parallel tools verified")
}
func TestContext_RequestIDGeneration(t *testing.T) {
t.Parallel()
// Test that request IDs are generated and tracked correctly
manager := setupMCPManager(t)
// Register tool
requestIDs := []string{}
idHandler := func(args any) (string, error) {
// In real implementation, would capture request ID from context
requestID := fmt.Sprintf("req_%d", len(requestIDs))
requestIDs = append(requestIDs, requestID)
return fmt.Sprintf(`{"request_id": "%s"}`, requestID), nil
}
idSchema := schemas.ChatTool{
Type: schemas.ChatToolTypeFunction,
Function: &schemas.ChatToolFunction{
Name: "id_tool",
Description: schemas.Ptr("Tracks request IDs"),
},
}
err := manager.RegisterTool("id_tool", "Tracks IDs", idHandler, idSchema)
require.NoError(t, err)
err = SetInternalClientAutoExecute(manager, []string{"*"})
require.NoError(t, err)
ctx := createTestContext()
// Execute multiple iterations
mockLLM := &MockLLMCaller{
chatResponses: []*schemas.BifrostChatResponse{
CreateChatResponseWithToolCalls([]schemas.ChatAssistantMessageToolCall{
{
ID: schemas.Ptr("call-1"),
Type: schemas.Ptr("function"),
Function: schemas.ChatAssistantMessageToolCallFunction{
Name: schemas.Ptr("bifrostInternal-id_tool"),
Arguments: "{}",
},
},
}),
CreateChatResponseWithToolCalls([]schemas.ChatAssistantMessageToolCall{
{
ID: schemas.Ptr("call-2"),
Type: schemas.Ptr("function"),
Function: schemas.ChatAssistantMessageToolCallFunction{
Name: schemas.Ptr("bifrostInternal-id_tool"),
Arguments: "{}",
},
},
}),
CreateChatResponseWithText("ID tracking complete"),
},
}
initialResponse := mockLLM.chatResponses[0]
mockLLM.chatCallCount = 1
originalReq := &schemas.BifrostChatRequest{
Provider: schemas.OpenAI,
Model: "gpt-4",
Input: []schemas.ChatMessage{
{
Role: schemas.ChatMessageRoleUser,
Content: &schemas.ChatMessageContent{
ContentStr: schemas.Ptr("Test request IDs"),
},
},
},
}
result, bifrostErr := manager.CheckAndExecuteAgentForChatRequest(
ctx,
originalReq,
initialResponse,
mockLLM.MakeChatRequest,
func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) {
return manager.ExecuteToolCall(ctx, request)
},
)
require.Nil(t, bifrostErr)
require.NotNil(t, result)
t.Logf("✅ Request ID generation verified")
t.Logf("Generated IDs: %v", requestIDs)
}
func TestContext_CleanupOnCompletion(t *testing.T) {
t.Parallel()
// Test that contexts are properly cleaned up after tool execution
manager := setupMCPManager(t)
cleanupCount := 0
cleanupHandler := func(args any) (string, error) {
cleanupCount++
// Simulate resource usage
return fmt.Sprintf(`{"execution": %d}`, cleanupCount), nil
}
cleanupSchema := schemas.ChatTool{
Type: schemas.ChatToolTypeFunction,
Function: &schemas.ChatToolFunction{
Name: "cleanup_tool",
Description: schemas.Ptr("Tests cleanup"),
},
}
err := manager.RegisterTool("cleanup_tool", "Tests cleanup", cleanupHandler, cleanupSchema)
require.NoError(t, err)
bifrost := setupBifrost(t)
bifrost.SetMCPManager(manager)
ctx := createTestContext()
// Execute tool multiple times
for i := 0; i < 3; i++ {
toolCall := schemas.ChatAssistantMessageToolCall{
ID: schemas.Ptr(fmt.Sprintf("call-%d", i)),
Type: schemas.Ptr("function"),
Function: schemas.ChatAssistantMessageToolCallFunction{
Name: schemas.Ptr("bifrostInternal-cleanup_tool"),
Arguments: "{}",
},
}
result, bifrostErr := bifrost.ExecuteChatMCPTool(ctx, &toolCall)
require.Nil(t, bifrostErr)
require.NotNil(t, result)
}
assert.Equal(t, 3, cleanupCount, "should have executed 3 times")
t.Logf("✅ Context cleanup verified after %d executions", cleanupCount)
}
func TestContext_ConcurrentAccess(t *testing.T) {
t.Parallel()
// Test concurrent access to context values
manager := setupMCPManager(t)
concurrentHandler := func(args any) (string, error) {
// Simulate concurrent access
time.Sleep(10 * time.Millisecond)
return `{"result": "concurrent access"}`, nil
}
concurrentSchema := schemas.ChatTool{
Type: schemas.ChatToolTypeFunction,
Function: &schemas.ChatToolFunction{
Name: "concurrent_tool",
Description: schemas.Ptr("Tests concurrent access"),
},
}
err := manager.RegisterTool("concurrent_tool", "Tests concurrent access", concurrentHandler, concurrentSchema)
require.NoError(t, err)
err = SetInternalClientAutoExecute(manager, []string{"*"})
require.NoError(t, err)
ctx := createTestContext()
// Execute multiple tools in parallel
toolCalls := []schemas.ChatAssistantMessageToolCall{}
for i := 0; i < 10; i++ {
toolCalls = append(toolCalls, schemas.ChatAssistantMessageToolCall{
ID: schemas.Ptr(fmt.Sprintf("call-concurrent-%d", i)),
Type: schemas.Ptr("function"),
Function: schemas.ChatAssistantMessageToolCallFunction{
Name: schemas.Ptr("bifrostInternal-concurrent_tool"),
Arguments: "{}",
},
})
}
mockLLM := &MockLLMCaller{
chatResponses: []*schemas.BifrostChatResponse{
CreateChatResponseWithToolCalls(toolCalls),
CreateChatResponseWithText("Concurrent test complete"),
},
}
initialResponse := mockLLM.chatResponses[0]
mockLLM.chatCallCount = 1
originalReq := &schemas.BifrostChatRequest{
Provider: schemas.OpenAI,
Model: "gpt-4",
Input: []schemas.ChatMessage{
{
Role: schemas.ChatMessageRoleUser,
Content: &schemas.ChatMessageContent{
ContentStr: schemas.Ptr("Test concurrent access"),
},
},
},
}
result, bifrostErr := manager.CheckAndExecuteAgentForChatRequest(
ctx,
originalReq,
initialResponse,
mockLLM.MakeChatRequest,
func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) {
return manager.ExecuteToolCall(ctx, request)
},
)
require.Nil(t, bifrostErr)
require.NotNil(t, result)
t.Logf("✅ Concurrent context access verified (10 parallel tools)")
}