diff --git a/sdk/api/handlers/handlers_stream.go b/sdk/api/handlers/handlers_stream.go index 6f0bfbc931b..ee69f89378b 100644 --- a/sdk/api/handlers/handlers_stream.go +++ b/sdk/api/handlers/handlers_stream.go @@ -488,17 +488,10 @@ func (h *BaseAPIHandler) executeStreamWithAuthManagerFormats(ctx context.Context var bootstrapErr *interfaces.ErrorMessage readInitialStreamChunks := func() { for { - var chunk coreexecutor.StreamChunk - var ok bool - if ctx != nil { - select { - case <-ctx.Done(): - streamCanceledBeforeRead = true - return - case chunk, ok = <-chunks: - } - } else { - chunk, ok = <-chunks + chunk, ok, canceled := nextStreamChunk(ctx, nil, nil, chunks) + if canceled { + streamCanceledBeforeRead = true + return } if !ok { streamClosedBeforeRead = true diff --git a/sdk/api/handlers/model_execution_delivery_test.go b/sdk/api/handlers/model_execution_delivery_test.go index 4724a06f572..252a6bc5119 100644 --- a/sdk/api/handlers/model_execution_delivery_test.go +++ b/sdk/api/handlers/model_execution_delivery_test.go @@ -3,11 +3,13 @@ package handlers import ( "context" "testing" + "time" coreauth "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/auth" coreexecutor "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/executor" "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/usage" sdkconfig "github.com/router-for-me/CLIProxyAPI/v8/sdk/config" + "github.com/router-for-me/CLIProxyAPI/v8/sdk/pluginapi" ) func TestInternalModelExecutionIsolatesDelivery(t *testing.T) { @@ -65,3 +67,43 @@ func TestInternalModelExecutionIsolatesDelivery(t *testing.T) { }) } } + +// An internal stream canceled before the handler has its first chunk: the auth manager's forwarder stops on the cancellation and +// closes its channel, so the close and the cancellation can be ready together at the handler's first read. Without delivery tracking +// nothing corrects the outcome later, so that read must report the cancellation. +func TestExecuteModelStreamCanceledBeforeFirstChunkCompletesAsCanceled(t *testing.T) { + model := "nested-cancel-first-read-model" + var cancel context.CancelFunc + executor := &modelExecutionCaptureExecutor{ + stream: func(context.Context, *coreauth.Auth, coreexecutor.Request, coreexecutor.Options) (*coreexecutor.StreamResult, error) { + cancel() + chunks := make(chan coreexecutor.StreamChunk, 1) + chunks <- coreexecutor.StreamChunk{Payload: []byte("data: {\"ok\":true}\n\n")} + close(chunks) + return &coreexecutor.StreamResult{Chunks: chunks}, nil + }, + } + h := newModelExecutionHandler(t, model, executor, &sdkconfig.SDKConfig{}) + completions := make(chan pluginapi.RequestCompletion, 1) + h.SetPluginHost(&handlerInterceptorTestHost{completeRequest: func(_ context.Context, completion pluginapi.RequestCompletion) { + completions <- completion + }}) + for round := range 2000 { + ctx, roundCancel := context.WithCancel(context.Background()) + cancel = roundCancel + result, errMsg := h.ExecuteModelStream(ctx, ModelExecutionRequest{Model: model, EntryProtocol: "openai", Stream: true, Body: []byte(`{"model":"` + model + `"}`)}) + if errMsg == nil { + for range result.Chunks { + } + } + roundCancel() + select { + case completion := <-completions: + if completion.Outcome != pluginapi.RequestCompletionCanceled { + t.Fatalf("round %d: completion = %+v, want a cancellation", round, completion) + } + case <-time.After(5 * time.Second): + t.Fatal("missing lifecycle completion") + } + } +}