Bug fixes for the wait option in replicate.run

index.d.ts

@@ -93,9 +93,9 @@ declare module "replicate" {
     model: string;
     version: string;
     input: object;
-    output?: any;
+    output?: any; // TODO: this should be `unknown`
     source: "api" | "web";
-    error?: any;
+    error?: unknown;
     logs?: string;
     metrics?: {
       predict_time?: number;
@@ -156,13 +156,13 @@ declare module "replicate" {
       identifier: `${string}/${string}` | `${string}/${string}:${string}`,
       options: {
         input: object;
-        wait?: boolean | number | { mode?: "poll"; interval?: number };
+        wait?: boolean | number | { interval?: number };
         webhook?: string;
         webhook_events_filter?: WebhookEventType[];
         signal?: AbortSignal;
       },
       progress?: (prediction: Prediction) => void
-    ): Promise<object>;
+    ): Promise<unknown>;
 
     stream(
       identifier: `${string}/${string}` | `${string}/${string}:${string}`,
@@ -215,7 +215,7 @@ declare module "replicate" {
             stream?: boolean;
             webhook?: string;
             webhook_events_filter?: WebhookEventType[];
-            block?: boolean;
+            wait?: boolean | number | { mode?: "poll"; interval?: number };
           }
         ): Promise<Prediction>;
       };
@@ -304,7 +304,7 @@ declare module "replicate" {
           stream?: boolean;
           webhook?: string;
           webhook_events_filter?: WebhookEventType[];
-          block?: boolean;
+          wait?: boolean | number | { mode?: "poll"; interval?: number };
         } & ({ version: string } | { model: string })
       ): Promise<Prediction>;
       get(prediction_id: string): Promise<Prediction>;

index.js

@@ -147,19 +147,20 @@ class Replicate {
     const { wait, signal, ...data } = options;
 
     const identifier = ModelVersionIdentifier.parse(ref);
+    const isBlocking = typeof wait === "boolean" || typeof wait === "number";
 
     let prediction;
     if (identifier.version) {
       prediction = await this.predictions.create({
         ...data,
         version: identifier.version,
-        wait: wait,
+        wait: isBlocking ? wait : false,
       });
     } else if (identifier.owner && identifier.name) {
       prediction = await this.predictions.create({
         ...data,
         model: `${identifier.owner}/${identifier.name}`,
-        wait: wait,
+        wait: isBlocking ? wait : false,
       });
     } else {
       throw new Error("Invalid model version identifier");
@@ -170,23 +171,26 @@ class Replicate {
       progress(prediction);
     }
 
-    prediction = await this.wait(
-      prediction,
-      wait || {},
-      async (updatedPrediction) => {
-        // Call progress callback with the updated prediction object
-        if (progress) {
-          progress(updatedPrediction);
+    const isDone = isBlocking && prediction.status !== "starting";
+    if (!isDone) {
+      prediction = await this.wait(
+        prediction,
+        isBlocking ? {} : wait,
+        async (updatedPrediction) => {
+          // Call progress callback with the updated prediction object
+          if (progress) {
+            progress(updatedPrediction);
+          }
+
+          // We handle the cancel later in the function.
+          if (signal && signal.aborted) {
+            return true; // stop polling
+          }
+
+          return false; // continue polling
         }
-
-        // We handle the cancel later in the function.
-        if (signal && signal.aborted) {
-          return true; // stop polling
-        }
-
-        return false; // continue polling
-      }
-    );
+      );
+    }
 
     if (signal && signal.aborted) {
       prediction = await this.predictions.cancel(prediction.id);