File tree 2 files changed +29
-4
lines changed
2 files changed +29
-4
lines changed Original file line number Diff line number Diff line change @@ -162,15 +162,19 @@ class Replicate {
162
162
progress ( updatedPrediction ) ;
163
163
}
164
164
165
- if ( signal && signal . aborted ) {
166
- await this . predictions . cancel ( updatedPrediction . id ) ;
165
+ // We handle the cancel later in the function.
166
+ if ( signal ?. aborted ) {
167
167
return true ; // stop polling
168
168
}
169
169
170
170
return false ; // continue polling
171
171
}
172
172
) ;
173
173
174
+ if ( signal ?. aborted ) {
175
+ prediction = await this . predictions . cancel ( prediction . id ) ;
176
+ }
177
+
174
178
// Call progress callback with the completed prediction object
175
179
if ( progress ) {
176
180
progress ( prediction ) ;
Original file line number Diff line number Diff line change @@ -1258,19 +1258,40 @@ describe("Replicate client", () => {
1258
1258
status : "canceled" ,
1259
1259
} ) ;
1260
1260
1261
- await client . run (
1261
+ const onProgress = jest . fn ( ) ;
1262
+ const output = await client . run (
1262
1263
"owner/model:5c7d5dc6dd8bf75c1acaa8565735e7986bc5b66206b55cca93cb72c9bf15ccaa" ,
1263
1264
{
1264
1265
input : { text : "Hello, world!" } ,
1265
1266
signal,
1266
- }
1267
+ } ,
1268
+ onProgress
1267
1269
) ;
1268
1270
1269
1271
expect ( body ) . toBeDefined ( ) ;
1270
1272
expect ( body ?. [ "signal" ] ) . toBeUndefined ( ) ;
1271
1273
expect ( signal . aborted ) . toBe ( true ) ;
1272
1274
expect ( output ) . toBeUndefined ( ) ;
1273
1275
1276
+ expect ( onProgress ) . toHaveBeenNthCalledWith (
1277
+ 1 ,
1278
+ expect . objectContaining ( {
1279
+ status : "processing" ,
1280
+ } )
1281
+ ) ;
1282
+ expect ( onProgress ) . toHaveBeenNthCalledWith (
1283
+ 2 ,
1284
+ expect . objectContaining ( {
1285
+ status : "processing" ,
1286
+ } )
1287
+ ) ;
1288
+ expect ( onProgress ) . toHaveBeenNthCalledWith (
1289
+ 3 ,
1290
+ expect . objectContaining ( {
1291
+ status : "canceled" ,
1292
+ } )
1293
+ ) ;
1294
+
1274
1295
scope . done ( ) ;
1275
1296
} ) ;
1276
1297
} ) ;
You can’t perform that action at this time.
0 commit comments