Enhance context cancellation support in OpenAI provider and tests
- Added context cancellation handling in the OpenAI provider's Stream and Post methods, allowing for graceful termination of operations when the context is cancelled. - Implemented a new test, TestOpenAIStreamContextCancellation, to validate that streaming respects context cancellation, ensuring proper error handling and event emission during cancellation scenarios. - Refactored context usage in existing methods to improve consistency and reliability in handling context across streaming operations.
This commit is contained in:
Max
parent
da53917d51
commit
8636b0c842
2 changed files with 174 additions and 9 deletions
|
|
@ -121,17 +121,39 @@ func (p *Provider) Stream(ctx *context.Context, messages []context.Message, opti
|
|||
maxValidationRetries := 3
|
||||
var lastErr error
|
||||
|
||||
// Get Go context for cancellation support
|
||||
goCtx := ctx.Context
|
||||
if goCtx == nil {
|
||||
goCtx = gocontext.Background()
|
||||
}
|
||||
|
||||
// Make a copy of messages to avoid modifying the original
|
||||
currentMessages := make([]context.Message, len(messages))
|
||||
copy(currentMessages, messages)
|
||||
|
||||
// Outer loop: handle network/API errors with exponential backoff
|
||||
for attempt := 0; attempt < maxRetries; attempt++ {
|
||||
// Check if context is cancelled before retry
|
||||
select {
|
||||
case <-goCtx.Done():
|
||||
return nil, fmt.Errorf("context cancelled: %w", goCtx.Err())
|
||||
default:
|
||||
}
|
||||
|
||||
if attempt > 0 {
|
||||
// Exponential backoff: 1s, 2s, 4s
|
||||
backoff := time.Duration(1<<uint(attempt-1)) * time.Second
|
||||
log.Warn("OpenAI stream request failed, retrying in %v (attempt %d/%d): %v", backoff, attempt+1, maxRetries, lastErr)
|
||||
time.Sleep(backoff)
|
||||
|
||||
// Sleep with context cancellation support
|
||||
timer := time.NewTimer(backoff)
|
||||
select {
|
||||
case <-timer.C:
|
||||
// Continue to retry
|
||||
case <-goCtx.Done():
|
||||
timer.Stop()
|
||||
return nil, fmt.Errorf("context cancelled during backoff: %w", goCtx.Err())
|
||||
}
|
||||
}
|
||||
|
||||
response, err := p.streamWithRetry(ctx, currentMessages, options, handler)
|
||||
|
|
@ -188,6 +210,19 @@ func (p *Provider) streamWithRetry(ctx *context.Context, messages []context.Mess
|
|||
streamStartTime := time.Now()
|
||||
requestID := fmt.Sprintf("req_%d", streamStartTime.UnixNano())
|
||||
|
||||
// Get Go context for cancellation support
|
||||
goCtx := ctx.Context
|
||||
if goCtx == nil {
|
||||
goCtx = gocontext.Background()
|
||||
}
|
||||
|
||||
// Check if context is already cancelled
|
||||
select {
|
||||
case <-goCtx.Done():
|
||||
return nil, fmt.Errorf("context cancelled before stream start: %w", goCtx.Err())
|
||||
default:
|
||||
}
|
||||
|
||||
// Send stream_start event
|
||||
if handler != nil {
|
||||
model, _ := p.GetModel()
|
||||
|
|
@ -256,6 +291,14 @@ func (p *Provider) streamWithRetry(ctx *context.Context, messages []context.Mess
|
|||
|
||||
// Stream handler
|
||||
streamHandler := func(data []byte) int {
|
||||
// Check for context cancellation
|
||||
select {
|
||||
case <-goCtx.Done():
|
||||
log.Warn("Stream cancelled by context")
|
||||
return http.HandlerReturnBreak
|
||||
default:
|
||||
}
|
||||
|
||||
if len(data) == 0 {
|
||||
return http.HandlerReturnOk
|
||||
}
|
||||
|
|
@ -401,13 +444,30 @@ func (p *Provider) streamWithRetry(ctx *context.Context, messages []context.Mess
|
|||
return http.HandlerReturnOk
|
||||
}
|
||||
|
||||
// Make streaming request
|
||||
goCtx := ctx.Context
|
||||
if goCtx == nil {
|
||||
goCtx = gocontext.Background()
|
||||
// Make streaming request (goCtx already set at function start)
|
||||
err = req.Stream(goCtx, "POST", requestBody, streamHandler)
|
||||
|
||||
// Check if error is due to context cancellation
|
||||
if err != nil && goCtx.Err() != nil {
|
||||
// End current group if active
|
||||
groupTracker.endGroup(handler)
|
||||
|
||||
// Send stream_end with cancellation status
|
||||
if handler != nil {
|
||||
endData := &context.StreamEndData{
|
||||
RequestID: requestID,
|
||||
Timestamp: time.Now().UnixMilli(),
|
||||
DurationMs: time.Since(streamStartTime).Milliseconds(),
|
||||
Status: "cancelled",
|
||||
Error: goCtx.Err().Error(),
|
||||
}
|
||||
if endJSON, err := jsoniter.Marshal(endData); err == nil {
|
||||
handler(context.ChunkStreamEnd, endJSON)
|
||||
}
|
||||
}
|
||||
return nil, fmt.Errorf("stream cancelled: %w", goCtx.Err())
|
||||
}
|
||||
|
||||
err = req.Stream(goCtx, "POST", requestBody, streamHandler)
|
||||
if err != nil {
|
||||
// End current group if active
|
||||
groupTracker.endGroup(handler)
|
||||
|
|
@ -540,17 +600,39 @@ func (p *Provider) Post(ctx *context.Context, messages []context.Message, option
|
|||
maxValidationRetries := 3
|
||||
var lastErr error
|
||||
|
||||
// Get Go context for cancellation support
|
||||
goCtx := ctx.Context
|
||||
if goCtx == nil {
|
||||
goCtx = gocontext.Background()
|
||||
}
|
||||
|
||||
// Make a copy of messages to avoid modifying the original
|
||||
currentMessages := make([]context.Message, len(messages))
|
||||
copy(currentMessages, messages)
|
||||
|
||||
// Outer loop: handle network/API errors with exponential backoff
|
||||
for attempt := 0; attempt < maxRetries; attempt++ {
|
||||
// Check if context is cancelled before retry
|
||||
select {
|
||||
case <-goCtx.Done():
|
||||
return nil, fmt.Errorf("context cancelled: %w", goCtx.Err())
|
||||
default:
|
||||
}
|
||||
|
||||
if attempt > 0 {
|
||||
// Exponential backoff
|
||||
backoff := time.Duration(1<<uint(attempt-1)) * time.Second
|
||||
log.Warn("OpenAI post request failed, retrying in %v (attempt %d/%d): %v", backoff, attempt+1, maxRetries, lastErr)
|
||||
time.Sleep(backoff)
|
||||
|
||||
// Sleep with context cancellation support
|
||||
timer := time.NewTimer(backoff)
|
||||
select {
|
||||
case <-timer.C:
|
||||
// Continue to retry
|
||||
case <-goCtx.Done():
|
||||
timer.Stop()
|
||||
return nil, fmt.Errorf("context cancelled during backoff: %w", goCtx.Err())
|
||||
}
|
||||
}
|
||||
|
||||
response, err := p.postWithRetry(ctx, currentMessages, options)
|
||||
|
|
|
|||
|
|
@ -1,10 +1,11 @@
|
|||
package openai_test
|
||||
|
||||
import (
|
||||
stdContext "context"
|
||||
gocontext "context"
|
||||
"encoding/json"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/yaoapp/gou/connector"
|
||||
"github.com/yaoapp/gou/plan"
|
||||
|
|
@ -1388,6 +1389,88 @@ func TestOpenAIStreamLifecycleEvents(t *testing.T) {
|
|||
t.Log("Lifecycle events test completed successfully")
|
||||
}
|
||||
|
||||
// TestOpenAIStreamContextCancellation tests that stream respects context cancellation
|
||||
func TestOpenAIStreamContextCancellation(t *testing.T) {
|
||||
test.Prepare(t, config.Conf)
|
||||
defer test.Clean()
|
||||
|
||||
conn, err := connector.Select("openai.gpt-4o")
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to select connector: %v", err)
|
||||
}
|
||||
|
||||
trueVal := true
|
||||
options := &context.CompletionOptions{
|
||||
Capabilities: &context.ModelCapabilities{
|
||||
Streaming: &trueVal,
|
||||
ToolCalls: &trueVal,
|
||||
},
|
||||
}
|
||||
|
||||
llmInstance, err := llm.New(conn, options)
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to create LLM instance: %v", err)
|
||||
}
|
||||
|
||||
messages := []context.Message{
|
||||
{
|
||||
Role: context.RoleUser,
|
||||
Content: "Write a very long essay about the history of computing", // Long task
|
||||
},
|
||||
}
|
||||
|
||||
// Create a context with a very short timeout
|
||||
ctx := newTestContext("test-cancel", "openai.gpt-4o")
|
||||
goCtx, cancel := gocontext.WithTimeout(gocontext.Background(), 100*time.Millisecond)
|
||||
defer cancel()
|
||||
ctx.Context = goCtx
|
||||
|
||||
var receivedChunks int
|
||||
var receivedStreamEnd bool
|
||||
|
||||
handler := func(chunkType context.StreamChunkType, data []byte) int {
|
||||
if chunkType == context.ChunkText || chunkType == context.ChunkToolCall {
|
||||
receivedChunks++
|
||||
}
|
||||
if chunkType == context.ChunkStreamEnd {
|
||||
receivedStreamEnd = true
|
||||
var endData context.StreamEndData
|
||||
if err := json.Unmarshal(data, &endData); err == nil {
|
||||
t.Logf("stream_end status: %s, error: %s", endData.Status, endData.Error)
|
||||
}
|
||||
}
|
||||
return 0
|
||||
}
|
||||
|
||||
response, err := llmInstance.Stream(ctx, messages, options, handler)
|
||||
|
||||
// Should get an error due to context cancellation
|
||||
if err == nil {
|
||||
t.Error("Expected error due to context cancellation, but got nil")
|
||||
} else {
|
||||
t.Logf("✓ Got expected cancellation error: %v", err)
|
||||
|
||||
// Check if error message indicates cancellation
|
||||
errStr := err.Error()
|
||||
if !strings.Contains(errStr, "context") && !strings.Contains(errStr, "cancel") {
|
||||
t.Errorf("Error should mention context/cancellation: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// Response should be nil due to cancellation
|
||||
if response != nil {
|
||||
t.Logf("Warning: Response is not nil despite cancellation (partial response)")
|
||||
}
|
||||
|
||||
// Should have received stream_end event (even for cancellation)
|
||||
if !receivedStreamEnd {
|
||||
t.Error("Expected stream_end event even for cancelled stream")
|
||||
}
|
||||
|
||||
t.Logf("Received %d chunks before cancellation", receivedChunks)
|
||||
t.Log("Context cancellation test completed successfully")
|
||||
}
|
||||
|
||||
// TestOpenAIStreamWithTemperature tests different temperature settings
|
||||
func TestOpenAIStreamWithTemperature(t *testing.T) {
|
||||
test.Prepare(t, config.Conf)
|
||||
|
|
@ -1476,7 +1559,7 @@ func TestOpenAIStreamWithTemperature(t *testing.T) {
|
|||
// newTestContext creates a real Context for testing OpenAI provider
|
||||
func newTestContext(chatID, connectorID string) *context.Context {
|
||||
return &context.Context{
|
||||
Context: stdContext.Background(),
|
||||
Context: gocontext.Background(),
|
||||
Space: plan.NewMemorySharedSpace(),
|
||||
ChatID: chatID,
|
||||
AssistantID: "test-assistant",
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue