From 8636b0c842a6251679947ca8b69d3a3ff5809918 Mon Sep 17 00:00:00 2001
From: Max <shadow.iqka@qq.com>
Date: Sat, 15 Nov 2025 15:43:57 +0800
Subject: [PATCH] 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.
---
 agent/llm/providers/openai/openai.go      | 96 +++++++++++++++++++++--
 agent/llm/providers/openai/openai_test.go | 87 +++++++++++++++++++-
 2 files changed, 174 insertions(+), 9 deletions(-)

diff --git a/agent/llm/providers/openai/openai.go b/agent/llm/providers/openai/openai.go
index 62d14b52..93917958 100644
--- a/agent/llm/providers/openai/openai.go
+++ b/agent/llm/providers/openai/openai.go
@@ -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)
diff --git a/agent/llm/providers/openai/openai_test.go b/agent/llm/providers/openai/openai_test.go
index 16b148aa..ca03e02c 100644
--- a/agent/llm/providers/openai/openai_test.go
+++ b/agent/llm/providers/openai/openai_test.go
@@ -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",