From 8df533b489d7b70e5207beddce5bf9a7178432cf Mon Sep 17 00:00:00 2001
From: Tat Dat Duong <david@duong.cz>
Date: Wed, 9 Oct 2024 20:34:19 +0200
Subject: [PATCH] fix(sdk-js): allow passing `checkpoint` when creating a run

---
 libs/sdk-js/src/client.ts | 21 +++++++++++++++++++++
 libs/sdk-js/src/types.ts  |  8 +++++++-
 2 files changed, 28 insertions(+), 1 deletion(-)

diff --git a/libs/sdk-js/src/client.ts b/libs/sdk-js/src/client.ts
index 05a369191..eb39d5a61 100644
--- a/libs/sdk-js/src/client.ts
+++ b/libs/sdk-js/src/client.ts
@@ -664,6 +664,12 @@ export class RunsClient extends BaseClient {
     event: StreamEvent;
     data: any;
   }> {
+    if (!!payload?.checkpointId && !!payload?.checkpoint) {
+      throw new Error(
+        "Cannot specify both a checkpoint and a checkpoint ID in the payload.",
+      );
+    }
+
     const json: Record<string, any> = {
       input: payload?.input,
       config: payload?.config,
@@ -674,6 +680,7 @@ export class RunsClient extends BaseClient {
       assistant_id: assistantId,
       interrupt_before: payload?.interruptBefore,
       interrupt_after: payload?.interruptAfter,
+      checkpoint: payload?.checkpoint,
       checkpoint_id: payload?.checkpointId,
       webhook: payload?.webhook,
       multitask_strategy: payload?.multitaskStrategy,
@@ -748,6 +755,12 @@ export class RunsClient extends BaseClient {
     assistantId: string,
     payload?: RunsCreatePayload,
   ): Promise<Run> {
+    if (!!payload?.checkpointId && !!payload?.checkpoint) {
+      throw new Error(
+        "Cannot specify both a checkpoint and a checkpoint ID in the payload.",
+      );
+    }
+
     const json: Record<string, any> = {
       input: payload?.input,
       config: payload?.config,
@@ -756,6 +769,7 @@ export class RunsClient extends BaseClient {
       interrupt_before: payload?.interruptBefore,
       interrupt_after: payload?.interruptAfter,
       webhook: payload?.webhook,
+      checkpoint: payload?.checkpoint,
       checkpoint_id: payload?.checkpointId,
       multitask_strategy: payload?.multitaskStrategy,
       after_seconds: payload?.afterSeconds,
@@ -815,6 +829,12 @@ export class RunsClient extends BaseClient {
     assistantId: string,
     payload?: RunsWaitPayload,
   ): Promise<ThreadState["values"]> {
+    if (!!payload?.checkpointId && !!payload?.checkpoint) {
+      throw new Error(
+        "Cannot specify both a checkpoint and a checkpoint ID in the payload.",
+      );
+    }
+
     const json: Record<string, any> = {
       input: payload?.input,
       config: payload?.config,
@@ -822,6 +842,7 @@ export class RunsClient extends BaseClient {
       assistant_id: assistantId,
       interrupt_before: payload?.interruptBefore,
       interrupt_after: payload?.interruptAfter,
+      checkpoint: payload?.checkpoint,
       checkpoint_id: payload?.checkpointId,
       webhook: payload?.webhook,
       multitask_strategy: payload?.multitaskStrategy,
diff --git a/libs/sdk-js/src/types.ts b/libs/sdk-js/src/types.ts
index c3ca16ef0..42181c368 100644
--- a/libs/sdk-js/src/types.ts
+++ b/libs/sdk-js/src/types.ts
@@ -1,4 +1,4 @@
-import { Config, Metadata } from "./schema.js";
+import { Checkpoint, Config, Metadata } from "./schema.js";
 
 export type StreamMode = "values" | "messages" | "updates" | "events" | "debug";
 export type MultitaskStrategy = "reject" | "interrupt" | "rollback" | "enqueue";
@@ -34,9 +34,15 @@ interface RunsInvokePayload {
 
   /**
    * Checkpoint ID for when creating a new run.
+   * @deprecated Use `checkpoint` instead.
    */
   checkpointId?: string;
 
+  /**
+   * Checkpoint for when creating a new run.
+   */
+  checkpoint?: Omit<Checkpoint, "thread_id">;
+
   /**
    * Interrupt execution before entering these nodes.
    */